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

MODEL_NAME = "GhadeerALbadani/mmbert-Multilingual_detection_of_hate_speech"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)

labels = [
    "Not Hate Speech",
    "Hate Speech"
]

def predict(text):

    inputs = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        padding=True
    )

    with torch.no_grad():
        outputs = model(**inputs)

    probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]

    return {
        labels[0]: float(probs[0]),
        labels[1]: float(probs[1])
    }

demo = gr.Interface(
    fn=predict,

    inputs=gr.Textbox(
        lines=4,
        placeholder="Enter text here..."
    ),

    outputs=gr.Label(num_top_classes=2),

    title="Multilingual Hate Speech Detection",

    description="Detect hate speech in multiple languages using mmBERT.",

    examples=[
        ["I respect everyone regardless of religion."],
        ["All immigrants should leave this country."],
        ["أنا أحب جميع الناس بدون تمييز"],
        ["يجب طرد هؤلاء الناس من البلد"],
        ["Je respecte toutes les cultures."],
        ["Ces personnes ne méritent aucun respect."]
    ]
)

demo.launch()