Spaces:
Running
on
Zero
Running
on
Zero
# 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 | |
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) | |
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() | |