---
license: openrail
---
import torch
from transformers import AutoTokenizer, MobileBertForSequenceClassification

# Load the saved model
model_name = 'harshith20/Emotion_predictor'
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = MobileBertForSequenceClassification.from_pretrained(model_name)

# Tokenize input text
input_text = "I am feeling happy today"
encoded_text = tokenizer.encode_plus(
    input_text,
    max_length=128,
    padding='max_length',
    truncation=True,
    return_attention_mask=True,
    return_tensors='pt'
)

# Predict emotion
with torch.no_grad():
    logits = model(**encoded_text)[0]
    predicted_emotion = torch.argmax(logits).item()
    emotion_labels = ['anger', 'fear', 'joy', 'love', 'sadness', 'surprise']
    predicted_emotion_label = emotion_labels[predicted_emotion]

print(f"Input text: {input_text}")
print(f"Predicted emotion: {predicted_emotion_label}")