Luigi's picture
bugfix for diverse search with zero diversity_penalty
3350989
raw
history blame
5.43 kB
# app.py
import spaces
import gradio as gr
from gradio import update
from functools import lru_cache
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from opencc import OpenCC # 用於簡體轉繁體
# 初始化簡體到繁體轉換器
cc = OpenCC('s2t')
# 可選模型列表
MODEL_LIST = [
"liswei/Taiwan-ELM-270M",
"Mxode/SmolLM-Chinese-180M",
"flyingfishinwater/chinese-baby-llama2",
"unsloth/gemma-3-1b-pt",
"ckiplab/gpt2-tiny-chinese",
"ckiplab/gpt2-base-chinese",
"liswei/Taiwan-ELM-1_1B",
"benchang1110/Qwen2.5-Taiwan-1.5B-Instruct",
"benchang1110/Taiwan-tinyllama-v1.0-base",
"lianghsun/Llama-3.2-Taiwan-3B",
"twinkle-ai/Llama-3.2-3B-F1-Instruct",
"Epiculous/Violet_Twilight-v0.2",
]
def merge_common_prefixes(suggestions, min_len=2):
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(
model_name, weights_only=False, trust_remote_code=True
)
mdl.to("cuda")
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):
gen_pipe = get_pipeline(model_name)
# 構造 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:
unique_suggestions.append(s)
# 合併共同前綴
final_suggestions = merge_common_prefixes(unique_suggestions, min_len=2)
return update(choices=final_suggestions, value=None)
def append_suggestion(text, choice):
return text + choice
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; }
""") as demo:
with gr.Column():
suggestions = gr.Radio(
[], label="", interactive=True, type="value",
elem_id="suggestions-bar", elem_classes="candidate-list"
)
input_text = gr.Textbox(
label="", placeholder="請輸入拼音或文字…",
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")
with gr.Accordion("進階設定", open=False):
model_selector = gr.Dropdown(
MODEL_LIST, value=MODEL_LIST[0], label="模型"
)
k_slider = gr.Slider(
minimum=1, maximum=50, step=1, value=10, label="K(最大新詞元數)"
)
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,
group_slider,
diversity_penalty_slider
],
outputs=suggestions,
)
input_text.change(
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(
fn=append_suggestion,
inputs=[input_text, suggestions],
outputs=input_text,
)
demo.launch()