update

README.md

@@ -11,24 +11,29 @@ license: mit
 short_description: Chinese input method accelerator
 ---
 
-# 台灣中文下段預測（ZeroGPU + Gradio v5）
+# 台灣中文輸入法加速器（ZeroGPU + Gradio v5）
 
 ## 一、專案概述  
 本示範結合多種小型中文語言模型，並透過 Hugging Face 的 **ZeroGPU**（H200）即時執行文字生成，模擬中文輸入法中的候選詞建議功能。
 
 ## 二、主要功能  
-1. **輸入文字區**：在此輸入您正在撰寫的中文片段。  
-2. **模型選擇**：可在七款預訓練模型間切換。  
-3. **參數滑桿**：  
-   - **K（最大新生成詞元）**：控制每次生成的字數上限。  
-   - **M（建議數量）**：控制同時產生的候選建議數量。  
-4. **使用 GPU 生成建議**：點擊後將在 H200 上啟動推理，並自動釋放資源。  
-5. **建議清單**：點選任一候選，該文字片段即會自動附加至輸入區。
+…  
+4. **使用 GPU 生成建議**：  
+   - 採用 **Beam Search**（`num_beams=M`）同時產出 M 條最可能的候選下段，並在 H200 上執行推理。  
+…  
 
 ## 三、運作原理  
-- **模型快取**：首次使用各模型時載入至 CPU 快取。  
-- **ZeroGPU 呼叫**：點擊「使用 GPU 生成建議」時，`@spaces.GPU` 裝飾的函式會自動分配 H200，完成推理後歸還。  
-- **即時拼接**：選擇建議即時追加，操作如同 IME 建議列。
+- 點擊「使用 GPU 生成建議」時，函式會以 **Beam Search** 模式呼叫模型：  
+  ```python
+  outs = gen_pipe(
+      text,
+      max_new_tokens=K,
+      num_beams=M,
+      num_return_sequences=M,
+      do_sample=False,
+      early_stopping=True
+  )
+
 
 ## 四、部署步驟  
 1. 在 Hugging Face Spaces 建立新 Space，框架選 **Gradio SDK**。  

app.py

@@ -4,6 +4,7 @@ import gradio as gr
 from functools import lru_cache
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
+# 可選模型列表
 MODEL_LIST = [
     "ckiplab/gpt2-tiny-chinese",
     "ckiplab/gpt2-base-chinese",
@@ -17,31 +18,50 @@ MODEL_LIST = [
 @lru_cache(maxsize=None)
 def get_pipeline(model_name):
     tok = AutoTokenizer.from_pretrained(model_name)
-    # By setting weights_only=False we bypass the torch.load(weights_only=True)
-    # path that is disallowed for torch<2.6 due to CVE-2025-32434 :contentReference[oaicite:1]{index=1}.
     mdl = AutoModelForCausalLM.from_pretrained(model_name, weights_only=False)
     mdl.to("cuda")
     return pipeline("text-generation", model=mdl, tokenizer=tok, device=0)
 
 @spaces.GPU
 def suggest_next(text, model_name, k, m):
-    outs = get_pipeline(model_name)(
-        text, max_new_tokens=k, num_return_sequences=m, do_sample=False
+    """
+    使用 Beam Search 產生 M 條最可能的下段建議（每條最多 K 個新詞元）。
+    """
+    gen_pipe = get_pipeline(model_name)
+    outs = gen_pipe(
+        text,
+        max_new_tokens=k,
+        num_beams=m,
+        num_return_sequences=m,
+        do_sample=False,
+        early_stopping=True
     )
+    # 只取掉 prompt 的部份
     return [out["generated_text"][len(text):] for out in outs]
 
 def append_suggestion(current, choice):
     return current + choice
 
 with gr.Blocks() as demo:
-    gr.Markdown("## 🇹🇼 台灣中文下段預測（ZeroGPU + Gradio v5）")
+    gr.Markdown(
+        "## 🇹🇼 台灣中文下段預測\n"
+        "結合小型語言模型與 ZeroGPU，提供 Beam Search 風格的多條下段建議。"
+    )
 
-    input_text = gr.TextArea(label="輸入文字", lines=4, placeholder="請在此輸入起始片段…")
+    input_text = gr.TextArea(
+        label="輸入文字", lines=4, placeholder="請在此輸入起始片段…"
+    )
 
     with gr.Row():
-        model_selector = gr.Dropdown(MODEL_LIST, value=MODEL_LIST[0], label="選擇模型")
-        k_slider = gr.Slider(1, 50, value=5, label="K（最大新生成詞元）")
-        m_slider = gr.Slider(1, 10, value=5, label="M（建議數量）")
+        model_selector = gr.Dropdown(
+            MODEL_LIST, value=MODEL_LIST[0], label="選擇模型"
+        )
+        k_slider = gr.Slider(
+            minimum=1, maximum=50, step=1, value=5, label="K（最大新生成詞元）"
+        )
+        m_slider = gr.Slider(
+            minimum=1, maximum=10, step=1, value=5, label="M（建議數量 / Beam 數）"
+        )
 
     suggestions = gr.Dropdown([], label="建議清單", interactive=True)
     gpu_button = gr.Button("使用 GPU 生成建議")