tweak

app.py

@@ -1,6 +1,8 @@
 import streamlit as st
-import torch
 import plotly.express as px
+import torch
+
+from torch import nn
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
 deftxt = "I hate you cancerous insects so much"
@@ -9,13 +11,15 @@ txt = st.text_area('Text to analyze', deftxt)
 # load tokenizer and model weights
 tokenizer = AutoTokenizer.from_pretrained("s-nlp/roberta_toxicity_classifier")
 model = AutoModelForSequenceClassification.from_pretrained("s-nlp/roberta_toxicity_classifier")
-
-# prepare the input
 batch = tokenizer.encode(txt, return_tensors='pt')
 
-# inference
+# e.g. "logits":"tensor([[ 4.8982, -5.1952]], grad_fn=<AddmmBackward0>)"
 result = model(batch)
-result
+
+# get probabilities
+prediction = nn.functional.softmax(result.logits, dim=-1)
+
+print(prediction)
 
 #fig = px.bar(result, x="", y="", orientation='h')
 #fig.show()