Spaces:
Running
Running
from transformers import AutoModelForSequenceClassification, AutoTokenizer | |
import torch | |
import torch.nn.functional as F | |
import const | |
def load_model(): | |
return AutoModelForSequenceClassification.from_pretrained( | |
"koshin2001/Japanese-to-emotions" | |
).eval() | |
def load_tokenizer(): | |
return AutoTokenizer.from_pretrained("koshin2001/Japanese-to-emotions") | |
def exec(text): | |
model = load_model() | |
tokenizer = load_tokenizer() | |
inputs = tokenizer( | |
text, | |
return_tensors="pt", | |
truncation=True, | |
return_token_type_ids=False, | |
max_length=512, | |
) | |
output = model(**inputs) | |
output_logits = torch.tensor(output.logits).clone().detach().requires_grad_(True) | |
probs = F.softmax(output_logits, dim=-1).tolist()[0] | |
emotion_probs = dict(zip(const.EMOTIONS, probs)) | |
return emotion_probs | |