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

REPO_ID = "Omarrran/koshur-diacritizer-byt5-small"

print("Loading model...")
tokenizer = AutoTokenizer.from_pretrained(REPO_ID)
model = AutoModelForSeq2SeqLM.from_pretrained(REPO_ID)
model.eval()
print("Model loaded.")


def diacritize(text: str, max_tokens: int) -> str:
    if not text or not text.strip():
        return ""
    inputs = tokenizer(text.strip(), return_tensors="pt", padding=True)
    with torch.no_grad():
        out = model.generate(**inputs, max_new_tokens=int(max_tokens))
    return tokenizer.decode(out[0], skip_special_tokens=True)


examples = [
    ["کاشر زبان", 256],
    ["میانی ہند", 256],
    ["سریںنگر شہر بوہت خوبصورت چھ", 256],
    ["اس کتاب منز بوہت ژھور معلومات چھ", 256],
    ["کشیر گرمی منز سبز تہ خوبصورت اوسان چھ", 256],
    ["زہ پرون شہر گوم", 256],
    ["امی گر کیاہ پیٹھ بنایو", 256],
]

description = """
## Koshur Diacritizer — ByT5-Small

This model restores **diacritical marks** (اِعراب) to undiacritized Kashmiri (کٲشُر) text written in Perso-Arabic script.

**Model:** [`Omarrran/koshur-diacritizer-byt5-small`](https://huggingface.co/Omarrran/koshur-diacritizer-byt5-small)

Enter raw Kashmiri text below or click an example to try instantly.
"""

demo = gr.Interface(
    fn=diacritize,
    inputs=[
        gr.Textbox(
            label="Input Text (undiacritized Kashmiri)",
            placeholder="یہاں کٲشُر متن لِکھِو…",
            lines=3,
            rtl=True,
        ),
        gr.Slider(
            minimum=64,
            maximum=512,
            value=256,
            step=32,
            label="Max New Tokens",
        ),
    ],
    outputs=gr.Textbox(
        label="Diacritized Output",
        lines=3,
        rtl=True,
        show_copy_button=True,
    ),
    examples=examples,
    title="کٲشُر ڈایاکرِٹایزر | Koshur Diacritizer",
    description=description,
    article="Built by [Omar Haq Nawaz Malik](https://huggingface.co/Omarrran) as part of the Kashmiri language AI infrastructure initiative.",
    theme=gr.themes.Soft(),
    cache_examples=False,
    flagging_mode="never",
)

if __name__ == "__main__":
    demo.launch()