add divsere beam search and merge common prefixes

README.md

@@ -11,27 +11,16 @@ license: mit
 short_description: Chinese input method accelerator
 ---
 
----
-title: Input Method Acceleration
-emoji: 💻
-colorFrom: blue
-colorTo: pink
-sdk: gradio
-sdk_version: 5.34.0
-app_file: app.py
-pinned: true
-license: mit
-short_description: Chinese input method accelerator
----
-
 # 台灣中文輸入法加速器（ZeroGPU + Gradio v5）
 
+結合多種小型中文語言模型，並透過 Hugging Face 的 **ZeroGPU** 即時生成候選，模擬中文輸入法下段建議功能。
+
 ## 一、專案概述  
-本示範結合多種小型中文語言模型，並透過 Hugging Face 的 **ZeroGPU**（H200）即時執行文字生成，模擬中文輸入法中的候選詞建議功能。
+本示範結合多種小型中文語言模型，並透過 Hugging Face 的 **ZeroGPU** 即時執行文字生成，模擬中文輸入法中的候選詞建議功能。
 
 ## 二、主要功能  
 1. **使用 GPU 生成建議**  
-   - 採用 **Beam Search**（`num_beams=M`）同時產出 M 條最可能的候選，下段建議即時回饋。  
+   - 採用 **Beam Search**（`num_beams=M`），一次產出 M 條最可能的候選。  
 2. **經典 IME 候選欄樣式**  
    - 橫向滑動、點擊選字，模擬傳統輸入法使用體驗。  
 3. **自動延展文字框**  
@@ -42,24 +31,28 @@ short_description: Chinese input method accelerator
    - 手動／自動預測互補，按鈕常駐不隱藏。  
 6. **消除重複建議**  
    - 自動去重，保留首見建議，避免候選列表出現相同項目。  
+7. **多樣性設置**  
+   - 在「進階設定」中可調整 **Beam 群組數**（`num_beam_groups`）與 **多樣性懲罰**（`diversity_penalty`），以控制輸出建議的多樣性。
 
 ## 三、運作原理  
-- 點擊「預測」或於文字變更時（若已開啟自動預測），函式會以 **Beam Search** 模式呼叫模型，並做以下處理：  
-  1. 生成 m 條下文；  
+- 點擊「預測」或於文字變更時（若已開啟自動預測），函式會以 **Diverse Beam Search** 模式呼叫模型，並做以下處理：  
+  1. 生成 m 條下文（可透過 `num_beam_groups` 與 `diversity_penalty` 提高多樣性）；  
   2. 過濾空字串並轉繁體；  
   3. 去重並保留順序；  
-  4. 最後編號顯示於候選欄。  
+  4. 合併共享的前綴片段作為建議，提高候選精簡度；  
+  5. 更新候選列表。
 
 ## 四、部署步驟  
 1. 在 Hugging Face Spaces 建立新 Space，框架選 **Gradio SDK**。  
 2. 在 **Hardware** 欄位選擇 **Zero GPU**（需 PRO 帳號）。  
 3. 設定 Python 版本為 3.10.13。  
-4. 上傳 `app.py`、`requirements.txt`、更新後的 `README.md`。Spaces 會自動安裝相依套件。  
-5. 點擊 **Deploy**，即可線上體驗！  
+4. 上傳 `app.py`、`requirements.txt`、`README.md`，Spaces 會自動安裝相依套件。  
+5. 點擊 **Deploy** 即可線上體驗！
 
 ## 五、使用效益  
 - **快速回應**：利用 GPU burst 加速生成速度。  
+- **多樣化候選**：透過 Diverse Beam Search 與共同前綴合併，提供更具多樣性和精簡性的建議。  
 - **零門檻**：免本地 GPU，開發者與使用者皆可輕鬆試用。  
-- **多模型對比**：可快速測試不同模型品質與效能。  
+- **多模型對比**：可快速測試不同模型品質與效能。
 
 歡迎體驗高效、無痛的台灣中文輸入法下段建議 Demo！  

app.py

@@ -1,3 +1,4 @@
+# app.py
 import spaces
 import gradio as gr
 from gradio import update
@@ -24,7 +25,35 @@ MODEL_LIST = [
     "Epiculous/Violet_Twilight-v0.2",
 ]
 
