# 使用transformers库部署Qwen3
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 加载模型
model_name = "Qwen/Qwen3-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto",
    trust_remote_code=True
)

# 对话函数
def chat_with_qwen3(message, history=[], thinking_mode=False):
    messages = history + [{"role": "user", "content": message}]

    # 添加思维模式提示
    if thinking_mode:
        system_msg = "请使用思维模式，展示你的推理过程。"
        messages.insert(0, {"role": "system", "content": system_msg})

    # 应用聊天模板
    input_ids = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt"
    ).to(model.device)

    # 生成回答
    with torch.no_grad():
        outputs = model.generate(
            input_ids,
            max_new_tokens=2000,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            pad_token_id=tokenizer.eos_token_id
        )

    response = tokenizer.decode(
        outputs[0][input_ids.shape[-1]:],
        skip_special_tokens=True
    )

    return response

# 使用示例
# 普通模式
response = chat_with_qwen3("请解释深度学习的基本概念")
print("普通模式:", response)

# 思维模式
response = chat_with_qwen3(
    "解决这个数学问题：如果一个数的平方等于它的两倍，这个数是多少？",
    thinking_mode=True
)
print("思维模式:", response)

# 部署MoE模型需要特殊配置
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 加载MoE模型
model_name = "Qwen/Qwen3-30B-A3B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# MoE模型需要更多内存和特殊配置
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto",
    trust_remote_code=True,
    # MoE特定配置
    load_in_8bit=False,  # MoE模型通常不建议使用8bit
    low_cpu_mem_usage=True
)

# MoE模型推理函数
def moe_inference(prompt, max_tokens=1000):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            # MoE特定参数
            use_cache=True,
            pad_token_id=tokenizer.eos_token_id
        )

    response = tokenizer.decode(
        outputs[0][inputs['input_ids'].shape[-1]:],
        skip_special_tokens=True
    )

    return response

# 使用示例
response = moe_inference("编写一个Python快速排序算法")
print(response)

# Qwen3 Agent集成示例
import json
import requests

class Qwen3Agent:
    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer
        self.tools = self._init_tools()

    def _init_tools(self):
        """初始化可用工具"""
        return {
            "web_search": self.web_search,
            "calculator": self.calculator,
            "code_executor": self.code_executor,
            "file_reader": self.file_reader
        }

    def web_search(self, query):
        """网络搜索工具"""
        # 模拟网络搜索
        return f"搜索结果：{query}的相关信息"

    def calculator(self, expression):
        """计算器工具"""
        try:
            result = eval(expression)
            return f"计算结果：{result}"
        except:
            return "计算错误"

    def code_executor(self, code):
        """代码执行工具"""
        try:
            # 安全的代码执行环境
            exec_globals = {"__builtins__": {}}
            exec(code, exec_globals)
            return "代码执行成功"
        except Exception as e:
            return f"代码执行错误：{str(e)}"

    def file_reader(self, filepath):
        """文件读取工具"""
        try:
            with open(filepath, 'r', encoding='utf-8') as f:
                content = f.read()[:1000]  # 限制读取长度
            return f"文件内容：{content}"
        except:
            return "文件读取失败"

    def process_request(self, user_input):
        """处理用户请求"""
        # 构建包含工具信息的提示
        tools_desc = json.dumps({
            name: func.__doc__ for name, func in self.tools.items()
        }, ensure_ascii=False, indent=2)

        system_prompt = f"""
        你是一个智能助手，可以使用以下工具：
        {tools_desc}

        当需要使用工具时，请按以下格式回答：
        <tool_call>
        {{"tool": "tool_name", "args": {{"param": "value"}}}}
        </tool_call>
        """

        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_input}
        ]

        response = chat_with_qwen3(user_input, [], thinking_mode=True)

        # 检查是否需要使用工具
        if "<tool_call>" in response:
            tool_result = self._execute_tool(response)
            # 将工具结果反馈给模型
            follow_up = f"工具执行结果：{tool_result}\n请基于这个结果回答用户的问题。"
            final_response = chat_with_qwen3(follow_up)
            return final_response

        return response

    def _execute_tool(self, response):
        """执行工具调用"""
        try:
            start = response.find("<tool_call>") + len("<tool_call>")
            end = response.find("</tool_call>")
            tool_call_str = response[start:end].strip()

            tool_call = json.loads(tool_call_str)
            tool_name = tool_call["tool"]
            args = tool_call.get("args", {})

            if tool_name in self.tools:
                return self.tools[tool_name](**args)
            else:
                return "未知工具"
        except:
            return "工具调用格式错误"

# 使用示例
agent = Qwen3Agent(model, tokenizer)
response = agent.process_request("帮我计算 15 * 23 + 7 的结果")
print(response)