Update app.py

app.py

@@ -17,7 +17,7 @@ model_name = ""
 with st.container():
     model_name = st.selectbox(
         "Select Model",
-        ("RobCaamano/toxicity",),
+        ("RobCaamano/toxicity", "distilbert-base-uncased-finetuned-sst-2-english"),
     )
     submit = st.button("Submit", type="primary")
 
@@ -31,21 +31,32 @@ input = tokenizer(text, return_tensors="tf")
 
 if submit:
     results = dict(d.values() for d in clf(text)[0])
-    classes = {k: results[k] for k in results.keys() if not k == "toxic"}
 
-    max_class = max(classes, key=classes.get)
-    probability = classes[max_class]
+    if model_name == "RobCaamano/toxicity":
+        classes = {k: results[k] for k in results.keys() if not k == "toxic"}
+
+        max_class = max(classes, key=classes.get)
+        probability = classes[max_class]
+
+        if results['toxic'] >= 0.5:
+            result_df = pd.DataFrame({
+                'Toxic': ['Yes'],
+                'Toxicity Class': [max_class],
+                'Probability': [probability]
+            })
+        else:
+            result_df = pd.DataFrame({
+                'Toxic': ['No'],
+                'Toxicity Class': 'This text is not toxic',
+            })
+
+    elif model_name == "distilbert-base-uncased-finetuned-sst-2-english":
+        result = max(results, key=results.get)
+        probability = results[result]
 
-    if results['toxic'] >= 0.5:
-        result_df = pd.DataFrame({
-            'Toxic': ['Yes'],
-            'Toxicity Class': [max_class],
-            'Probability': [probability]
-        })
-    else:
         result_df = pd.DataFrame({
-            'Toxic': ['No'],
-            'Toxicity Class': 'This text is not toxic',
+            'Result': [result],
+            'Probability': [probability],
         })
 
     st.table(result_df)