edit

app.py

@@ -1,6 +1,7 @@
 import numpy as np
 import streamlit as st
 from transformers import pipeline
+from transformers import DistilBertTokenizerFast, DistilBertForSequenceClassification
 
 import torch
 
@@ -36,19 +37,23 @@ def siebert(data):
     return label, score
 
 def finetuned(data):
-    specific_model = pipeline(model='dahongj/finetuned_toxictweets')
-    result = specific_model(data)
-    maxres = result[0]['label']
-    maxscore = result[0]['score']
-    sec = result[1]['label']
-    secscore = result[1]['score']
+    model_name = "dahongj/finetuned_toxictweets"
+    tokenizer = DistilBertTokenizerFast.from_pretrained(model_name)
+    model = DistilBertForSequenceClassification.from_pretrained(model_name)
+    tokenized_text = tokenizer(data, return_tensors="pt")
+    res = model(**tokenized_text)
+    mes = torch.sigmoid(res.logits)
+
+    Dict = {0: "toxic", 1: "severe_toxic", 2: "obscene", 3: "threat", 4: "insult", 5: "identity_hate"}
+
+    maxres, maxscore, sec, secscore = Dict[0], mes[0][0].item(), 0, 0
 
-    for i in result:
-        if i['score'] > secscore:
-            sec = i['label']
-            secscore = i['score']
+    for i in range(1,6):
+        if mes[0][i].item() > secscore:
+            sec = i
+            secscore = mes[0][i].item()
 
-    return maxres, maxscore, sec, secscore
+    return maxres, maxscore, Dict[sec], secscore
 
 def getSent(data, model):
     if(model == 'Bertweet'):