Update app.py

app.py

@@ -1,12 +1,28 @@
 import os
+import time
+from collections import defaultdict, deque
+from typing import Tuple
+
 import gradio as gr
 from huggingface_hub import InferenceClient
 
-# -------- Settings --------
+# =========================
+# CONFIG & AUTH
+# =========================
+# Lấy token từ biến môi trường (khuyên dùng: add tại Settings → Repository secrets)
+# Nếu cần hardcode để test nhanh (KHÔNG an toàn repo public):
+# HF_TOKEN = "hf_XXXXXXXXXXXXXXXXXXXXXXXX"  # <--- DÁN TOKEN Ở ĐÂY (không khuyến nghị)
+HF_TOKEN = os.getenv("HF_TOKEN", None)
+
+# Model mặc định cho UI (có thể đổi qua dropdown)
 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)
+# Model cố định cho API
+FIXED_MODEL = "tencent/Hunyuan-MT-7B-fp8"
+
+# =========================
+# NGÔN NGỮ & PROMPT
+# =========================
 LANGS = [
     ("Chinese (简体中文)", "zh"),
     ("Traditional Chinese (繁體中文)", "zh-Hant"),
@@ -47,29 +63,36 @@ LANGS = [
     ("Mongolian (Монгол)", "mn"),
     ("Uyghur (ئۇيغۇرچە)", "ug"),
 ]
-
+LABEL2CODE = {label: code for label, code in LANGS}
 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
+    Prompt template tham khảo từ model card:
+    - Nếu có tiếng Trung (zh/zh-Hant/yue) ở nguồn hoặc đích -> dùng template tiếng Trung
+    - Nếu không -> template tiếng Anh
     """
+    txt = (text or "").strip()
+    if not txt:
+        return ""
     if src_lang in ZH_CODES or tgt_lang in ZH_CODES:
-        # Template ZH <=> XX
-        return f"把下面的文本翻译成{tgt_lang}，不要额外解释。\n\n{text.strip()}"
+        return f"把下面的文本翻译成{tgt_lang}，不要额外解释。\n\n{txt}"
     else:
-        # Template XX <=> XX (không có ZH)
-        return f"Translate the following segment into {tgt_lang}, without additional explanation.\n\n{text.strip()}"
+        return f"Translate the following segment into {tgt_lang}, without additional explanation.\n\n{txt}"
+
 
+# =========================
+# INFERENCE HELPER
+# =========================
 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.
+    Gọi Hugging Face Serverless Inference API (text-generation).
+    Không cần GPU, phù hợp Space free.
     """
+    if not prompt:
+        return ""
     client = InferenceClient(token=HF_TOKEN)
-    # Tham số khuyến nghị từ model card
     try:
         out = client.text_generation(
             model=model,
@@ -79,32 +102,69 @@ def call_hf_inference(model: str, prompt: str) -> str:
             top_p=0.6,
             repetition_penalty=1.05,
             stream=False,
-            # truncate không bật để tránh cắt prompt
         )
-        return out.strip()
+        return (out or "").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):
+
+def translate(text: str, src_code: str, tgt_code: str, model_choice: str) -> str:
     if not text or not text.strip():
         return "Vui lòng nhập nội dung cần dịch."
-    if src == tgt:
+    if not src_code or not tgt_code:
+        return "Thiếu mã ngôn ngữ nguồn/đích."
+    if src_code == tgt_code:
         return text.strip()
+    prompt = build_prompt(src_code, tgt_code, text)
+    return call_hf_inference(model_choice, prompt)
+
+
+# =========================
+# RATE LIMIT THEO IP (IN-MEMORY)
+# =========================
+# SỬA TẠI ĐÂY nếu muốn: tối đa bao nhiêu request / mỗi cửa sổ thời gian
+RATE_WINDOW_SEC = int(os.getenv("RATE_WINDOW_SEC", "60"))  # ví dụ: 60 giây
+RATE_MAX_REQ = int(os.getenv("RATE_MAX_REQ", "10"))        # ví dụ: 10 request / IP / 60s
 
-    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
+_ip_buckets: dict[str, deque] = defaultdict(deque)
 
