Update README.md

README.md

@@ -13,17 +13,17 @@ This model is trained for toxicity classification task. The dataset used for tra
 
 ## How to use
 ```python
+import torch
 from transformers import RobertaTokenizer, RobertaForSequenceClassification
 
-# load tokenizer and model weights
 tokenizer = RobertaTokenizer.from_pretrained('s-nlp/roberta_toxicity_classifier')
 model = RobertaForSequenceClassification.from_pretrained('s-nlp/roberta_toxicity_classifier')
 
-# prepare the input
-batch = tokenizer.encode('you are amazing', return_tensors='pt')
+batch = tokenizer.encode("You are amazing!", return_tensors="pt")
 
-# inference
-model(batch)
+output = model(batch)
+predicted_label = torch.sigmoid(output.logits).argmax().item()
+# 0 for neutral, 1 for toxic
 ```
 
 ## Citation