Update README.md

README.md

@@ -68,14 +68,36 @@ pip install transformers torch accelerate
 
 Then the model can be downloaded and used for inference:
 ```py
+import torch
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
-model = AutoModelForSequenceClassification.from_pretrained("identrics/wasper_propaganda_classifier_en", num_labels=5)
+labels = [
+    "Legitimisation Techniques",
+    "Rhetorical Devices",
+    "Logical Fallacies",
+    "Self-Identification Techniques",
+    "Defamation Techniques",
+]
+
+model = AutoModelForSequenceClassification.from_pretrained(
+    "identrics/wasper_propaganda_classifier_en", num_labels=5
+)
 tokenizer = AutoTokenizer.from_pretrained("identrics/wasper_propaganda_classifier_en")
 
-tokens = tokenizer("Our country is the most powerful country in the world!", return_tensors="pt")
-output = model(**tokens)
-print(output.logits)
+text = "Газа евтин, американското ядрено гориво евтино, пълно с фотоволтаици а пък тока с 30% нагоре. Защо ?"
+
+inputs = tokenizer(text, padding=True, truncation=True, return_tensors="pt")
+
+with torch.no_grad():
+    outputs = model(**inputs)
+    logits = outputs.logits
+
+probabilities = torch.sigmoid(logits).cpu().numpy().flatten()
+
+# Format predictions
+
+predictions = {labels[i]: probabilities[i] for i in range(len(labels))}
+print(predictions)
 ```