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