nppmatt commited on
Commit
c5be7aa
1 Parent(s): c125d59

update indices

Browse files
Files changed (1) hide show
  1. app.py +8 -2
app.py CHANGED
@@ -15,15 +15,21 @@ txt = st.text_area("Text to analyze", defaultTxt)
15
  if (option == "RoBERTa"):
16
  tokenizerPath = "s-nlp/roberta_toxicity_classifier"
17
  modelPath = "s-nlp/roberta_toxicity_classifier"
 
 
18
  elif (option == "DistilBERT"):
19
  tokenizerPath = "citizenlab/distilbert-base-multilingual-cased-toxicity"
20
  modelPath = "citizenlab/distilbert-base-multilingual-cased-toxicity"
 
 
21
  elif (option == "XLM-RoBERTa"):
22
  tokenizerPath = "unitary/multilingual-toxic-xlm-roberta"
23
  modelPath = "unitary/multilingual-toxic-xlm-roberta"
24
  else:
25
  tokenizerPath = "s-nlp/roberta_toxicity_classifier"
26
  modelPath = "s-nlp/roberta_toxicity_classifier"
 
 
27
 
28
  tokenizer = AutoTokenizer.from_pretrained(tokenizerPath)
29
  model = AutoModelForSequenceClassification.from_pretrained(modelPath)
@@ -35,8 +41,8 @@ result = model(encoding)
35
 
36
  # transform logit to get probabilities
37
  prediction = nn.functional.softmax(result.logits, dim=-1)
38
- neutralProb = prediction.data[0][0]
39
- toxicProb = prediction.data[0][1]
40
 
41
  # Expected returns from RoBERTa on default text:
42
  # Neutral: 0.0052
 
15
  if (option == "RoBERTa"):
16
  tokenizerPath = "s-nlp/roberta_toxicity_classifier"
17
  modelPath = "s-nlp/roberta_toxicity_classifier"
18
+ neutralIndex = 0
19
+ toxicIndex = 1
20
  elif (option == "DistilBERT"):
21
  tokenizerPath = "citizenlab/distilbert-base-multilingual-cased-toxicity"
22
  modelPath = "citizenlab/distilbert-base-multilingual-cased-toxicity"
23
+ neutralIndex = 1
24
+ toxicIndex = 0
25
  elif (option == "XLM-RoBERTa"):
26
  tokenizerPath = "unitary/multilingual-toxic-xlm-roberta"
27
  modelPath = "unitary/multilingual-toxic-xlm-roberta"
28
  else:
29
  tokenizerPath = "s-nlp/roberta_toxicity_classifier"
30
  modelPath = "s-nlp/roberta_toxicity_classifier"
31
+ neutralIndex = 0
32
+ toxicIndex = 1
33
 
34
  tokenizer = AutoTokenizer.from_pretrained(tokenizerPath)
35
  model = AutoModelForSequenceClassification.from_pretrained(modelPath)
 
41
 
42
  # transform logit to get probabilities
43
  prediction = nn.functional.softmax(result.logits, dim=-1)
44
+ neutralProb = prediction.data[0][neutralIndex]
45
+ toxicProb = prediction.data[0][toxicIndex]
46
 
47
  # Expected returns from RoBERTa on default text:
48
  # Neutral: 0.0052