Update app.py

app.py

@@ -1,122 +1,149 @@
-# app.py — HF Spaces Free (CPU), Hunyuan-MT 7B-fp8, đa ngôn ngữ, chia đoạn, UI + API
-import os, re
-from typing import List, Optional
-
+import os
 import gradio as gr
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
-# ===== Cấu hình =====
-DEFAULT_MODEL = "tencent/Hunyuan-MT-7B-fp8"  # đổi bằng env MODEL_NAME nếu muốn
-MODEL_NAME = os.getenv("MODEL_NAME", DEFAULT_MODEL)
-
-GEN_KW = dict(  # tham số sinh nhẹ cho CPU
-    max_new_tokens=256,
-    top_k=20,
-    top_p=0.6,
-    repetition_penalty=1.05,
-    temperature=0.7,
-    do_sample=True,
-)
-
-MAX_INPUT_TOKENS = int(os.getenv("MAX_INPUT_TOKENS", "800"))  # giới hạn input mỗi mảnh
-
-# ===== Load tokenizer & model (fp8 bằng dict quantization_config) =====
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
-quant_cfg = {"quantization_method": "fp8", "ignore": []}  # tránh lỗi ignore=None
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    trust_remote_code=True,
-    quantization_config=quant_cfg,
-)
-DEVICE = getattr(model, "device", torch.device("cpu"))
-
-# ===== Chuẩn hóa tên ngôn ngữ =====
-LANG_ALIASES = {
-    "vi": "Vietnamese", "vie": "Vietnamese", "vietnamese": "Vietnamese", "tiếng việt": "Vietnamese",
-    "zh": "Chinese", "chi": "Chinese", "zho": "Chinese", "chinese": "Chinese", "tiếng trung": "Chinese", "hán ngữ": "Chinese", "mandarin": "Chinese",
-    "en": "English", "eng": "English", "tiếng anh": "English", "english": "English",
-    "ja": "Japanese", "jpn": "Japanese", "tiếng nhật": "Japanese", "japanese": "Japanese",
-    "ko": "Korean", "kor": "Korean", "tiếng hàn": "Korean", "korean": "Korean",
-    "fr": "French", "fra": "French", "fre": "French", "tiếng pháp": "French", "french": "French",
-    "de": "German", "deu": "German", "ger": "German", "tiếng đức": "German", "german": "German",
-    "es": "Spanish", "spa": "Spanish", "tiếng tây ban nha": "Spanish", "spanish": "Spanish",
-    "th": "Thai", "tha": "Thai", "tiếng thái": "Thai", "thai": "Thai",
-    "id": "Indonesian", "ind": "Indonesian", "tiếng indonesia": "Indonesian", "indonesian": "Indonesian",
-    "ms": "Malay", "msa": "Malay", "tiếng malaysia": "Malay", "malay": "Malay",
-    "pt": "Portuguese", "por": "Portuguese", "tiếng bồ đào nha": "Portuguese", "portuguese": "Portuguese",
-    "ru": "Russian", "rus": "Russian", "tiếng nga": "Russian", "russian": "Russian",
-}
-LANG_CHOICES = sorted(set(LANG_ALIASES.values()))
-def norm_lang(s: Optional[str]) -> Optional[str]:
-    if not s: return None
-    k = s.strip().lower()
-    return LANG_ALIASES.get(k, s.strip())
-
-# ===== Chia văn bản theo token =====
-def chunk_by_tokens(text: str, max_tokens: int) -> List[str]:
-    text = text.strip()
-    if not text: return []
-    rough = re.split(r"(?<=[\.!?。！？])\s+", text)
-    chunks, buf = [], ""
-    def tok_len(s: str) -> int:
-        return tokenizer(s, add_special_tokens=False, return_length=True)["length"]
-    for part in rough:
-        cand = (buf + " " + part).strip() if buf else part
-        if tok_len(cand) <= max_tokens:
-            buf = cand
-        else:
-            if buf: chunks.append(buf); buf = ""
-            if tok_len(part) <= max_tokens:
-                buf = part
-            else:
-                ids = tokenizer(part, add_special_tokens=False)["input_ids"]
-                for i in range(0, len(ids), max_tokens):
-                    piece = tokenizer.decode(ids[i:i+max_tokens], skip_special_tokens=True)
-                    if piece.strip(): chunks.append(piece.strip())
-    if buf: chunks.append(buf)
-    return [c for c in chunks if c.strip()]
-
-# ===== Core translate (chat template) =====
-@torch.inference_mode()
-def translate_text(text: str, target_lang: str, source_lang: Optional[str]=None) -> str:
-    tgt = norm_lang(target_lang) or "Vietnamese"
-    src = norm_lang(source_lang)
-    sys_prompt = (f"Translate the following segment from {src} into {tgt}, without additional explanation."
-                  if src else
-                  f"Translate the following segment into {tgt}, without additional explanation.")
-    outs = []
-    for piece in chunk_by_tokens(text, MAX_INPUT_TOKENS):
-        msgs = [{"role":"user","content": f"{sys_prompt}\n\n{piece}"}]
-        inputs = tokenizer.apply_chat_template(msgs, tokenize=True, add_generation_prompt=False, return_tensors="pt")
-        out_ids = model.generate(inputs.to(DEVICE), **GEN_KW)
-        outs.append(tokenizer.decode(out_ids[0], skip_special_tokens=True).strip())
-    return "\n".join(outs).strip()
-
-def translate_batch(texts: List[str], target_lang: str, source_lang: Optional[str]=None) -> List[str]:
-    return [translate_text(t, target_lang, source_lang) for t in texts]
-
-# ===== Gradio UI + API =====
-with gr.Blocks() as demo:
-    gr.Markdown("## Hunyuan-MT 7B-fp8 — Multilingual Translation (HF Free CPU)\nChia đoạn theo token, UI + API (Gradio).")
-
-    with gr.Tab("Single"):
-        src = gr.Textbox(label="Văn bản nguồn", lines=10, placeholder="Dán văn bản cần dịch…")
+from huggingface_hub import InferenceClient
+
+# -------- Settings --------
+DEFAULT_MODEL = os.getenv("HYMT_MODEL", "tencent/Hunyuan-MT-7B-fp8")
+HF_TOKEN = os.getenv("HF_TOKEN", None)  # có thể để trống (ẩn danh, sẽ bị rate-limit)
+
+# Ngôn ngữ được model hỗ trợ (trích từ model card)
+LANGS = [
+    ("Chinese (简体中文)", "zh"),
+    ("Traditional Chinese (繁體中文)", "zh-Hant"),
+    ("Cantonese (粤语)", "yue"),
+    ("English (English)", "en"),
+    ("Vietnamese (Tiếng Việt)", "vi"),
+    ("Japanese (日本語)", "ja"),
+    ("Korean (한국어)", "ko"),
+    ("Thai (ไทย)", "th"),
+    ("French (Français)", "fr"),
+    ("Spanish (Español)", "es"),
+    ("Portuguese (Português)", "pt"),
+    ("Italian (Italiano)", "it"),
+    ("German (Deutsch)", "de"),
+    ("Russian (Русский)", "ru"),
+    ("Arabic (العربية)", "ar"),
+    ("Turkish (Türkçe)", "tr"),
+    ("Indonesian (Bahasa Indonesia)", "id"),
+    ("Malay (Bahasa Melayu)", "ms"),
+    ("Filipino (Filipino)", "tl"),
+    ("Hindi (हिन्दी)", "hi"),
+    ("Polish (Polski)", "pl"),
+    ("Czech (Čeština)", "cs"),
+    ("Dutch (Nederlands)", "nl"),
+    ("Khmer (ភាសាខ្មែរ)", "km"),
+    ("Burmese (မြန်မာ)", "my"),
+    ("Persian (فارسی)", "fa"),
+    ("Gujarati (ગુજરાતી)", "gu"),
+    ("Urdu (اردو)", "ur"),
+    ("Telugu (తెలుగు)", "te"),
+    ("Marathi (मराठी)", "mr"),
+    ("Hebrew (עברית)", "he"),
+    ("Bengali (বাংলা)", "bn"),
+    ("Tamil (தமிழ்)", "ta"),
+    ("Ukrainian (Українська)", "uk"),
+    ("Tibetan (བོད་ཡིག)", "bo"),
+    ("Kazakh (Қазақша)", "kk"),
+    ("Mongolian (Монгол)", "mn"),
+    ("Uyghur (ئۇيغۇرچە)", "ug"),
+]
+
+ZH_CODES = {"zh", "zh-Hant", "yue"}
+
+def build_prompt(src_lang: str, tgt_lang: str, text: str) -> str:
+    """
+    Theo gợi ý prompt trong model card:
+    - ZH <=> XX: dùng template tiếng Trung
+    - XX <=> XX (không có ZH): dùng template tiếng Anh
+    """
+    if src_lang in ZH_CODES or tgt_lang in ZH_CODES:
+        # Template ZH <=> XX
+        return f"把下面的文本翻译成{tgt_lang}，不要额外解释。\n\n{text.strip()}"
+    else:
+        # Template XX <=> XX (không có ZH)
+        return f"Translate the following segment into {tgt_lang}, without additional explanation.\n\n{text.strip()}"
+
+def call_hf_inference(model: str, prompt: str) -> str:
+    """
+    Gọi Serverless Inference API (text-generation).
+    Không cần GPU trên Space. Có thể dùng ẩn danh hoặc set HF_TOKEN trong Secrets.
+    """
+    client = InferenceClient(token=HF_TOKEN)
+    # Tham số khuyến nghị từ model card
+    try:
+        out = client.text_generation(
+            model=model,
+            prompt=prompt,
+            max_new_tokens=512,
+            temperature=0.7,
+            top_p=0.6,
+            repetition_penalty=1.05,
+            stream=False,
+            # truncate không bật để tránh cắt prompt
+        )
+        return out.strip()
+    except Exception as e:
+        return f"[Lỗi] Không thể gọi Inference API: {e}"
+
+def translate(text: str, src: str, tgt: str, model_choice: str):
+    if not text or not text.strip():
+        return "Vui lòng nhập nội dung cần dịch."
+    if src == tgt:
+        return text.strip()
+
+    prompt = build_prompt(src, tgt, text)
+    # Lưu ý: Hunyuan-MT là causal LM định hướng prompt, không yêu cầu định dạng chat đặc biệt
+    result = call_hf_inference(model_choice, prompt)
+    return result
+
+def ui():
+    with gr.Blocks(title="Hunyuan-MT Translation (HF Inference API)", fill_height=True) as demo:
+        gr.Markdown(
+            """
+            # Tencent Hunyuan-MT (Serverless)
+            Chạy trên **Hugging Face Space (CPU free)** bằng **Serverless Inference API**.
+            - Chọn mô hình `tencent/Hunyuan-MT-7B` hoặc `tencent/Hunyuan-MT-7B-fp8`.
+            - Chọn ngôn ngữ nguồn/đích rồi bấm **Dịch**.
+            > Gợi ý: vào *Settings → Repository secrets* thêm `HF_TOKEN` để tăng hạn mức.
+            """
+        )
+
         with gr.Row():
