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

MODELS = {
    "Moderate sentences (v1)": "Omarrran/koshur-diacritizer-byt5-small",
    "Longer sentences (v2)": "Omarrran/koshur-diacritizer-byt5-small-v2",
}

# Load both models once at startup and cache them
loaded = {}
for label, repo_id in MODELS.items():
    print(f"Loading {label} -> {repo_id} ...")
    tok = AutoTokenizer.from_pretrained(repo_id)
    mdl = AutoModelForSeq2SeqLM.from_pretrained(repo_id)
    mdl.eval()
    loaded[label] = (tok, mdl)
    print(f"  loaded {label}.")
print("All models loaded.")


def diacritize(text: str, model_choice: str, max_tokens: int) -> str:
    if not text or not text.strip():
        return ""
    tokenizer, model = loaded[model_choice]
    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 = [
    ["بہ چھس بتہ کھٮوان۔", "Moderate sentences (v1)", 256],
    [" تم چھ ٹی وی وچھان۔", "Moderate sentences (v1)", 256],
    ["اسۍ چھ کھیل گندان۔", "Moderate sentences (v1)", 256],
    ["کاشرۍ پلو چھ اکثر خطہ کس تاریخچ دلیل ونان، یس مختلف ثقافتن ہندۍ اثرات ظاہر کران ", "Longer sentences (v2)", 256],
    ["امہ سند مجموعہ چھ خاص تقریبن خاطر اکھ نفیس تہ شوخ شکل پیش کران۔", "Longer sentences (v2)", 256],
    ["امۍ سندس مجموعس منز چھ شاندار کپر تہ خوبصورت ڈیزاین شامل، یم خاص موقعن خاطر بہترین چھ۔", "Longer sentences (v2)", 256],
    ["انٹر کمیونٹی ڈاییلاگ: خاندرن اندۍ پکہ بین الکمیونٹی مکالمہ چھ مختلف ثقافتی گروپن درمیان تفہیم تہ تعاونس فروغ دیوان۔", "Longer sentences (v2)", 256],
]

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

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

**Choose a model:**
- **Moderate sentences (v1):** [`Omarrran/koshur-diacritizer-byt5-small`](https://huggingface.co/Omarrran/koshur-diacritizer-byt5-small)
- **Longer sentences (v2):** [`Omarrran/koshur-diacritizer-byt5-small-v2`](https://huggingface.co/Omarrran/koshur-diacritizer-byt5-small-v2)

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

Note: The Model is about 75.5% Correct. It can make mistakes. Hope To overcome these in next versions.
"""

demo = gr.Interface(
    fn=diacritize,
    inputs=[
        gr.Textbox(
            label="Input Text (undiacritized Kashmiri)",
            placeholder="یہاں کٲشُر متن لِکھِو…",
            lines=3,
            rtl=True,
        ),
        gr.Radio(
            choices=list(MODELS.keys()),
            value="Moderate sentences (v1)",
            label="Model",
            info="Pick v1 for moderate-length sentences, v2 for longer sentences.",
        ),
        gr.Slider(
            minimum=64,
            maximum=512,
            value=256,
            step=32,
            label="Max New Tokens",
        ),
    ],
    outputs=gr.Textbox(
        label="Diacritized Output",
        lines=3,
        rtl=True,
    ),
    examples=examples,
    title="کٲشُر ڈایاکرِٹایزر | Koshur Diacritizer",
    description=description,
    article="Built by [ 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()