Update README.md

README.md

@@ -2,9 +2,12 @@
 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)
@@ -12,21 +15,20 @@ 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'
-)
+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():
-    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]
+        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}")```