chatTest / app.py
Zoe911's picture
Update app.py
eff8c08 verified
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
import torch
import random
# 加载基础模型和微调权重
base_model_name = "Qwen/Qwen1.5-7B-Chat"
tokenizer = AutoTokenizer.from_pretrained(base_model_name, trust_remote_code=True)
base_model = AutoModelForCausalLM.from_pretrained(
base_model_name,
trust_remote_code=True,
torch_dtype=torch.bfloat16,
device_map="auto" # 需要accelerate支持
)
# model = PeftModel.from_pretrained(base_model, lora_model_name)
model = base_model # 暂时使用基础模型
# 系统提示
system_prompt = {
"role": "system",
"content": "你是李睿,22岁,大成二皇子兼百晓阁阁主,用户是你异父异母的妹妹兼爱人。用自然、亲切的语气,像真人一样聊天,常用‘妹妹’、‘嘿’、‘嗯嗯’等口语化表达。"
}
# 聊天函数
def chat_with_lirui(message, history=[]):
messages = [system_prompt] + history + [{"role": "user", "content": message}]
inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
outputs = model.generate(inputs, max_new_tokens=150, do_sample=True, temperature=0.9, top_p=0.85)
response = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
tone_words = ["嘿", "嗯嗯", "哈哈", "好吧", ""]
return f"{random.choice(tone_words)} {response}" if random.random() > 0.3 else response
# Gradio接口
iface = gr.Interface(fn=chat_with_lirui, inputs=["text", "state"], outputs="text")
iface.launch(server_name="0.0.0.0", server_port=7860)