Revert "add suggestions cleanning"

This reverts commit a94f020eba50b5aec3746c327d42f6b8c6344ac1.

app.py

@@ -10,7 +10,6 @@ from termcolor import cprint
 
 # 初始化簡體到繁體轉換器
 cc = OpenCC('s2t')
-tokenizer = None
 
 # 可選模型列表
 MODEL_LIST = [
@@ -28,82 +27,10 @@ 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):
-    global tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    tok = AutoTokenizer.from_pretrained(model_name)
     mdl = AutoModelForCausalLM.from_pretrained(
         model_name, weights_only=False, trust_remote_code=True
     )
@@ -111,10 +38,10 @@ def get_pipeline(model_name):
         mdl.to("cuda")
     except Exception as e:
         print(f'Error: {e}')
-    return pipeline("text-generation", model=mdl, tokenizer=tokenizer, device=0)
+    return pipeline("text-generation", model=mdl, tokenizer=tok, device=0)
 
 @spaces.GPU
-def suggest_next(text, model_name, k, m, num_beam_groups, diversity_penalty, max_prefix_levels=2):
+def suggest_next(text, model_name, k, m, num_beam_groups, diversity_penalty):
     """
     使用 Diverse Beam Search 產生 m 條候選：
      - num_beams = m
@@ -131,7 +58,7 @@ def suggest_next(text, model_name, k, m, num_beam_groups, diversity_penalty, max
         "early_stopping": True,
     }
     if diversity_penalty and diversity_penalty > 0:
-        valid_group = max(gcd(m, num_beam_groups),2)
+        valid_group = gcd(m, num_beam_groups)
         gen_kwargs["num_beam_groups"] = valid_group
         gen_kwargs["diversity_penalty"] = float(diversity_penalty)
 
@@ -146,7 +73,6 @@ def suggest_next(text, model_name, k, m, num_beam_groups, diversity_penalty, max
         converted = cc.convert(snippet).strip()
         suggestions.add(converted)
     suggestions = list(suggestions)
-    suggestions = clean_suggestions(suggestions, max_prefix_levels)
 
     return update(choices=suggestions, value=None)
 
@@ -269,10 +195,6 @@ 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(
@@ -283,14 +205,13 @@ with gr.Blocks(css=custom_css) as demo:
             k_slider,
             m_slider,
             group_slider,
-            diversity_penalty_slider,
-            prefix_levels_slider  # 新增
+            diversity_penalty_slider
         ],
         outputs=suggestions,
     )
     input_text.change(
-        fn=lambda txt, mdl, k, m, g, d, auto, pl: (
-            suggest_next(txt, mdl, k, m, g, d, pl)
+        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=[
@@ -300,8 +221,7 @@ with gr.Blocks(css=custom_css) as demo:
             m_slider,
             group_slider,
             diversity_penalty_slider,
-            auto_predict,
-            prefix_levels_slider  # 新增
+            auto_predict
         ],
         outputs=suggestions,
     )