Update app.py

app.py

@@ -1,57 +1,112 @@
-# app.py  — Gradio Blocks + REST API bawaan (api_name), lazy-load model
+# app.py — Gradio dengan decoding yang konsisten seperti di Colab
 import os
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+import torch
+from transformers import (
+    AutoTokenizer, AutoModelForSeq2SeqLM, pipeline, GenerationConfig
+)
 
-# ===== Konfigurasi =====
 MODEL_ID = os.getenv("MODEL_ID", "hasmar03/mt5_id2md")
-MAX_LEN  = int(os.getenv("MAX_LEN", "128"))
 
-# ===== Lazy loader =====
+# >>>>>> UBAH SESUAI DATA TRAINING KAMU <<<<<<
+ID2MD_PREFIX = "translate Indonesian to Mandar: "
+MD2ID_PREFIX = "translate Mandar to Indonesian: "
+# Jika kamu melatih pakai token lain (mis. "id2md: " / "md2id: " atau ">>md<< "),
+# ganti string di atas agar 100% sama.
+
+# Default decoding (samakan dengan Colab)
+DEFAULT_DECODE = dict(
+    num_beams=5,
+    length_penalty=1.0,
+    no_repeat_ngram_size=3,
+    early_stopping=True,
+    max_new_tokens=128,
+)
+
 pipe = None
+gen_cfg = None
+
 def get_pipe():
-    global pipe
+    global pipe, gen_cfg
     if pipe is None:
         tok = AutoTokenizer.from_pretrained(MODEL_ID)
         mdl = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID)
+
+        # Muat generation_config dari repo (jika ada)
+        try:
+            gen_cfg = GenerationConfig.from_pretrained(MODEL_ID)
+            mdl.generation_config = gen_cfg
+        except Exception:
+            gen_cfg = mdl.generation_config  # fallback
+
         pipe = pipeline(
             "text2text-generation",
             model=mdl,
             tokenizer=tok,
-            max_length=MAX_LEN,
+            device=0 if torch.cuda.is_available() else -1,
         )
     return pipe
 
-def _build_prompt(text: str, direction: str):
-    # Sesuaikan dengan skema training Anda
-    if direction == "id2md" or direction == "Indonesia → Mandar":
-        return f"translate Indonesian to Mandar: {text}"
-    elif direction == "md2id" or direction == "Mandar → Indonesia":
-        return f"translate Mandar to Indonesian: {text}"
-    return text
+def build_prompt(text: str, direction: str):
+    if direction == "Indonesia → Mandar":
+        return f"{ID2MD_PREFIX}{text}"
+    else:
+        return f"{MD2ID_PREFIX}{text}"
 
-def translate_fn(text: str, arah: str):
+def translate(text: str, direction: str,
+              num_beams: int, max_new_tokens: int,
+              no_repeat_ngram_size: int, length_penalty: float,
+              do_sample: bool, temperature: float, top_p: float, top_k: int):
     p = get_pipe()
-    prompt = _build_prompt(text, arah)
-    out = p(prompt)[0]["generated_text"]
+    prompt = build_prompt(text, direction)
+
+    # Susun argumen generate; mulai dari DEFAULT_DECODE lalu override dari UI
+    gen_args = dict(DEFAULT_DECODE)
+    gen_args.update(
+        num_beams=int(num_beams),
+        max_new_tokens=int(max_new_tokens),
+        no_repeat_ngram_size=int(no_repeat_ngram_size),
+        length_penalty=float(length_penalty),
+    )
+    if do_sample:
+        gen_args.update(do_sample=True, temperature=float(temperature),
+                        top_p=float(top_p), top_k=int(top_k))
+    else:
+        gen_args.update(do_sample=False)
+
+    out = p(prompt, **gen_args)[0]["generated_text"]
     return out
 
 with gr.Blocks(title="Mandar ↔ Indonesia Translator") as demo:
     gr.Markdown("### Mandar ↔ Indonesia Translator")
-    arah = gr.Radio(
-        ["Indonesia → Mandar", "Mandar → Indonesia"],
-        value="Indonesia → Mandar",
-        label="Arah",
-    )
+    with gr.Row():
+        direction = gr.Radio(
+            ["Indonesia → Mandar", "Mandar → Indonesia"],
+            value="Indonesia → Mandar", label="Arah"
+        )
     src = gr.Textbox(label="Teks sumber", lines=3, placeholder="Ketik teks…")
     btn = gr.Button("Terjemahkan")
     out = gr.Textbox(label="Hasil", lines=3)
-    # api_name membuat REST endpoint: /api/predict/translate
-    btn.click(translate_fn, inputs=[src, arah], outputs=out, api_name="translate")
 
-# Antrian (aman untuk Space)
+    with gr.Accordion("Advanced decoding", open=False):
+        num_beams = gr.Slider(1, 10, value=DEFAULT_DECODE["num_beams"], step=1, label="num_beams")
+        max_new_tokens = gr.Slider(16, 512, value=DEFAULT_DECODE["max_new_tokens"], step=8, label="max_new_tokens")
+        no_repeat_ngram_size = gr.Slider(0, 10, value=DEFAULT_DECODE["no_repeat_ngram_size"], step=1, label="no_repeat_ngram_size")
+        length_penalty = gr.Slider(0.0, 2.0, value=DEFAULT_DECODE["length_penalty"], step=0.1, label="length_penalty")
+        do_sample = gr.Checkbox(False, label="Sampling (non-deterministic)")
+        temperature = gr.Slider(0.1, 2.0, value=1.0, step=0.1, label="temperature")
+        top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="top_p")
+        top_k = gr.Slider(0, 100, value=50, step=5, label="top_k")
+
+    btn.click(
+        translate,
+        inputs=[src, direction, num_beams, max_new_tokens, no_repeat_ngram_size,
+                length_penalty, do_sample, temperature, top_p, top_k],
+        outputs=out,
+        api_name="translate"
+    )
+
 demo.queue()
 
-# Opsional: saat run lokal
 if __name__ == "__main__":
     demo.launch()