Update app.py

app.py

@@ -29,25 +29,28 @@ if selected_model in ["Fine-tuned Toxicity Model"]:
     toxicity_classes = ["toxic", "severe_toxic", "obscene", "threat", "insult", "identity_hate"]
     model.config.id2label = {i: toxicity_classes[i] for i in range(model.config.num_labels)}
 
-def get_toxicity_class(predictions, threshold=0.3):
-    return {model.config.id2label[i]: pred for i, pred in enumerate(predictions) if pred >= threshold}
+def get_most_toxic_class(predictions):
+    return {model.config.id2label[i]: pred for i, pred in enumerate(predictions)}
 
 input = tokenizer(text, return_tensors="tf")
 
 if st.button("Submit", type="primary"):
     results = dict(d.values() for d in clf(text)[0])
-    toxic_labels = {k: results[k] for k in results.keys() if not k == "toxic"}
+    toxic_labels = get_most_toxic_class(results)
 
     tweet_portion = text[:50] + "..." if len(text) > 50 else text
 
     if len(toxic_labels) == 0:
         st.write("This text is not toxic.")
     else:
+        max_toxic_class = max(toxic_labels, key=toxic_labels.get)
+        max_probability = toxic_labels[max_toxic_class]
+
         df = pd.DataFrame(
             {
-                "Text (portion)": [tweet_portion] * len(toxic_labels),
-                "Toxicity Class": list(toxic_labels.keys()),
-                "Probability": list(toxic_labels.values()),
+                "Text (portion)": [tweet_portion],
+                "Toxicity Class": [max_toxic_class],
+                "Probability": [max_probability],
             }
         )
         st.table(df)