-@lru_cache(maxsize=None)
+def merge_common_prefixes(suggestions, min_len=2):
+    """
+    合併具有共同前綴的建議：
+    - 找出所有長度 ≥ min_len 的共同前綴
+    - 將這些前綴作為新建議，移除原有被合併的項目
+    """
+    prefixes = []
+    to_remove = set()
+
+    for i in range(len(suggestions)):
+        for j in range(i+1, len(suggestions)):
+            s1, s2 = suggestions[i], suggestions[j]
+            # 計算字元級共同前綴
+            common = ''.join(c1 for c1, c2 in zip(s1, s2) if c1 == c2)
+            if len(common) >= min_len:
+                prefixes.append(common)
+                to_remove.update([s1, s2])
+
+    # 去重前綴
+    unique_prefixes = []
+    for p in prefixes:
+        if p not in unique_prefixes:
+            unique_prefixes.append(p)
+
+    # 剩下沒有被合併的建議
+    remainder = [s for s in suggestions if s not in to_remove]
+    return unique_prefixes + remainder
+
+@lru_cache(maxsize=8)
 def get_pipeline(model_name):
     tok = AutoTokenizer.from_pretrained(model_name)
     mdl = AutoModelForCausalLM.from_pretrained(
@@ -34,15 +63,20 @@ def get_pipeline(model_name):
     return pipeline("text-generation", model=mdl, tokenizer=tok, device=0)
 
 @spaces.GPU
-def suggest_next(text, model_name, k, m):
+def suggest_next(text, model_name, k, m, num_beam_groups, diversity_penalty):
     """
-    使用 Beam Search 產生 m 條候選，並一次更新候選列表，轉繁體並去除重複。
+    使用 Diverse Beam Search 產生 m 條候選：
+     - num_beams = m
+     - num_beam_groups, diversity_penalty 可調整多樣性
+    之後轉繁體、去重、合併共同前綴後回傳。
     """
     gen_pipe = get_pipeline(model_name)
     outs = gen_pipe(
         text,
         max_new_tokens=k,
         num_beams=m,
+        num_beam_groups=num_beam_groups,
+        diversity_penalty=diversity_penalty,
         num_return_sequences=m,
         do_sample=False,
         early_stopping=True
@@ -56,8 +90,11 @@ def suggest_next(text, model_name, k, m):
     for s in suggestions:
         if s not in unique_suggestions:
             unique_suggestions.append(s)
-    # 不編號，直接顯示文字
-    return update(choices=unique_suggestions, value=None)
+
+    # 合併共同前綴
+    final_suggestions = merge_common_prefixes(unique_suggestions, min_len=2)
+
+    return update(choices=final_suggestions, value=None)
 
 
 def append_suggestion(current, choice):
@@ -170,16 +207,42 @@ with gr.Blocks(css=custom_css) as demo:
         m_slider = gr.Slider(
             minimum=1, maximum=30, step=1, value=30, label="M（建議數/Beam 數）"
         )
+        group_slider = gr.Slider(
+            minimum=1, maximum=30, step=1, value=30,
+            label="Beam 群組數 (num_beam_groups)"
+        )
+        diversity_penalty_slider = gr.Slider(
+            minimum=0.0, maximum=2.0, step=0.1, value=1.0,
+            label="多樣性懲罰 (diversity_penalty)"
+        )
 
     # 綁定事件
     predict_button.click(
         fn=suggest_next,
-        inputs=[input_text, model_selector, k_slider, m_slider],
+        inputs=[
+            input_text,
+            model_selector,
+            k_slider,
+            m_slider,
+            group_slider,
+            diversity_penalty_slider
+        ],
         outputs=suggestions,
     )
     input_text.change(
-        fn=lambda txt, mdl, k, m, auto: suggest_next(txt, mdl, k, m) if auto else update(choices=[], value=None),
-        inputs=[input_text, model_selector, k_slider, m_slider, auto_predict],
+        fn=lambda txt, mdl, k, m, g, d, auto: (
+            suggest_next(txt, mdl, k, m, g, d)
+            if auto else update(choices=[], value=None)
+        ),
+        inputs=[
+            input_text,
+            model_selector,
+            k_slider,
+            m_slider,
+            group_slider,
+            diversity_penalty_slider,
+            auto_predict
+        ],
         outputs=suggestions,
     )
     suggestions.change(