Update README.md

README.md

@@ -39,14 +39,21 @@ pip install transformers torch accelerate
 
 Then the model can be downloaded and used for inference:
 ```py
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
+from transformers import pipeline
 
-model = AutoModelForSequenceClassification.from_pretrained("identrics/wasper_propaganda_detection_en", num_labels=2)
-tokenizer = AutoTokenizer.from_pretrained("identrics/wasper_propaganda_detection_en")
+labels_map = {"LABEL_0": "No Propaganda", "LABEL_1": "Propaganda"}
+
+pipe = pipeline(
+    "text-classification",
+    model="identrics/wasper_propaganda_detection_en",
+    tokenizer="identrics/wasper_propaganda_detection_en",
+)
+
+text = "Our country is the most powerful country in the world!"
+
+prediction = pipe(text)
+print(labels_map[prediction[0]["label"]])
 
-tokens = tokenizer("Our country is the most powerful country in the world!", return_tensors="pt")
-output = model(**tokens)
-print(output.logits)
 ```
 
 ## Training Details