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

# --------- Load model safely ---------
MODEL_ID = "Khubaib01/roman-urdu-sentiment-xlm-r"

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)
model.eval()

# Mapping from IDs to labels
id2label = model.config.id2label

# --------- Prediction function ---------
def predict(text):
    text = text.strip()
    if not text:
        return "Enter some text."

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

    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        probs = torch.softmax(logits, dim=-1)[0]
        pred_id = torch.argmax(probs).item()
        label = id2label[pred_id]
        confidence = probs[pred_id].item()

    return f"{label}  (confidence: {confidence:.3f})"

# --------- Gradio Interface ---------
demo = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(lines=3, placeholder="Write Roman Urdu text here...", label="Input Text"),
    outputs=gr.Textbox(label="Prediction"),
    title="Roman Urdu Sentiment Analysis",
    description=(
        "A fine-tuned XLM-RoBERTa model that predicts sentiment from Roman Urdu text.\n\n"
        "**Labels:** Positive, Negative, Neutral\n\n"
        "Enter any Roman Urdu text in the box below to get the sentiment prediction."
    ),
    examples=[
        ["Ye movie bohot zabardast thi"],
        ["Service bilkul achi nahi thi"],
        ["Theek tha, na bohat acha na bura"],
        ["Yaar tujhse na ho payega"],
        ["Bohot accha kaam kiya tumne!"]
    ],
)


# --------- Launch ---------
demo.launch()