1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3import json
45# 1. 加载模型和分词器6model_path ="your_model_path"# 替换为你的模型路径7tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)8model = AutoModelForCausalLM.from_pretrained(9 model_path,10 torch_dtype=torch.bfloat16,# 根据你的硬件调整11 device_map="auto"12)1314# 2. 定义你的工具集 (Tools)15tools =[16{17"name":"get_current_weather",18"description":"获取指定城市的实时天气信息",19"parameters":{20"type":"object",21"properties":{22"city":{23"type":"string",24"description":"城市名称,例如:北京、上海"25},26"unit":{27"type":"string",28"enum":["celsius","fahrenheit"],29"description":"温度单位"30}31},32"required":["city"]33}34},35{36"name":"send_email",37"description":"发送一封电子邮件",38"parameters":{39"type":"object",40"properties":{41"recipient":{42"type":"string",43"description":"收件人邮箱地址"44},45"subject":{46"type":"string",47"description":"邮件主题"48},49"body":{50"type":"string",51"description":"邮件正文内容"52}53},54"required":["recipient","subject","body"]55}56}57]5859# 3. 构建 Prompt60query ="帮我查一下北京今天的天气,用摄氏度显示"61system_prompt =f"You are a helpful assistant with access to the following tools. Use them if required to answer the user's query.\n{json.dumps(tools, indent=2)}"6263# 使用 Yi-Chat 模型的对话模板64messages =[65{"role":"system","content": system_prompt},66{"role":"user","content": query}67]6869# 将 messages 转换为模型期望的输入格式70prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)7172# 4. 模型推理73inputs = tokenizer(prompt, return_tensors="pt").to(model.device)74outputs = model.generate(75**inputs,76 max_new_tokens=256,77 eos_token_id=tokenizer.eos_token_id,# 根据你的 tokenizer 设置78 pad_token_id=tokenizer.pad_token_id if tokenizer.pad_token_id isnotNoneelse tokenizer.eos_token_id,79 do_sample=True,80 top_p=0.8,81 temperature=0.782)8384response_text = tokenizer.decode(outputs[0][len(inputs.input_ids[0]):], skip_special_tokens=True)8586print("--- Model Output ---")87print(response_text)8889# 5. 解析并执行工具调用90# !!! 警告:绝不要直接执行模型生成的代码或字符串。始终先进行解析和验证。91try:9293 tool_call_json = json.loads(response_text)9495 tool_name = tool_call_json.get("name")96 tool_args = tool_call_json.get("arguments",{})9798print(f"\n--- Tool Call Parsed ---")99print(f"Tool Name: {tool_name}")100print(f"Arguments: {tool_args}")101102# 在这里添加你的工具执行逻辑103# if tool_name == "get_current_weather":104# result = get_current_weather(**tool_args)105# ...106107except json.JSONDecodeError:108print("\n--- Final Answer (No Tool Call) ---")109print(response_text)110
Prompt 格式
为了触发工具调用,模型期望的输入遵循特定的格式。在微调期间,我们使用了包含系统指令的对话模板。
System Prompt: 包含一个引导指令和 JSON 格式的工具定义列表。
User Prompt: 用户的原始请求。
模板示例:
<|im_start|>system
You are a helpful assistant with access to the following tools. Use them if required to answer the user's query.
[
{
"name": "get_current_weather",
"description": "获取指定城市的实时天气信息",
"parameters": {
"type": "object",
"properties": {
"city": {
"type": "string",
"description": "城市名称,例如:北京、上海"
}
},
"required": ["city"]
}
}
]
<|im_end|>
<|im_start|>user
上海今天天气怎么样?<|im_end|>
<|im_start|>assistant