Update app.py

app.py

@@ -23,7 +23,8 @@ submit = False
 model_name = ""
 
 model_mapping = {
-    "Toxicity": "RobCaamano/toxicity",
+    "Toxicity - 1 Epoch": "RobCaamano/toxicity",
+    "Toxicity - 8 Epochs": "RobCaamano/toxicity_update",
     "DistilBERT Base Uncased (SST-2)": "distilbert-base-uncased-finetuned-sst-2-english",
 }
 
@@ -46,7 +47,7 @@ input = tokenizer(text, return_tensors="tf")
 if submit:
     results = dict(d.values() for d in clf(text)[0])
 
-    if model_name == "RobCaamano/toxicity":
+    if model_name in ["RobCaamano/toxicity", "RobCaamano/toxicity_update"]:
         classes = {k: results[k] for k in results.keys() if not k == "toxic"}
 
         max_class = max(classes, key=classes.get)