import gradio as gr
import spaces
import torch
from transformers import AutoTokenizer, AutoModelForTokenClassification

model_id = "HaniaEmaan/code-switching-codesaviours-si26-hania"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForTokenClassification.from_pretrained(model_id)

@spaces.GPU
def classify_text(text):
    if not text.strip():
        return ""

    words = text.strip().split()

    tokenized = tokenizer(
        words,
        is_split_into_words=True,
        truncation=True,
        max_length=128,
        return_tensors="pt"
    )

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

    probs = torch.softmax(outputs.logits, dim=2)[0]
    predictions = torch.argmax(outputs.logits, dim=2)[0]
    word_ids = tokenized.word_ids(batch_index=0)

    id2label = model.config.id2label
    results = []
    seen = set()
    for idx, word_idx in enumerate(word_ids):
        if word_idx is None or word_idx in seen:
            continue
        seen.add(word_idx)
        label = id2label[predictions[idx].item()]
        score = probs[idx][predictions[idx]].item()
        results.append(f"Word: '{words[word_idx]}' | Label: {label} | Score: {score:.2f}")

    return "\n".join(results)

demo = gr.Interface(
    fn=classify_text,
    inputs=gr.Textbox(placeholder="Type here... e.g. you don't listen me. me tmse bohat tang hun"),
    outputs="textbox",
    title="Code-Switched Language Identification (ZeroGPU)",
    description="XLM-RoBERTa Token Classification for Roman Urdu and English — word-level tagging."
)
demo.launch()