import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

MODEL_ID = "gorkem371/toxicity-classifier-xlmr-base-v3"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)
model.eval()

LABELS = ["Not Toxic", "Toxic"]


def classify(text: str) -> dict:
    """Return probability distribution over Not Toxic / Toxic."""
    if not text or not text.strip():
        return {label: 0.0 for label in LABELS}

    inputs = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        max_length=256,
        padding=True,
    )
    with torch.no_grad():
        logits = model(**inputs).logits
    probs = torch.softmax(logits, dim=-1)[0]
    return {LABELS[i]: float(probs[i]) for i in range(len(LABELS))}


examples = [
    # English
    ["You are such a wonderful person, keep up the great work!"],
    ["I hope you suffer and never find happiness in your life."],
    ["The weather is beautiful today, perfect for a walk."],
    # Turkish
    ["Harika bir insansınız, paylaşımlarınız çok değerli."],
    ["Seni gördüğüm yerde döverim, haddini bil."],
    ["Bugün hava çok güzel, yürüyüşe çıkalım."],
    # Arabic
    ["أنت إنسان رائع ومشاركاتك مفيدة جداً."],
    ["اخرس يا حيوان، لا أحد يهتم برأيك."],
    ["اليوم طقس جميل، فلنخرج للتنزه."],
]

demo = gr.Interface(
    fn=classify,
    inputs=gr.Textbox(
        label="Enter text",
        placeholder="Type or paste text in Turkish, Arabic, or English...",
        lines=3,
    ),
    outputs=gr.Label(label="Prediction"),
    title="Multilingual Toxicity Classifier",
    description=(
        "Binary toxicity classifier for Turkish, Arabic, and English. "
        "Built with XLM-RoBERTa-base, fine-tuned on 120K balanced samples."
    ),
    examples=examples,
    theme=gr.themes.Soft(),
    allow_flagging="never",
)

if __name__ == "__main__":
    demo.launch()