| from transformers import AutoTokenizer, AutoModelForCausalLM |
| import torch |
| import gradio as gr |
|
|
| model_name = "Qwen/Qwen2.5-0.5B-Instruct" |
|
|
| tokenizer = AutoTokenizer.from_pretrained(model_name) |
| model = AutoModelForCausalLM.from_pretrained( |
| model_name, |
| torch_dtype=torch.float16, |
| device_map="auto" |
| ) |
|
|
| def chat(user_input): |
| |
| messages = [ |
| {"role": "system", "content": "你是一个中文助理，回答要简短清晰，不超过50字。"}, |
| {"role": "user", "content": user_input} |
| ] |
| |
| |
| text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) |
| inputs = tokenizer(text, return_tensors="pt").to(model.device) |
|
|
| |
| outputs = model.generate( |
| **inputs, |
| max_new_tokens=50, |
| do_sample=True, |
| temperature=0.7, |
| ) |
| |
| |
| answer = tokenizer.decode(outputs[0], skip_special_tokens=True) |
| |
| |
| if "user" in answer: |
| answer = answer.split("user:")[-1] |
| |
| return answer.strip() |
|
|
| iface = gr.Interface( |
| fn=chat, |
| inputs=gr.Textbox(lines=3, label="输入你的问题"), |
| outputs=gr.Textbox(lines=10, label="回答"), |
| title="Qwen 0.5B 中文问答机器人" |
| ) |
|
|
| iface.launch() |
|
|
