metadata
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}")```