iteration

app.py

@@ -5,26 +5,38 @@ import torch
 from torch import nn
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
+option = st.selectbox("Select a toxicity analysis model:", ("RoBERTa", "DistilBERT", "XLM-RoBERTa"))
 defaultTxt = "I hate you cancerous insects so much"
-txt = st.text_area('Text to analyze', defaultTxt)
+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')
+# Load tokenizer and model weights, try to default to RoBERTa.
+match option:
+    case "RoBERTa":
+        tokenizerPath = "s-nlp/roberta_toxicity_classifier"
+        modelPath = "s-nlp/roberta_toxicity_classifier"
+    case "DistilBERT":
+        tokenizerPath = "citizenlab/distilbert-base-multilingual-cased-toxicity"
+        modelPath = "citizenlab/distilbert-base-multilingual-cased-toxicity"
+    case "XLM-RoBERTa":
+        tokenizerPath = "unitary/multilingual-toxic-xlm-roberta"
+        modelPath = "unitary/multilingual-toxic-xlm-roberta"
+    case _:
+        tokenizerPath = "s-nlp/roberta_toxicity_classifier"
+        modelPath = "s-nlp/roberta_toxicity_classifier"
+tokenizer = AutoTokenizer.from_pretrained(tokenizerPath)
+model = AutoModelForSequenceClassification.from_pretrained(modelPath)
 
 # run encoding through model to get classification output
-# e.g. "logits": tensor([[ 4.8982, -5.1952]], grad_fn=<AddmmBackward0>)
-result = model(batch)
+# RoBERTA: [0]: neutral, [1]: toxic
+encoding = tokenizer.encode(txt, return_tensors='pt')
+result = model(encoding)
 
 # 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]
 
-# default text input ought to return:
+# Expected returns from RoBERTa on default text:
 # Neutral: 0.0052
 # Toxic: 0.9948
 st.write("Classification Probabilities")