update indices

app.py

@@ -15,15 +15,21 @@ txt = st.text_area("Text to analyze", defaultTxt)
 if (option == "RoBERTa"):
     tokenizerPath = "s-nlp/roberta_toxicity_classifier"
     modelPath = "s-nlp/roberta_toxicity_classifier"
+    neutralIndex = 0
+    toxicIndex = 1
 elif (option == "DistilBERT"):
     tokenizerPath = "citizenlab/distilbert-base-multilingual-cased-toxicity"
     modelPath = "citizenlab/distilbert-base-multilingual-cased-toxicity"
+    neutralIndex = 1
+    toxicIndex = 0
 elif (option == "XLM-RoBERTa"):
     tokenizerPath = "unitary/multilingual-toxic-xlm-roberta"
     modelPath = "unitary/multilingual-toxic-xlm-roberta"
 else:
     tokenizerPath = "s-nlp/roberta_toxicity_classifier"
     modelPath = "s-nlp/roberta_toxicity_classifier"
+    neutralIndex = 0
+    toxicIndex = 1
 
 tokenizer = AutoTokenizer.from_pretrained(tokenizerPath)
 model = AutoModelForSequenceClassification.from_pretrained(modelPath)
@@ -35,8 +41,8 @@ result = model(encoding)
 
 # transform logit to get probabilities
 prediction = nn.functional.softmax(result.logits, dim=-1)
-neutralProb = prediction.data[0][0]
-toxicProb = prediction.data[0][1]
+neutralProb = prediction.data[0][neutralIndex]
+toxicProb = prediction.data[0][toxicIndex]
 
 # Expected returns from RoBERTa on default text:
 # Neutral: 0.0052