Update app.py

app.py

@@ -25,6 +25,7 @@ model_name = ""
 model_mapping = {
     "Toxicity - 1 Epoch": "RobCaamano/toxicity",
     "Toxicity - 8 Epochs": "RobCaamano/toxicity_update",
+    "Toxicity - Weighted": "RobCaamano/toxicity_weighted",
     "DistilBERT Base Uncased (SST-2)": "distilbert-base-uncased-finetuned-sst-2-english",
 }
 
@@ -47,7 +48,7 @@ input = tokenizer(text, return_tensors="tf")
 if submit:
     results = dict(d.values() for d in clf(text)[0])
 
-    if model_name in ["RobCaamano/toxicity", "RobCaamano/toxicity_update"]:
+    if model_name in ["RobCaamano/toxicity", "RobCaamano/toxicity_update", "RobCaamano/toxicity_weighted"]:
         classes = {k: results[k] for k in results.keys() if not k == "toxic"}
 
         max_class = max(classes, key=classes.get)
@@ -77,4 +78,4 @@ if submit:
     st.table(result_df)
 
     expander = st.expander("View Raw output")
-    expander.write(results)
+    expander.write(results)