banochka/ArticleTopicsClassification

1import torch2from utils import *3from transformers import AutoModelForSequenceClassification, AutoTokenizer4import json5 6checkpoint = 'roberta-base'7model_saved = AutoModelForSequenceClassification.from_pretrained('model_trained/checkpoint-133500')8tokenizer = AutoTokenizer.from_pretrained(checkpoint, do_lower_case=True)9 10model_saved.eval()11 12 13def get_answer(all_text):14    with torch.no_grad():15        logits = model_saved(**tokenizer(all_text, return_tensors='pt')).logits16 17    indexes = torch.softmax(logits, dim=-1).numpy().flatten().argsort()[::-1]18    logits = torch.softmax(logits, dim=-1).numpy().flatten()19 20    name_dict = json.load(open('name_dict.json', 'r'))21 22    args, probs = get_predictions(logits, indexes)23    out = return_pred_name(name_dict, args, probs)24    return ', '.join(out)25 26 27 28 29