bugfix for diverse search with zero diversity_penalty

app.py

@@ -26,30 +26,19 @@ MODEL_LIST = [
 ]
 
 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
 
@@ -64,28 +53,29 @@ def get_pipeline(model_name):
 
 @spaces.GPU
 def suggest_next(text, model_name, k, m, num_beam_groups, diversity_penalty):
-    """
-    使用 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
-    )
-    # 提取新生成文本，過濾空字串，轉繁體
-    suggestions = [out["generated_text"][len(text):].strip() for out in outs]
-    suggestions = [s for s in suggestions if s]
-    suggestions = [cc.convert(s) for s in suggestions]
-    # 去除重複，保留順序
+    # 構造 generate 參數字典，僅在 penalty>0 時加入 diversity 相關
+    gen_kwargs = {
+        "max_new_tokens": k,
+        "num_beams": m,
+        "num_return_sequences": m,
+        "do_sample": False,
+        "early_stopping": True,
+    }
+    if diversity_penalty and diversity_penalty > 0:
+        gen_kwargs["num_beam_groups"] = num_beam_groups
+        gen_kwargs["diversity_penalty"] = diversity_penalty
+
+    outs = gen_pipe(text, **gen_kwargs)
+
+    # 提取純下文、過濾空字串、繁體化
+    suggestions = [
+        cc.convert(out["generated_text"][len(text):].strip())
+        for out in outs
+        if out["generated_text"][len(text):].strip()
+    ]
+
+    # 去重
     unique_suggestions = []
     for s in suggestions:
         if s not in unique_suggestions:
@@ -96,88 +86,18 @@ def suggest_next(text, model_name, k, m, num_beam_groups, diversity_penalty):
 
     return update(choices=final_suggestions, value=None)
 
+def append_suggestion(text, choice):
+    return text + choice
 
-def append_suggestion(current, choice):
-    if choice is None:
-        return current
-    # 直接插入選中的候選文字
-    return current + choice
-
-# 自訂 CSS：模擬經典中文輸入法候選欄樣式，並優化手機響應與自動高度
-custom_css = """
-#suggestions-bar {
-    width: 100%;
-    margin-bottom: 8px;
-}
+with gr.Blocks(css="""
+#suggestions-bar { width: 100%; margin-bottom: 8px; }
 #suggestions-bar .candidate-list {
-    display: flex;
-    gap: 8px;
-    background: #fff;
-    border: 1px solid #999;
-    border-radius: 4px;
-    padding: 6px;
-    overflow-x: auto;
-    white-space: nowrap;
-}
-#suggestions-bar .candidate-list label {
-    cursor: pointer;
-    padding: 6px 10px;
-    font-size: 16px;
+    display: flex; gap: 8px; background: #fff;
+    border: 1px solid #999; border-radius: 4px;
+    padding: 6px; overflow-x: auto; white-space: nowrap;
 }
-#suggestions-bar .candidate-list label:hover {
-    background: #f5f5f5;
-}
-#suggestions-bar .candidate-list input[type=radio]:checked + label {
-    background: #e6f7ff;
-    border: 1px solid #1890ff;
-}
-#input-box textarea {
-    width: 100%;
-    font-size: 16px;
-    padding: 6px;
-    box-sizing: border-box;
-    overflow: hidden;
-    resize: none;
-}
-#predict-button {
-    margin-top: 8px;
-    width: 100%;
-}
-/* 手機響應式 */
-@media only screen and (max-width: 600px) {
-    #suggestions-bar .candidate-list label {
-        padding: 8px;
-        font-size: 18px;
-    }
-    #predict-button {
-        font-size: 18px;
-    }
-}
-"""
-
-# 自動增高腳本
-auto_height_js = """
-<script>
-  window.addEventListener('load', () => {
-    const textarea = document.querySelector('#input-box textarea');
-    if (!textarea) return;
-    textarea.style.height = 'auto';
-    textarea.addEventListener('input', function() {
-      this.style.height = 'auto';
-      this.style.height = this.scrollHeight + 'px';
-    });
-  });
-</script>
-"""
-
-with gr.Blocks(css=custom_css) as demo:
-    gr.HTML(auto_height_js)
-    gr.Markdown(
-        "## 🇹🇼 繁體中文 IME 加速器  \
-"
-        "結合小型語言模型與 ZeroGPU，提供即時輸入法風格候選欄。"
-    )
-
+#suggestions-bar .candidate-list label { cursor: pointer; }
+""") as demo:
     with gr.Column():
         suggestions = gr.Radio(
             [], label="", interactive=True, type="value",
@@ -188,14 +108,11 @@ with gr.Blocks(css=custom_css) as demo:
             lines=1, max_lines=20, elem_id="input-box"
         )
 
-    # 永遠顯示預測按鈕
     with gr.Row():
         auto_predict = gr.Checkbox(
             value=True, label="自動預測（內容變更時觸發）", elem_id="auto-predict"
         )
-        predict_button = gr.Button(
-            "預測", elem_id="predict-button"
-        )
+        predict_button = gr.Button("預測", elem_id="predict-button")
 
     with gr.Accordion("進階設定", open=False):
         model_selector = gr.Dropdown(
@@ -216,7 +133,6 @@ with gr.Blocks(css=custom_css) as demo:
             label="多樣性懲罰 (diversity_penalty)"
         )
 
-    # 綁定事件
     predict_button.click(
         fn=suggest_next,
         inputs=[
@@ -251,4 +167,4 @@ with gr.Blocks(css=custom_css) as demo:
         outputs=input_text,
     )
 
-    demo.launch()
+    demo.launch()