| from transformers import pipeline | |
| import gradio as gr | |
| generator = pipeline("text-generation", model="ckiplab/gpt2-base-chinese", | |
| tokenizer="ckiplab/gpt2-base-chinese") | |
| def chat_fn(message, history): | |
| history = history or [] | |
| input_text = "\n".join(history + [f"你: {message}", "AI:"]) | |
| output = generator(input_text, max_new_tokens=80, pad_token_id=0)[0]["generated_text"] | |
| response = output.split("AI:")[-1].strip().split("你:")[0].strip() | |
| history.append(f"你: {message}") | |
| history.append(f"AI: {response}") | |
| messages = [(history[i], history[i+1]) for i in range(0, len(history)-1, 2)] | |
| return messages, history | |
| with gr.Blocks() as demo: | |
| chatbot = gr.Chatbot(label="中文聊天機器人", type="tuples") | |
| state = gr.State([]) | |
| textbox = gr.Textbox(placeholder="請輸入訊息") | |
| textbox.submit(chat_fn, [textbox, state], [chatbot, state]) | |
| textbox.submit(lambda: "", None, textbox) | |
| demo.launch() | |