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,
    )
    
    # 解码生成的文本，并去掉 "user" 和 "system" 的部分，只返回纯粹的回答
    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # 只保留模型的回答部分，可以根据需要进一步清理
    if "user" in answer:  # 如果模型回答中还包含“user”，去掉
        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()