metadata
license: openrail
import torch
from transformers import AutoTokenizer, MobileBertForSequenceClassification
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 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"
input_ids = tokenizer.encode(input_text, add_special_tokens=True, truncation=True, max_length=128)
input_tensor = torch.tensor([input_ids]).to(device)
# Predict emotion
with torch.no_grad():
outputs = model(input_tensor)
logits = outputs[0]
# Get the predicted label
predicted_emotion = torch.argmax(logits, dim=1).item()
emotion_labels = {0:'sadness',1:'joy',2:'love',3:'anger',4:'fear',5:'surprise'}
predicted_emotion_label = emotion_labels[predicted_emotion]
print(f"Input text: {input_text}")
print(f"Predicted emotion: {predicted_emotion_label}")```