app.py

import os
import gradio as gr
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():
            model_choice = gr.Dropdown(
                choices=[
                    "tencent/Hunyuan-MT-7B-fp8",
                    "tencent/Hunyuan-MT-7B",
                ],
                value=DEFAULT_MODEL,
                label="Model (Serverless)"
            )

        with gr.Row():
            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()