import spaces

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

MODEL_NAME = "122Uswa/code-switching-codesaviours-si26-uswa"

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForTokenClassification.from_pretrained(MODEL_NAME)

model.eval()


@spaces.GPU(duration=60)
def predict_language(text):
    if not text.strip():
        return "Please enter a sentence."

    words = text.split()

    # Tokenize words separately so word_ids() can map predictions
    inputs = tokenizer(
        words,
        is_split_into_words=True,
        return_tensors="pt",
        truncation=True
    )

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

    predictions = torch.argmax(outputs.logits, dim=2)[0].tolist()

    word_ids = tokenizer(
        words,
        is_split_into_words=True,
        truncation=True
    ).word_ids()

    results = []
    previous_word = None

    for prediction, word_id in zip(predictions, word_ids):

        if word_id is None or word_id == previous_word:
            continue

        word = words[word_id]
        label = model.config.id2label[prediction]

        results.append(f"{word} → {label}")

        previous_word = word_id

    return "\n".join(results)


demo = gr.Interface(
    fn=predict_language,
    inputs=gr.Textbox(
        label="Roman Urdu / English Sentence",
        placeholder="Yaar mujhe thora time chahiye for this assignment",
        lines=3
    ),
    outputs=gr.Textbox(
        label="Language Identification"
    ),
    title="Roman Urdu–English Language Identifier",
    description="Identifies each word as URD, ENG, or MIX."
)

demo.launch()