NLPclass

/

parsBERT_text_emotion_classification

like 0

Hello,
I checked your model but i recieved only 'sad' label. I think I do wrong. please check this code:

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch


tokenizer = AutoTokenizer.from_pretrained("NLPclass/parsBERT_text_emotion_classification")
model = AutoModelForSequenceClassification.from_pretrained("NLPclass/parsBERT_text_emotion_classification")

emotion_labels = [
    "joy",
    "sad",
    "anger",
    "disgust",
    "fear",
    "surprise"
]

sentence =["من امروز واقعا ناراحتم.", "من از اینجا خیلی می ترسم",
                         "ما خیلی از دیدنت خوشحالیم", "می شه ما رو نجات بدی", "خسته شدم دیگه"]

inputs = tokenizer(sentence, return_tensors="pt", truncation=True, padding=True)

outputs = model(**inputs)
logits = outputs.logits

predicted_class = torch.argmax(logits, dim=1)
for i in range(len(predicted_class)):
    predicted_emotion = emotion_labels[predicted_class[i]]
    print(f"Emotion: {predicted_emotion}")

Output is :
Sentence: من امروز واقعا ناراحتم.
Emotion: sad

Sentence: من از اینجا خیلی می ترسم
Emotion: sad

Sentence: ما خیلی از دیدنت خوشحالیم
Emotion: sad

Sentence: می شه ما رو نجات بدی
Emotion: sad

Sentence: خسته شدم دیگه
Emotion: sad