Update README.md

README.md

@@ -83,22 +83,22 @@ For more detailed control over predictions:
 ```python
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch.nn.functional as F
 
-# Load model and tokenizer, set device
-model_path = "MidhunKanadan/roberta-large-fallacy-classification"
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = AutoModelForSequenceClassification.from_pretrained(model_path).to(device).eval()
-tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
-
-# Tokenize input and get probabilities
+model_path = "MidhunKanadan/roberta-large-logical-fallacy-classifier"
 text = "The rooster crows always before the sun rises, therefore the crowing rooster causes the sun to rise."
-inputs = tokenizer(text, return_tensors="pt").to(device)
+
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+model = AutoModelForSequenceClassification.from_pretrained(model_path).to("cuda")
+inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=128).to("cuda")
+
 with torch.no_grad():
-    probabilities = torch.nn.functional.softmax(model(**inputs).logits, dim=1)[0]
+    probs = F.softmax(model(**inputs).logits, dim=-1)
+    results = {model.config.id2label[i]: score.item() for i, score in enumerate(probs[0])}
 
-# Output sorted results
-for label, score in sorted(zip(model.config.id2label.values(), probabilities), key=lambda x: x[1], reverse=True):
-    print(f"{label}: {score.item():.4f}")
+# Print scores for all labels
+for label, score in sorted(results.items(), key=lambda x: x[1], reverse=True):
+    print(f"{label}: {score:.4f}")
 ```
 
 Expected Output: