Update README.md

README.md

@@ -50,11 +50,21 @@ import torch
 model = AutoModelForSequenceClassification.from_pretrained("mervp/SentimentBERT")
 tokenizer = AutoTokenizer.from_pretrained("mervp/SentimentBERT")
 
-text = "The government’s response to the crisis was surprisingly effective."
-inputs = tokenizer(text, return_tensors="pt")
-
-with torch.no_grad():
-    logits = model(**inputs).logits
-
-predicted_class = torch.argmax(logits, dim=1).item()
-print(["Negative", "Neutral", "Positive"][predicted_class])
+def predict_sentiment(text):
+    model.eval()
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+    with torch.no_grad():
+        outputs = model(**inputs)
+        logits = outputs.logits
+        prediction = torch.argmax(logits, dim=-1).item()
+    label = model.config.id2label[prediction]
+    return label
+
+print(predict_sentiment("What a beautiful day."))               # positive
+print(predict_sentiment("The service was excellent."))          # positive
+print(predict_sentiment("He did a fantastic job."))             # positive
+print(predict_sentiment("The experience was terrible."))        # negative
+print(predict_sentiment("Everything went wrong."))              # negative
+print(predict_sentiment("He opened the door and walked in."))   # neutral
+print(predict_sentiment("They are meeting at 5 PM."))           # neutral
+print(predict_sentiment("She has a cat."))                      # neutral