import os
# 设置环境变量禁用 TensorFlow
os.environ["USE_TORCH"] = "1"
os.environ["USE_TF"] = "0"

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

def load_model(model_path):
    """
    加载模型和分词器
    """
    print(f"正在从 {model_path} 加载模型...")
    
    # 检查 CUDA 是否可用
    if torch.cuda.is_available():
        device = torch.device("cuda")
        print("使用 CUDA 后端加速")
    # 检查 MPS 是否可用
    elif torch.backends.mps.is_available():
        device = torch.device("mps")
        print("使用 MPS 后端加速")
    else:
        device = torch.device("cpu")
        print("CUDA 和 MPS 均不可用，使用 CPU")
    
    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        torch_dtype=torch.float16,
        trust_remote_code=True
    ).to(device)
    
    print("模型加载完成！")
    return model, tokenizer, device

def chat_with_model(model, tokenizer, device):
    """
    与模型进行对话
    """
    history = []
    print("开始对话，输入 'exit' 退出")
    
    while True:
        user_input = input("\n用户: ")
        if user_input.lower() == 'exit':
            print("对话结束")
            break
        
        # 为 Qwen2 模型构建对话格式
        if not history:
            messages = [{"role": "user", "content": user_input}]
        else:
            messages = []
            for i, (user_msg, assistant_msg) in enumerate(history):
                messages.append({"role": "user", "content": user_msg})
                messages.append({"role": "assistant", "content": assistant_msg})
            messages.append({"role": "user", "content": user_input})
        
        # 使用 tokenizer 处理对话
        inputs = tokenizer.apply_chat_template(
            messages, 
            return_tensors="pt"
        ).to(device)
        
        # 生成回复
        outputs = model.generate(
            inputs,
            max_new_tokens=2048,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
        )
        
        # 解码回复
        response = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
        print(f"\n助手: {response}")
        
        # 更新历史记录
        history.append((user_input, response))

if __name__ == "__main__":
    model_path = "../qwen-medical"
    model, tokenizer, device = load_model(model_path)
    chat_with_model(model, tokenizer, device)