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)