-            src_lang = gr.Textbox(label="Ngôn ngữ nguồn (tùy chọn)", placeholder="Ví dụ: Vietnamese/Chinese/English…")
-            tgt_lang = gr.Dropdown(label="Ngôn ngữ đích", choices=LANG_CHOICES, value="Vietnamese")
-        out = gr.Textbox(label="Bản dịch", lines=10)
-        gr.Button("Dịch").click(translate_text, inputs=[src, tgt_lang, src_lang], outputs=out, api_name="translate_text")
+            model_choice = gr.Dropdown(
+                choices=[
+                    "tencent/Hunyuan-MT-7B-fp8",
+                    "tencent/Hunyuan-MT-7B",
+                ],
+                value=DEFAULT_MODEL,
+                label="Model (Serverless)"
+            )
 
-    with gr.Tab("Batch"):
-        src_list = gr.Textbox(label="Mỗi dòng 1 câu/đoạn", lines=10)
         with gr.Row():
-            src_lang_b = gr.Textbox(label="Ngôn ngữ nguồn (tùy chọn)")
-            tgt_lang_b = gr.Dropdown(label="Ngôn ngữ đích", choices=LANG_CHOICES, value="Vietnamese")
-        out_list = gr.Textbox(label="Kết quả (mỗi dòng tương ứng 1 đầu vào)", lines=10)
-        def _batch(txts_raw: str, tgt: str, src_: Optional[str]):
-            texts = [x for x in txts_raw.splitlines() if x.strip()]
-            return "\n".join(translate_batch(texts, tgt, src_))
-        gr.Button("Dịch Batch").click(_batch, inputs=[src_list, tgt_lang_b, src_lang_b], outputs=out_list, api_name="translate_batch")
-
-demo.queue(concurrency_count=1, max_size=2).launch()
+            src = gr.Dropdown(choices=[l for l, _ in LANGS], value="English (English)", label="Nguồn")
+            tgt = gr.Dropdown(choices=[l for l, _ in LANGS], value="Vietnamese (Tiếng Việt)", label="Đích")
+
+        # Map label -> code cho back-end
+        label2code = {label: code for label, code in LANGS}
+
+        def _on_translate(text, src_label, tgt_label, model_id):
+            src_code = label2code[src_label]
+            tgt_code = label2code[tgt_label]
+            return translate(text, src_code, tgt_code, model_id)
+
+        inp = gr.Textbox(label="Nội dung cần dịch", lines=8, placeholder="Nhập văn bản…")
+        btn = gr.Button("Dịch", variant="primary")
+        out = gr.Textbox(label="Kết quả", lines=8)
+
+        btn.click(_on_translate, [inp, src, tgt, model_choice], [out])
+
+        gr.Markdown(
+            """
+            #### Lưu ý
+            - Đây là demo qua **Serverless Inference API** nên tốc độ/phản hồi phụ thuộc hạn mức serverless.
+            - Với lượng lớn/nhanh hơn, hãy nâng cấp phần cứng (GPU) hoặc tự triển khai TGI/vLLM.
+            """
+        )
+    return demo
+
+if __name__ == "__main__":
+    ui().launch()