-def ui():
+
+def _rate_limited(request: gr.Request) -> Tuple[bool, str]:
+    """
+    Trả (ok, msg). ok=False nếu vượt ngưỡng.
+    Lưu ý: in-memory -> reset khi container restart.
+    """
+    try:
+        ip = (request.client.host if request and request.client else "unknown") or "unknown"
+    except Exception:
+        ip = "unknown"
+
+    now = time.time()
+    dq = _ip_buckets[ip]
+
+    # Bỏ các dấu thời gian ngoài cửa sổ
+    while dq and (now - dq[0] > RATE_WINDOW_SEC):
+        dq.popleft()
+
+    if len(dq) >= RATE_MAX_REQ:
+        wait_sec = max(1, int(RATE_WINDOW_SEC - (now - dq[0])))
+        return False, f"Bạn đã vượt giới hạn {RATE_MAX_REQ} yêu cầu / {RATE_WINDOW_SEC}s. Hãy thử lại sau ~{wait_sec}s."
+    dq.append(now)
+    return True, ""
+
+
+# =========================
+# UI (GRADIO BLOCKS)
+# =========================
+def build_ui() -> gr.Blocks:
     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 mô hình `tencent/Hunyuan-MT-7B-fp8` (khuyến nghị) hoặc `tencent/Hunyuan-MT-7B`.
             - 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.
+            > Mẹo: vào *Settings → Repository secrets* thêm `HF_TOKEN` để tăng hạn mức.
             """
         )
 
@@ -119,31 +179,97 @@ def ui():
             )
 
         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)
+            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"
+            )
 
         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 = gr.Button("Dịch", variant="primary")
+
+        def _on_translate(text, src_label, tgt_label, model_id, request: gr.Request = None):
+            ok, msg = _rate_limited(request)
+            if not ok:
+                return msg
+            src_code = LABEL2CODE[src_label]
+            tgt_code = LABEL2CODE[tgt_label]
+            return translate(text, src_code, tgt_code, model_id)
 
         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.
+            - Demo gọi **Serverless Inference API** nên tốc độ phụ thuộc hạn mức.
+            - Cần throughput cao hơn? Hãy cân nhắc GPU hoặc tự triển khai TGI/vLLM.
             """
         )
     return demo
 
+
+# =========================
+# API CỐ ĐỊNH (MODEL KHÓA)
+# =========================
+def api_translate_fixed(text: str, src_code: str, tgt_code: str, request: gr.Request = None) -> str:
+    """
+    API cho website: nhận mã ngôn ngữ (vd: 'en', 'vi', 'zh', ...),
+    dùng model cố định Hunyuan-MT-7B-fp8.
+    """
+    ok, msg = _rate_limited(request)
+    if not ok:
+        return msg
+    if not text or not text.strip():
+        return ""
+    if not src_code or not tgt_code:
+        return "Thiếu mã ngôn ngữ nguồn/đích."
+    if src_code == tgt_code:
+        return text.strip()
+    prompt = build_prompt(src_code, tgt_code, text)
+    return call_hf_inference(FIXED_MODEL, prompt)
+
+
+def build_api_interface() -> gr.Interface:
+    """
+    Tạo Interface riêng để có endpoint REST.
+    - Endpoint kiểu /run/<function_name> (nếu Gradio/HF hỗ trợ)
+    - Hoặc /run/predict với fn_index tương ứng
+    """
+    return gr.Interface(
+        fn=api_translate_fixed,
+        inputs=[
+            gr.Textbox(label="text"),
+            gr.Textbox(label="src_code"),
+            gr.Textbox(label="tgt_code"),
+        ],
+        outputs=gr.Textbox(label="translation"),
+        title="Hunyuan-MT Fixed API",
+        description="POST JSON tới endpoint để nhận bản dịch. Model cố định: tencent/Hunyuan-MT-7B-fp8."
+    )
+
+
+# =========================
+# MAIN
+# =========================
 if __name__ == "__main__":
-    ui().launch()
+    ui_app = build_ui()
+    api_iface = build_api_interface()
+
+    # G��p UI + API vào cùng server
+    demo = gr.TabbedInterface([ui_app, api_iface], tab_names=["App", "API"])
+
+    # Giới hạn đồng thời & hàng chờ (có thể chỉnh bằng biến môi trường)
+    # SỬA TẠI ĐÂY nếu muốn thay đổi:
+    CONCURRENCY = int(os.getenv("GRADIO_CONCURRENCY", "2"))  # ví dụ: 2 job chạy song song
+    QUEUE_MAX = int(os.getenv("GRADIO_QUEUE_MAX", "20"))     # ví dụ: 20 job có thể chờ
+
+    demo = demo.queue(concurrency_count=CONCURRENCY, max_size=QUEUE_MAX, status_update_rate=2)
+
+    # Bật REST API
+    demo.launch(enable_api=True)