from datasets import load_dataset
from transformers import pipeline
import gradio as gr

# Load dataset
dataset = load_dataset("Koushim/processed-jigsaw-toxic-comments", split="train", streaming=True)

# Sample examples
green, yellow, red = [], [], []
for example in dataset:
    score = example['toxicity']
    text = example['text']
    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: 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_flag(text):
    preds = classifier(text)[0]
    score = 0.0
    for pred in preds:
        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("# 💌 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="📩 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()