Update app.py

app.py

@@ -1,12 +1,17 @@
 import gradio as gr
 from transformers import pipeline
 
-# Load the model
-MODEL_PATH = "unitary/toxic-bert"
-classifier = pipeline("text-classification", model=MODEL_PATH, tokenizer=MODEL_PATH)
+# Load the models
+MODEL_PATHS = {
+    "Unitary Toxic BERT": "unitary/toxic-bert",
+    "Martin HA Toxic Comment": "martin-ha/toxic-comment-model"
+}
 
-def predict_toxicity(text):
+classifiers = {name: pipeline("text-classification", model=path, tokenizer=path) for name, path in MODEL_PATHS.items()}
+
+def predict_toxicity(text, model_choice):
     # Get predictions
+    classifier = classifiers[model_choice]
     predictions = classifier(text, return_all_scores=True)[0]
     
     # Format results
@@ -19,16 +24,19 @@ def predict_toxicity(text):
 # Create the Gradio interface
 iface = gr.Interface(
     fn=predict_toxicity,
-    inputs=gr.Textbox(lines=5, label="Enter text to analyze"),
+    inputs=[
+        gr.Textbox(lines=5, label="Enter text to analyze"),
+        gr.Radio(choices=list(MODEL_PATHS.keys()), label="Choose a model", value="Toxic BERT")
+    ],
     outputs=gr.Label(num_top_classes=6, label="Toxicity Scores"),
     title="Toxicity Prediction",
-    description="This POC uses the model based onm toxic-bert to predict toxicity in text. Multi-class response.",
+    description="This POC uses trained&pre-trained models to predict toxicity in text. Choose between two models: 'Toxic BERT' based and 'Martin HA Toxic Comment'.",
     examples=[
-        ["Great game everyone!"],
-        ["You're such a noob, uninstall please."],
-        ["I hope you die in real life, loser."],
-        ["Nice move! How did you do that?"],
-        ["Go back to the kitchen where you belong."],
+        ["Great game everyone!", "Toxic BERT"],
+        ["You're such a noob, uninstall please.", "Martin HA Toxic Comment"],
+        ["I hope you die in real life, loser.", "Toxic BERT"],
+        ["Nice move! How did you do that?", "Martin HA Toxic Comment"],
+        ["Go back to the kitchen where you belong.", "Toxic BERT"],
     ]
 )