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

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
LABELS = ["Not Hate", "Hate"]
MODEL_ID = "Harikrishna-Srinivasan/Hate-Speech-RoBERTa"

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


def predict(text: str):
    inputs = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        max_length=512,
        token_type_ids=False
    )
    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}

    with torch.inference_mode():
        logits = model(**inputs).logits
        probs = torch.softmax(logits, dim=-1)[0]

    not_hate, hate = probs[0].item(), probs[1].item()

    return {
        "Not Hate": not_hate,
        "Hate": hate
    }


gr.Interface(
    fn=predict,
    inputs=gr.Textbox(lines=3, placeholder="Enter the text"),
    outputs=gr.Label(num_top_classes=2),
    title="RoBERTa Hate Speech Classifier"
).launch()