| from transformers import AutoTokenizer, AutoModelForCausalLM |
| import torch |
| import gradio as gr |
|
|
| model_name = "tosei0000/chatbot" |
|
|
| tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) |
| model = AutoModelForCausalLM.from_pretrained( |
| model_name, |
| torch_dtype=torch.bfloat16, |
| trust_remote_code=True |
| ) |
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| model = model.to(device) |
|
|
| tokenizer.pad_token_id = tokenizer.eos_token_id |
| model.config.pad_token_id = tokenizer.eos_token_id |
|
|
| def chat(user_input, history): |
| prompt = "".join( |
| f"User: {u}\nAssistant: {a}\n" for u, a in history |
| ) + f"User: {user_input}\nAssistant:" |
| inputs = tokenizer(prompt, return_tensors="pt").to(device) |
| output = model.generate( |
| **inputs, |
| max_new_tokens=256, |
| do_sample=True, |
| temperature=0.7, |
| top_p=0.9, |
| pad_token_id=tokenizer.pad_token_id, |
| eos_token_id=tokenizer.eos_token_id |
| ) |
| text = tokenizer.decode(output[0], skip_special_tokens=True) |
| reply = text[len(prompt):].strip().split("\n")[0] |
| history.append((user_input, reply)) |
| return history, history |
|
|
| with gr.Blocks(title="Qwen2 Chatbot") as demo: |
| gr.Markdown("## 🤖 杜靖 聊天机器人") |
| chatbot = gr.Chatbot() |
| msg = gr.Textbox(label="输入你的问题") |
| clear = gr.Button("清除对话") |
| state = gr.State([]) |
|
|
| msg.submit(chat, [msg, state], [chatbot, state]) |
| clear.click(lambda: ([], []), None, [chatbot, state]) |
|
|
| if __name__ == "__main__": |
| demo.launch() |
|
|
|
|
|
|
| |
| |
|
|
| |
|
|
| |
| |
|
|
| |
| |
|
|
| |
| |
| |
| |
|
|
| |
| |
|
|
