Update README.md

README.md

@@ -23,6 +23,7 @@ An example usage of the model is below.
 
 ```
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
 
 tokenizer = AutoTokenizer.from_pretrained("emmatliu/language-agency-classifier")
 model = AutoModelForSequenceClassification.from_pretrained("emmatliu/language-agency-classifier")
@@ -33,9 +34,14 @@ inputs = tokenizer(sentence, return_tensors="pt")
 outputs = model(**inputs)
 probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
 
-predicted_class = torch.argmax(probabilities)
+predicted_class = torch.argmax(probabilities).item()
 
-print(f"Predicted class: {predicted_class}")
+labels = {
+    1: 'agentic',
+    0: 'communal'
+}
+
+print(f"Predicted class: {labels[predicted_class]}")
 ```
 
 ### Model Sources