import gradio as gr
from functools import lru_cache
from transformers import pipeline

MODEL_ID = "hasmar03/mt5_id2md"

@lru_cache(maxsize=1)
def get_pipe():
    """
    Cache pipeline supaya model hanya dimuat 1x.
    Pakai CPU (device=-1) agar gratisan di Spaces jalan.
    """
    return pipeline(
        task="text2text-generation",
        model=MODEL_ID,
        tokenizer=MODEL_ID,
        device=-1,                 # CPU
        # framework & dtype otomatis oleh HF/Transformers
    )

def translate(direction: str, text: str, max_new_tokens: int = 64) -> str:
    if not text or not text.strip():
        return ""

    # Prefix sesuai format training mT5 kamu
    prompt = f"translate {direction}: {text}"

    # Setting decoding yang umum untuk terjemahan
    # - num_beams memperbaiki kualitas (sedikit lebih lambat tapi OK di CPU kecil)
    # - no_repeat_ngram_size mengurangi pengulangan
    # - early_stopping untuk hentikan beam kalau sudah cukup baik
    out = get_pipe()(
        prompt,
        max_new_tokens=int(max_new_tokens),
        num_beams=5,
        do_sample=False,
        no_repeat_ngram_size=3,
        early_stopping=True,
    )[0]["generated_text"]

    return out

with gr.Blocks(title="mT5 id↔md Translator (HF Space API)") as demo:
    gr.Markdown("# mT5 id↔md Translator (HF Space API)")

    with gr.Row():
        direction = gr.Dropdown(["id2md", "md2id"], value="id2md", label="Arah")
        max_tok   = gr.Slider(16, 128, value=64, step=1, label="max_new_tokens")

    inp = gr.Textbox(label="Teks sumber", lines=3, placeholder="Ketik kalimat...")
    out = gr.Textbox(label="Terjemahan", lines=3)
    btn = gr.Button("Terjemahkan", variant="primary")

    btn.click(translate, [direction, inp, max_tok], [out], api_name="translate")

    gr.Examples(
        examples=[
            ["id2md", "ia terus pulang begitu saja", 64],
            ["md2id", "tarrus i pole tia", 64],
        ],
        inputs=[direction, inp, max_tok],
        outputs=[out],
        label="Contoh cepat",
    )

# ==== Queue versi-agnostik (jalan di Gradio 4.x & 5.x) ====
try:
    # Gradio 5.x (beberapa build menolak argumen apa pun)
    demo.queue()
except TypeError:
    try:
        # Gradio 5.x lain (punya max_size)
        demo.queue(max_size=12)
    except TypeError:
        # Gradio 4.x (pakai concurrency_count)
        demo.queue(concurrency_count=1, max_size=12)

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