Japanese-to-Emotions / inference.py
koshin2001's picture
πŸŽ‰ Add app
6364f02
raw
history blame contribute delete
860 Bytes
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