# 使用Hugging Face Transformers部署Mistral 7B
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 加载模型和分词器
model_name = "mistralai/Mistral-7B-Instruct-v0.1"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto"
)

# 对话函数
def chat_with_mistral(message, system_prompt="You are a helpful assistant."):
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": message}
    ]

    # 应用聊天模板
    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=1000,
            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_mistral("请解释什么是机器学习？")
print(response)

# 安装Ollama
curl -fsSL https://ollama.ai/install.sh | sh

# 下载并运行Mistral 7B
ollama pull mistral
ollama run mistral

# 在Python中使用Ollama API
import requests
import json

def ollama_chat(message):
    url = "http://localhost:11434/api/generate"
    data = {
        "model": "mistral",
        "prompt": message,
        "stream": False
    }

    response = requests.post(url, json=data)
    return response.json()["response"]

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

# Mistral 7B函数调用示例
import json

# 定义工具函数
def get_weather(location):
    """获取指定地点的天气信息"""
    # 模拟天气API调用
    return f"{location}的天气：晴天，温度25°C"

def calculate(expression):
    """计算数学表达式"""
    try:
        result = eval(expression)
        return f"计算结果：{result}"
    except:
        return "计算错误"

# 工具描述
tools = [
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "获取天气信息",
            "parameters": {
                "type": "object",
                "properties": {
                    "location": {
                        "type": "string",
                        "description": "地点名称"
                    }
                },
                "required": ["location"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "calculate",
            "description": "计算数学表达式",
            "parameters": {
                "type": "object",
                "properties": {
                    "expression": {
                        "type": "string",
                        "description": "数学表达式"
                    }
                },
                "required": ["expression"]
            }
        }
    }
]

# 函数调用处理
def process_function_call(message):
    # 构建包含工具信息的提示
    system_prompt = f"""
    你是一个有用的助手，可以调用以下工具：
    {json.dumps(tools, ensure_ascii=False, indent=2)}

    当需要使用工具时，请按以下格式回答：
    <function_call>
    {{"name": "function_name", "arguments": {{"param": "value"}}}}
    </function_call>
    """

    response = chat_with_mistral(message, system_prompt)

    # 检查是否包含函数调用
    if "<function_call>" in response:
        # 提取函数调用信息
        start = response.find("<function_call>") + len("<function_call>")
        end = response.find("</function_call>")
        function_call_str = response[start:end].strip()

        try:
            function_call = json.loads(function_call_str)
            function_name = function_call["name"]
            arguments = function_call["arguments"]

            # 执行函数
            if function_name == "get_weather":
                result = get_weather(arguments["location"])
            elif function_name == "calculate":
                result = calculate(arguments["expression"])
            else:
                result = "未知函数"

            return result
        except:
            return "函数调用格式错误"

    return response

# 使用示例
print(process_function_call("北京的天气怎么样？"))
print(process_function_call("计算 15 * 23 + 7"))