Update app.py

app.py

@@ -6,58 +6,64 @@ import gradio as gr
 dataset = load_dataset("Koushim/processed-jigsaw-toxic-comments", split="train", streaming=True)
 
 # Sample examples
-low, medium, high = [], [], []
+green, yellow, red = [], [], []
 for example in dataset:
     score = example['toxicity']
     text = example['text']
-    if score < 0.3 and len(low) < 3:
-        low.append((text, score))
-    elif 0.3 <= score < 0.7 and len(medium) < 3:
-        medium.append((text, score))
-    elif score >= 0.7 and len(high) < 3:
-        high.append((text, score))
-    if len(low) == 3 and len(medium) == 3 and len(high) == 3:
+    if score < 0.3 and len(green) < 3:
+        green.append((text, score))
+    elif 0.3 <= score < 0.7 and len(yellow) < 3:
+        yellow.append((text, score))
+    elif score >= 0.7 and len(red) < 3:
+        red.append((text, score))
+    if len(green) == 3 and len(yellow) == 3 and len(red) == 3:
         break
 
 examples_html = f"""
-### 🧪 Examples of Toxicity Levels
-
-#### 🔷 Low Toxicity
-- {low[0][0]} (score: {low[0][1]:.2f})
-- {low[1][0]} (score: {low[1][1]:.2f})
-- {low[2][0]} (score: {low[2][1]:.2f})
-
-#### 🟠 Medium Toxicity
-- {medium[0][0]} (score: {medium[0][1]:.2f})
-- {medium[1][0]} (score: {medium[1][1]:.2f})
-- {medium[2][0]} (score: {medium[2][1]:.2f})
-
-#### 🔴 High Toxicity
-- {high[0][0]} (score: {high[0][1]:.2f})
-- {high[1][0]} (score: {high[1][1]:.2f})
-- {high[2][0]} (score: {high[2][1]:.2f})
+### 🥰 Examples: Is your partner a Green Flag or Red Flag?
+
+#### 💚 Green Flag (Wholesome vibes 🌸)
+- {green[0][0]} (toxicity: {green[0][1]:.2f})
+- {green[1][0]} (toxicity: {green[1][1]:.2f})
+- {green[2][0]} (toxicity: {green[2][1]:.2f})
+
+#### 🟡 Yellow Flag (Eh… watch out 👀)
+- {yellow[0][0]} (toxicity: {yellow[0][1]:.2f})
+- {yellow[1][0]} (toxicity: {yellow[1][1]:.2f})
+- {yellow[2][0]} (toxicity: {yellow[2][1]:.2f})
+
+#### ❤️ Red Flag (🚨 Run bestie, run! 🚨)
+- {red[0][0]} (toxicity: {red[0][1]:.2f})
+- {red[1][0]} (toxicity: {red[1][1]:.2f})
+- {red[2][0]} (toxicity: {red[2][1]:.2f})
 """
 
 # Load toxicity detection pipeline
 classifier = pipeline("text-classification", model="cardiffnlp/twitter-roberta-base-offensive", top_k=None)
 
-def predict_toxicity(text):
+def predict_flag(text):
     preds = classifier(text)[0]
-    results = []
+    score = 0.0
     for pred in preds:
-        label = pred['label']
-        score = pred['score']
-        results.append(f"**{label}**: {score:.2f}")
-    return "\n".join(results)
+        if pred['label'].lower() in ['toxic', 'offensive', 'abusive']:
+            score = pred['score']
+            break
+    # Decide flag
+    if score < 0.3:
+        return f"💚 **Green Flag!**\nNot toxic at all. Keep them! 🌷 (toxicity: {score:.2f})"
+    elif 0.3 <= score < 0.7:
+        return f"🟡 **Yellow Flag!**\nHmm… could be better. Watch out. 👀 (toxicity: {score:.2f})"
+    else:
+        return f"❤️ **Red Flag!**\n🚨 Yikes, that’s toxic! 🚨 (toxicity: {score:.2f})"
 
 with gr.Blocks() as demo:
-    gr.Markdown("# 🧹 Hate Speech & Toxicity Monitor")
-    gr.Markdown("This tool shows examples of toxic comments and lets you check your own text for toxicity using a Hugging Face model.")
+    gr.Markdown("# 💌 Green Flag or Red Flag?")
+    gr.Markdown("Ever wondered if your partner’s texts are a green flag 💚 or a 🚨 red flag? Paste their messages below and let AI judge. Just for fun 😉")
     gr.Markdown(examples_html)
 
-    inp = gr.Textbox(label="🔷 Enter your comment")
-    out = gr.Markdown(label="🔷 Toxicity Scores")
-    btn = gr.Button("Check Toxicity")
-    btn.click(fn=predict_toxicity, inputs=inp, outputs=out)
+    inp = gr.Textbox(label="📩 Paste your partner's message here")
+    out = gr.Markdown(label="🧪 Verdict")
+    btn = gr.Button("👀 Check Now")
+    btn.click(fn=predict_flag, inputs=inp, outputs=out)
 
 demo.launch()