import gradio as gr
from transformers import AutoTokenizer, T5ForConditionalGeneration

# HF Hub path config
model_path = "crossroderick/aramt5"

# Load model once
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = T5ForConditionalGeneration.from_pretrained(model_path)


def transliterate(text: str, dialect: str) -> str:
    """Transliterate Syriac text to Latin script."""
    if not text.strip():
        return ""

    prefix = "Syriac2EastLatin: " if dialect == "east" else "Syriac2WestLatin: "
    inputs = tokenizer(prefix + text, return_tensors="pt")
    outputs = model.generate(**inputs, num_beams=4)
    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return result.strip()


# Interface
with gr.Blocks() as demo:
    gr.Markdown("# AramT5")
    gr.Markdown("Syriac → Latin script transliterator")
    gr.Markdown("---")

    text_input = gr.Textbox(
        label="Syriac text",
        placeholder="ܫܠܡܐ",
        lines=4,
        rtl=True
    )
    with gr.Row():
        west_btn = gr.Button("West (Serto)")
        east_btn = gr.Button("East (Madnḥaya)")

    output = gr.Textbox(label="Latin output", lines=2)

    west_btn.click(fn=lambda t: transliterate(t, "west"), inputs=text_input, outputs=output)
    east_btn.click(fn=lambda t: transliterate(t, "east"), inputs=text_input, outputs=output)


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=2)
    demo.launch()