add suggestions cleanning

app.py

@@ -10,6 +10,7 @@ from termcolor import cprint
 
 # 初始化簡體到繁體轉換器
 cc = OpenCC('s2t')
+tokenizer = None
 
 # 可選模型列表
 MODEL_LIST = [
@@ -27,10 +28,82 @@ MODEL_LIST = [
     "Epiculous/Violet_Twilight-v0.2",
 ]
 
+def clean_suggestions(suggestions: list[str], max_levels: int) -> list[str]:
+    """
+    清洗建议列表：
+    1. 对每条建议用 tokenizer.tokenize 得到 token 序列。
+    2. 构建前缀树，将所有 token 序列插入。
+    3. 遍历前缀树，仅在深度 <= max_levels 且该节点有子节点时，提取对应 token 前缀。
+    4. 将这些 token 前缀转换回文本并去重，返回列表。
+    """
+    # 定义 Trie 节点结构
+    class TrieNode:
+        __slots__ = ("children", "count")
+        def __init__(self):
+            self.children: dict[str, TrieNode] = {}
+            self.count: int = 0  # 可以记录有多少序列经过此节点（可选）
+
+    # 构建前缀树
+    root = TrieNode()
+    token_seqs: list[list[str]] = []
+
+    for text in suggestions:
+        # tokenizer.tokenize 可能返回子词 token 列表
+        try:
+            toks = tokenizer.tokenize(text)
+        except Exception:
+            # 如果 tokenizer 不支持直接 tokenize raw text，可以先用 basic tokenization，如按空白分割
+            toks = text.split()
+        if not toks:
+            continue
+        token_seqs.append(toks)
+        node = root
+        node.count += 1
+        for tok in toks:
+            if tok not in node.children:
+                node.children[tok] = TrieNode()
+            node = node.children[tok]
+            node.count += 1
+
+    # 遍历 Trie，收集深度 <= max_levels 且有子节点的前缀序列
+    results_prefix_tokens: set[tuple[str, ...]] = set()
+
+    def dfs(node: TrieNode, path: list[str], depth: int):
+        # node: 当前 TrieNode； path: 已走过的 token 列表; depth: len(path)
+        if depth > max_levels:
+            return
+        # 如果当前节点有子节点，且 depth>0 （排除根节点本身），则为一个候选前缀
+        if depth > 0 and node.children:
+            results_prefix_tokens.add(tuple(path))
+        # 继续往下遍历，直到 depth == max_levels
+        if depth == max_levels:
+            return
+        for tok, child in node.children.items():
+            path.append(tok)
+            dfs(child, path, depth + 1)
+            path.pop()
+
+    dfs(root, [], 0)
+
+    # 将 token 前缀转换回字符串
+    cleaned: set[str] = set()
+    for tok_prefix in results_prefix_tokens:
+        try:
+            # tokenizer.convert_tokens_to_string 在大多数 tokenizer 支持
+            text_pref = tokenizer.convert_tokens_to_string(list(tok_prefix)).strip()
+        except Exception:
+            # fallback: 直接拼接 token（可能需要根据 tokenizer 规范加空格或直接连起来）
+            text_pref = "".join(tok_prefix).strip()
+        if text_pref:
+            cleaned.add(text_pref)
+
+    # 返回去重之后的列表
+    return list(cleaned)
 
 @lru_cache(maxsize=8)
 def get_pipeline(model_name):
-    tok = AutoTokenizer.from_pretrained(model_name)
+    global tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
     mdl = AutoModelForCausalLM.from_pretrained(
         model_name, weights_only=False, trust_remote_code=True
     )
@@ -38,10 +111,10 @@ def get_pipeline(model_name):
         mdl.to("cuda")
     except Exception as e:
         print(f'Error: {e}')
-    return pipeline("text-generation", model=mdl, tokenizer=tok, device=0)
+    return pipeline("text-generation", model=mdl, tokenizer=tokenizer, device=0)
 
 @spaces.GPU
-def suggest_next(text, model_name, k, m, num_beam_groups, diversity_penalty):
+def suggest_next(text, model_name, k, m, num_beam_groups, diversity_penalty, max_prefix_levels=2):
     """
     使用 Diverse Beam Search 產生 m 條候選：
      - num_beams = m
@@ -58,7 +131,7 @@ def suggest_next(text, model_name, k, m, num_beam_groups, diversity_penalty):
         "early_stopping": True,
     }
     if diversity_penalty and diversity_penalty > 0:
-        valid_group = gcd(m, num_beam_groups)
+        valid_group = max(gcd(m, num_beam_groups),2)
         gen_kwargs["num_beam_groups"] = valid_group
         gen_kwargs["diversity_penalty"] = float(diversity_penalty)
 
@@ -73,6 +146,7 @@ def suggest_next(text, model_name, k, m, num_beam_groups, diversity_penalty):
         converted = cc.convert(snippet).strip()
         suggestions.add(converted)
     suggestions = list(suggestions)
+    suggestions = clean_suggestions(suggestions, max_prefix_levels)
 
     return update(choices=suggestions, value=None)
 
@@ -195,6 +269,10 @@ with gr.Blocks(css=custom_css) as demo:
             minimum=0.0, maximum=2.0, step=0.1, value=1.0,
             label="多樣性懲罰 (diversity_penalty)"
         )
+        prefix_levels_slider = gr.Slider(
+            minimum=1, maximum=5, step=1, value=2,
+            label="Clean 前綴深度 (max_levels)"
+        )
 
     # 綁定事件
     predict_button.click(
@@ -205,13 +283,14 @@ with gr.Blocks(css=custom_css) as demo:
             k_slider,
             m_slider,
             group_slider,
-            diversity_penalty_slider
+            diversity_penalty_slider,
+            prefix_levels_slider  # 新增
         ],
         outputs=suggestions,
     )
     input_text.change(
-        fn=lambda txt, mdl, k, m, g, d, auto: (
-            suggest_next(txt, mdl, k, m, g, d)
+        fn=lambda txt, mdl, k, m, g, d, auto, pl: (
+            suggest_next(txt, mdl, k, m, g, d, pl)
             if auto else update(choices=[], value=None)
         ),
         inputs=[
@@ -221,7 +300,8 @@ with gr.Blocks(css=custom_css) as demo:
             m_slider,
             group_slider,
             diversity_penalty_slider,
-            auto_predict
+            auto_predict,
+            prefix_levels_slider  # 新增
         ],
         outputs=suggestions,
     )