Update app.py

app.py

@@ -5,6 +5,7 @@ import pandas as pd
 from scipy.special import softmax
 import torch
 
+
 def preprocess(text):
     new_text = []
     for t in text.split(" "):
@@ -13,19 +14,16 @@ def preprocess(text):
         new_text.append(t)
     return " ".join(new_text)
 
-# Define toxicity class labels
-toxicity_labels = ["toxic", "severe toxic", "obscene", "threat", "insult", "identity hate"]
 
 st.title("Toxicity Classification App")
 
 user_input = st.text_input("Input texts to analyze", "Cody Jiang is a fantastic student in CS-UY-4613!")
 
-# Define the models to choose from
 model_names = ['distilbert-base-uncased-finetuned-sst-2-english', 'bert-base-uncased', 'roberta-base', 'Codys-Finetuning-Language-Model']
 model_name = st.selectbox("Select a pretrained model", model_names)
 
 if model_name == "Codys-Finetuning-Language-Model":
-    model = transformers.BertForSequenceClassification.from_pretrained("./myModel/")
+    model = transformers.BertForSequenceClassification.from_pretrained("./model/")
     tokenizer = transformers.BertTokenizerFast.from_pretrained("bert-base-uncased")
 else:
     tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
@@ -38,10 +36,11 @@ if st.button("Submit"):
         output_val = model(**input_val)
         probabilities = torch.sigmoid(output_val.logits)
         result_list = probabilities.tolist()[0]
-        toxicity_class = toxicity_labels[np.argmax(result_list)]
-        toxicity_prob = np.max(result_list)
-        df_result = pd.DataFrame([[user_input, toxicity_class, toxicity_prob]], columns=["Tweet", "Toxicity Class", "Probability"])
-        st.table(df_result)
+        columns = ["Tweet", "Toxicity Class", "Probability"]
+        toxicity_class = ["toxic", "severe toxic", "obscene", "threat", "insult", "identity hate"]
+        result_df = pd.DataFrame(columns=columns)
+        result_df.loc[0] = [user_input, toxicity_class[result_list.index(max(result_list))], max(result_list)]
+        st.table(result_df)
 
     else:
         text = preprocess(user_input)
@@ -50,12 +49,17 @@ if st.button("Submit"):
         scores = output[0][0].detach().numpy()
         scores = softmax(scores)
 
+        ranking = np.argsort(scores)
+        ranking = ranking[::-1]
+        columns = ["Tweet", "Toxicity Class", "Probability"]
+        toxicity_class = config.id2label
         result_list = []
+        columns_list = []
         for i in range(scores.shape[0]):
-            l = toxicity_labels[i]
-            s = scores[i]
+            l = toxicity_class[ranking[i]]
+            s = scores[ranking[i]]
             result_list.append(s)
-        toxicity_class = toxicity_labels[np.argmax(result_list)]
-        toxicity_prob = np.max(result_list)
-        df_result = pd.DataFrame([[user_input, toxicity_class, toxicity_prob]], columns=["Tweet", "Toxicity Class", "Probability"])
-        st.table(df_result)
+            columns_list.append(l)
+        result_df = pd.DataFrame(columns=columns)
+        result_df.loc[0] = [user_input, toxicity_class[result_list.index(max(result_list))], max(result_list)]
+        st.table(result_df)