tweak

app.py

@@ -5,31 +5,28 @@ import torch
 from torch import nn
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-deftxt = "I hate you cancerous insects so much"
-txt = st.text_area('Text to analyze', deftxt)
+defaultTxt = "I hate you cancerous insects so much"
+txt = st.text_area('Text to analyze', defaultTxt)
 
 # load tokenizer and model weights
 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')
 
-# run model e.g. "logits": tensor([[ 4.8982, -5.1952]], grad_fn=<AddmmBackward0>)
+# run encoding through model to get classification output
+# e.g. "logits": tensor([[ 4.8982, -5.1952]], grad_fn=<AddmmBackward0>)
 result = model(batch)
 
-# get probabilities e.g. tensor([[9.9996e-01, 4.2627e-05]], grad_fn=<SoftmaxBackward0>)
+# transform logit to 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)
-
 neutralProb = prediction.data[0][0]
 toxicProb = prediction.data[0][1]
 
-neutralProb = torch.round(neutralProb, decimals=4)
-toxicProb = torch.round(toxicProb, decimals=4)
-
 # default text input ought to return:
 # Neutral: 0.0052
 # Toxic: 0.9948
-
 st.write("Classification Probabilities")
-st.write(f"{neutralProb:.4}  -  NEUTRAL")
-st.write(f"{toxicProb:.4}  -  TOXIC")
+st.write(f"{neutralProb:.4} - NEUTRAL")
+st.write(f"{toxicProb:.4} - TOXIC")