tweak

app.py

@@ -13,13 +13,16 @@ tokenizer = AutoTokenizer.from_pretrained("s-nlp/roberta_toxicity_classifier")
 model = AutoModelForSequenceClassification.from_pretrained("s-nlp/roberta_toxicity_classifier")
 batch = tokenizer.encode(txt, return_tensors='pt')
 
-# e.g. "logits":"tensor([[ 4.8982, -5.1952]], grad_fn=<AddmmBackward0>)"
+# run model e.g. "logits": tensor([[ 4.8982, -5.1952]], grad_fn=<AddmmBackward0>)
 result = model(batch)
 
-# get probabilities
+# get probabilities e.g. tensor([[9.9996e-01, 4.2627e-05]], grad_fn=<SoftmaxBackward0>)
+# first indice is neutral, second is toxic
 prediction = nn.functional.softmax(result.logits, dim=-1)
 
-prediction
+neutralProb = round(prediction[0][0], 4)
+toxicProb = round(prediction[0][1], 4)
 
-#fig = px.bar(result, x="", y="", orientation='h')
-#fig.show()
+print("Classification Probabilities")
+print(f"Neutral: {neutralProb}")
+print(f"Toxic: {toxicProb}")