nppmatt commited on
Commit
67d3612
1 Parent(s): dd26b28

add padding to xlm

Browse files
Files changed (1) hide show
  1. app.py +5 -2
app.py CHANGED
@@ -26,8 +26,8 @@ elif (option == "DistilBERT"):
26
  elif (option == "XLM-RoBERTa"):
27
  tokenizerPath = "unitary/multilingual-toxic-xlm-roberta"
28
  modelPath = "unitary/multilingual-toxic-xlm-roberta"
29
- neutralIndex = 0
30
- toxicIndex = 1
31
  else:
32
  tokenizerPath = "s-nlp/roberta_toxicity_classifier"
33
  modelPath = "s-nlp/roberta_toxicity_classifier"
@@ -42,6 +42,9 @@ model = AutoModelForSequenceClassification.from_pretrained(modelPath)
42
  encoding = tokenizer.encode(txt, return_tensors='pt')
43
  result = model(encoding)
44
  result
 
 
 
45
 
46
  # transform logit to get probabilities
47
  prediction = nn.functional.softmax(result.logits, dim=-1)
 
26
  elif (option == "XLM-RoBERTa"):
27
  tokenizerPath = "unitary/multilingual-toxic-xlm-roberta"
28
  modelPath = "unitary/multilingual-toxic-xlm-roberta"
29
+ neutralIndex = 1
30
+ toxicIndex = 0
31
  else:
32
  tokenizerPath = "s-nlp/roberta_toxicity_classifier"
33
  modelPath = "s-nlp/roberta_toxicity_classifier"
 
42
  encoding = tokenizer.encode(txt, return_tensors='pt')
43
  result = model(encoding)
44
  result
45
+ pad = (0, 1)
46
+ nn.functional.pad(result.logits, pad, "constant", 0)
47
+ result
48
 
49
  # transform logit to get probabilities
50
  prediction = nn.functional.softmax(result.logits, dim=-1)