ToolGen：让大语言模型将工具API作为词汇直接生成与调用

1. 项目概述当大语言模型学会“认工具”在构建AI智能体的漫长探索中我们一直面临一个核心难题如何让大语言模型LLM精准地“想起”并“使用”成千上万的外部工具传统的做法无论是基于描述检索还是函数调用都像是在让一个只懂语言的专家去翻阅一本厚厚的、不断更新的工具百科全书过程笨拙且容易出错。检索可能不准调用格式可能不对整个流程充满了割裂感。今天要聊的ToolGen就是为解决这个问题而生的一项有趣工作。它来自ICLR 2025核心思想非常直观为什么不直接把工具“刻”进模型的“字典”里想象一下如果模型词汇表中不仅有“苹果”、“运行”这些词还有独一无二的、代表特定API的“工具词”比如WeatherAPIget_forecast那么当用户问“明天会下雨吗”时模型就能像生成一个普通单词那样直接“说出”这个工具词从而无缝触发对应的工具调用。这不仅仅是换个马甲那么简单。ToolGen通过一套名为“工具虚拟化”的技术将来自ToolBench的47,000多个工具API映射为模型词汇表中的特殊token。然后通过三个阶段的高效训练让模型首先“记住”所有这些工具token工具记忆再学会根据问题“选出”正确的工具token检索训练最后掌握完整使用工具链完成任务的能力端到端智能体调优。实测下来这种“生成即检索生成即调用”的统一范式在工具检索准确率和任务完成率上都展现出了显著优势。无论你是正在研究工具学习的研究者还是苦恼于如何为自家LLM注入强大工具使用能力的工程师ToolGen都提供了一个清晰、可复现且性能强大的新思路。2. 核心思路拆解为什么是“生成”而不是“检索”在深入代码之前我们有必要先厘清ToolGen设计哲学背后的“为什么”。这能帮助我们在后续复现或改进时做出更明智的决策。2.1 传统范式的瓶颈检索与生成的割裂主流的工具使用框架通常遵循“检索-生成”的流水线模式检索阶段根据用户查询从一个庞大的工具库常以向量数据库存储描述中召回最相关的几个工具。生成阶段LLM基于检索到的工具描述生成符合其调用规范的代码如JSON、Python函数。这个流程存在几个固有痛点误差累积检索阶段一旦出错工具不相关或描述不准后续生成再优秀也无济于事。上下文浪费工具描述有时很长需要被塞进有限的上下文窗口挤占了用于复杂推理的令牌。格式僵化模型需要额外学习严格的工具调用格式如tool_call(name‘xxx‘, arguments{...})增加了训练和推理的复杂度。2.2 ToolGen的破局点工具即词汇ToolGen的核心创新在于它跳出了“描述-检索”的框架将工具本身视为一种特殊的“语言单元”。虚拟化Virtualization每个工具API被编码成一个唯一的、人类不可读的虚拟令牌Virtual Token例如Spotifyget_playlist_tracks。这个令牌被直接添加到LLM的词汇表中。生成即调用当模型在推理过程中需要调用工具时它不再需要生成一段描述或代码而是直接生成这个虚拟令牌。下游系统如ToolGen框架捕获到这个特殊令牌后即可解析并执行对应的真实API调用。这样做带来了几个关键优势统一性工具检索和工具调用被统一为同一个“文本生成”任务。模型只需要做它最擅长的事——预测下一个token。效率一个虚拟token如API就能精准指向一个工具比携带冗长的工具描述要节省大量上下文空间。精度由于工具token在训练中被强制记忆和关联模型对其的“回忆”更类似于对事实知识的回忆而非基于相似度的模糊匹配从而提高了检索的精确度。可扩展性理论上只要词汇表允许可以持续添加新的工具token。训练新工具时类似于让模型学习一个新“单词”及其用法。注意这种方法的成功依赖于一个关键前提——模型在预训练阶段已经具备了强大的语言理解和生成能力。ToolGen的工作是在此基础上为其扩充一套“工具方言”的词汇和语法而不是从零开始教它语言。2.3 三阶段训练循序渐进的“教学”过程直接将4.7万个新词扔给模型是行不通的。ToolGen采用了精心设计的三阶段课程学习策略工具记忆Tool Memorization目标让模型“认识”每一个工具token。输入是工具的名称和描述输出就是对应的虚拟token。这就像在教学生背单词表确保它看到“QR Code Quality Check”就能反应出QRCheckquality_v1_quality_post这个“单词”。检索训练Retrieval Training目标让模型学会根据用户意图选出正确的工具token。输入是真实的用户查询如“帮我从Medium找一些关于Web开发的热门文章”输出是应该调用的工具token。这相当于教学生用刚背的单词来造句回答问题。端到端智能体调优End-to-End Agent Tuning目标让模型掌握连续使用多个工具来完成复杂任务的能力。输入是用户查询输出是一个完整的“思考-行动-观察”轨迹其中包含多个工具token的生成和中间的自然语言推理。这是最高阶的应用教学生写一篇使用多种工具解决问题的短文。这种由简到繁的训练方式符合认知规律能显著提升训练稳定性和最终效果。3. 环境准备与数据获取理论很美妙现在我们来动手搭建。复现ToolGen的第一步是准备好它的“食材”——模型、数据和运行环境。3.1 基础环境配置ToolGen基于PyTorch和Hugging Facetransformers库构建。建议使用Python 3.9版本。创建一个新的conda环境是一个好习惯。# 创建并激活环境 conda create -n toolgen python3.9 conda activate toolgen # 安装核心依赖 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 # 请根据你的CUDA版本调整 pip install transformers4.36.0 pip install datasets pip install accelerate pip install sentencepiece # 如果使用Llama等模型需要 pip install tiktoken # 如果使用Qwen等模型需要此外由于ToolGen的评估严重依赖ToolBench你还需要克隆并设置ToolBench环境以获取API调用密钥。这不是运行核心训练和推理的必须项但如果你想完整复现论文中的端到端评估则是必要的。git clone https://github.com/OpenBMB/ToolBench.git cd ToolBench # 按照ToolBench的README安装依赖并申请API Key # 你将得到一个toolbench_key保存好它。3.2 获取ToolGen模型与数据ToolGen的官方资源都托管在Hugging Face上获取非常方便。1. 下载预训练模型作者提供了基于Llama-3-8B和Qwen2.5-3B的两个版本。你可以直接使用transformers库加载。from transformers import AutoModelForCausalLM, AutoTokenizer model_name “reasonwang/ToolGen-Llama-3-8B” # 或 “reasonwang/ToolGen-Qwen2.5-3B” model AutoModelForCausalLM.from_pretrained(model_name, torch_dtypetorch.bfloat16) tokenizer AutoTokenizer.from_pretrained(model_name)2. 下载训练与评估数据集所有数据集都集中在reasonwang/ToolGen-Datasets仓库中。# 使用 huggingface-cli (需要先 pip install huggingface-hub) huggingface-cli download reasonwang/ToolGen-Datasets --repo-type dataset --local-dir ./toolgen_data # 或者直接下载压缩包 # wget https://huggingface.co/datasets/reasonwang/ToolGen-Datasets/resolve/main/data.tar.gz # tar -xzvf data.tar.gz下载后你会得到几个核心JSON文件toolgen_atomic_memorization.json: 工具记忆阶段数据。toolgen_atomic_retrieval_G123.json: 检索训练阶段数据。toolgen_atomic_G123_dfs.json: 端到端智能体调优数据。virtual_tokens.txt: 包含所有47k工具虚拟token的列表是“工具虚拟化”的基石。3.3 理解数据格式以工具记忆为例我们打开toolgen_atomic_memorization.json看一眼。它被处理成了类似ShareGPT的对话格式这极大地便利了与主流训练框架如FastChat、LLaMA-Factory的集成。[ { conversations: [ { role: user, content: Tool Name: QRCheck. Tool Description: Check the quality of any QRCode Api Name: quality_v1_quality_post Api Description: None., loss: false }, { role: assistant, content: QRCheckquality_v1_quality_post, loss: true } ] } ]loss: false表示在计算损失时这一部分内容通常是用户输入或系统指令会被忽略。模型不会因生成这些内容而受到惩罚。loss: true表示这一部分内容模型的输出将参与损失计算。在这里模型唯一需要学习的就是输出正确的虚拟token。这种格式非常清晰user部分提供了工具的元信息名称、描述、API名assistant部分给出了唯一正确的答案——工具虚拟token。实操心得在实际训练中确保你的训练脚本能正确识别和处理loss标签至关重要。许多训练框架默认对所有token计算损失你需要修改损失掩码loss mask的逻辑使其只对loss为true的token生效。否则模型会在学习工具token的同时被迫去“学习”那些作为条件输入的工具描述这会干扰主要目标。4. 核心实现从工具虚拟化到模型训练现在我们进入最核心的环节一步步拆解ToolGen是如何实现的。4.1 工具虚拟化将API“注入”词汇表这是整个项目的起点也是最精妙的一步。代码在项目提供的virtual_tokens.txt和示例脚本中已有体现我们来详细解读其意图和细节。import torch import transformers from unidecode import unidecode # 1. 加载虚拟token列表 with open(‘data/virtual_tokens.txt‘, ‘r‘) as f: virtual_tokens f.readlines() # 使用unidecode处理确保所有token都是ASCII字符避免分词器出现未知字符问题 virtual_tokens [unidecode(vt.strip()) for vt in virtual_tokens] # 2. 加载基础模型和分词器 model_name_or_path “meta-llama/Meta-Llama-3-8B-Instruct” # 以Llama 3为例 tokenizer transformers.AutoTokenizer.from_pretrained(model_name_or_path) model transformers.AutoModelForCausalLM.from_pretrained( model_name_or_path, torch_dtypetorch.bfloat16, device_map“auto” # 使用多GPU自动分配 ) # 3. 关键步骤扩展词汇表和嵌入层 original_vocab_size len(tokenizer) print(f“Original vocab size: {original_vocab_size}”) # 将虚拟token作为普通token非特殊token加入分词器 tokenizer.add_tokens(new_tokensvirtual_tokens, special_tokensFalse) new_vocab_size len(tokenizer) print(f“New vocab size: {new_vocab_size}, added {new_vocab_size - original_vocab_size} tokens”) # 相应地必须扩展模型词嵌入矩阵的大小 model.resize_token_embeddings(new_vocab_size) # 4. 初始化新token的嵌入向量 # 虚拟token的格式是 ToolNameApiName我们需要一个合理的初始化值 for vt in virtual_tokens: # 提取工具名和API名例如从 ‘QRCheckquality_v1_quality_post‘ 提取出 [‘QRCheck‘, ‘quality_v1_quality_post‘] combined_token vt[2:-2].split(““) # 去掉‘‘和‘‘再分割 # 将这些组成部分编码成token ID combined_token_ids tokenizer(“ “.join(combined_token), add_special_tokensFalse).input_ids # 获取虚拟token本身的ID应该只有一个 virtual_token_id tokenizer(vt, add_special_tokensFalse).input_ids assert len(virtual_token_id) 1, f“Virtual token {vt} should be tokenized into a single id, but got {virtual_token_id}“ # 计算组成部分token嵌入向量的平均值 with torch.no_grad(): # 重要在初始化阶段不计算梯度 combined_token_embeddings model.model.embed_tokens(torch.tensor(combined_token_ids).to(model.device)) # 对组成词的嵌入取平均作为新token的初始化值。这是一种常见且有效的策略。 mean_embedding torch.mean(combined_token_embeddings, dim0) # 将计算出的平均值赋值给新token的嵌入向量 model.model.embed_tokens.weight.data[virtual_token_id[0]] mean_embedding print(“Tool virtualization and embedding initialization completed.“)为什么这么做unidecode确保所有token是ASCII防止分词器Tokenizer遇到多语言字符时将其拆分成多个子词subword破坏“一个工具对应一个token”的设定。special_tokensFalse我们将工具token视为普通词汇而不是像[CLS]、[SEP]这样的特殊控制token。这样模型在生成文本时可以像生成普通单词一样自然地生成它们。平均初始化直接将一个随机向量赋给新token是低效的。用其组成部分工具名和API名的平均嵌入来初始化为这个新“单词”提供了一个语义丰富的起点能大幅加速后续训练收敛。注意事项resize_token_embeddings会改变模型参数保存模型时务必使用tokenizer.save_pretrained()和model.save_pretrained()一起保存否则下次加载时分词器和模型的词汇表会不匹配导致灾难性错误。4.2 三阶段训练实战完成虚拟化后我们就得到了一个“认识”了所有工具单词但还不会用的模型。接下来是三轮训练。阶段一工具记忆训练目标让模型建立工具描述 - 工具虚拟token的强关联。数据使用toolgen_atomic_memorization.json。关键配置损失函数仅计算assistant回复部分即虚拟token的交叉熵损失。学习率通常使用较小的学习率如5e-6到1e-5因为这是在微调不是从头训练。训练轮数论文中训练了8个epoch。由于任务简单近乎于配对任务模型会很快过拟合需要监控验证集loss在准确率接近100%后即可停止。训练脚本思路你可以使用LLaMA-Factory它原生支持这种带loss标签的对话格式。关键是在modeling_utils.py或相关代码中确保根据loss字段构建正确的labels和attention_mask。阶段二检索训练目标让模型学会用户查询 - 工具虚拟token的映射。数据使用toolgen_atomic_retrieval_G123.json。关键变化输入从“工具描述”变成了多样化的“用户自然语言查询”。这要求模型真正理解用户意图。训练技巧从第一阶段训练好的检查点checkpoint开始而不是原始模型。学习率可以比第一阶段稍高一点如1e-5因为任务变得更复杂。通常训练1个epoch即可。这是一个典型的检索精调任务目的是在已“背熟单词”的基础上学会“用词答意”。阶段三端到端智能体调优目标训练模型进行多轮工具调用和自然语言推理。数据使用toolgen_atomic_G123_dfs.json。数据格式更复杂包含多轮对话和|im_end|等特殊标记。训练挑战这是最复杂的阶段数据包含了完整的智能体轨迹思考、行动、观察。需要模型学会规划、调用工具、处理工具返回结果并生成最终答案。实操建议强烈建议使用作者提供的训练脚本training/目录下因为它们已经处理好了这种复杂轨迹数据的格式和损失计算逻辑。如果使用其他框架你需要仔细设计数据加载器确保能正确解析多轮对话并只在模型需要生成的部分如assistant的回复、工具token、思考内容计算损失。4.3 使用ToolGen进行推理训练完成后如何使用这个“武装到牙齿”的模型呢项目提供了一个简洁的OpenAgent接口。import json from OpenAgent import ToolGen, RapidAPIWrapper # 1. 初始化工具执行器需要提前部署StableToolBench并获取key with open(“keys.json“, ‘r‘) as f: keys json.load(f) toolbench_key keys[‘TOOLBENCH_KEY‘] rapidapi_wrapper RapidAPIWrapper( toolbench_keytoolbench_key, rapidapi_key““, # 如果需要特定RapidAPI key ) # 2. 初始化ToolGen智能体 # 选择模型和对应的模板 agent ToolGen( model_name_or_path“reasonwang/ToolGen-Llama-3-8B“, template“llama-3“, # 必须与模型匹配Qwen模型用“qwen-7b-chat“ indexing“Atomic“, # 使用Atomic索引模式与训练一致 toolsrapidapi_wrapper, # 传入工具执行器 ) # 3. 构建对话并运行 messages [ {“role“: “user“, “content“: “帮我查一下北京明天和上海的天气对比。“} ] agent.restart() # 清空对话历史 try: agent.start( single_chain_max_step16, # 最大推理步数防止死循环 start_messagesmessages ) # 运行后agent会自行调用工具、获取结果、并生成最终回答 final_response agent.get_response() print(“最终回答“, final_response) except Exception as e: print(f“推理过程中出现错误{e}“)在这个流程中ToolGen类内部会处理工具token检测监控模型生成的文本一旦出现...格式的虚拟token就将其截获。工具解析与调用将虚拟token解析为对应的工具名和API名然后通过RapidAPIWrapper调用真实的API。结果整合将API返回的结果格式化后作为“观察”重新插入到对话上下文中让模型继续生成后续内容可能是下一个工具调用也可能是最终答案。5. 评估与结果分析它到底有多强论文中展示了ToolGen在多个基准测试上的优异表现。作为实践者我们更关心如何复现这些评估以及结果背后的含义。5.1 工具检索评估这是评估模型“找工具”能力的核心。使用项目中的evaluation/retrieval/eval_toolgen.py脚本。python -m evaluation.retrieval.eval_toolgen \ --model_name_or_path “./checkpoints/toolgen_retriever“ \ # 你第二阶段训练出的检索模型 --indexing “Atomic“ \ --stage “G1“ \ # 评估集分组G1/G2/G3代表不同难度或领域 --split “instruction“ \ # 使用指令分割的数据集 --result_path “./eval_results/retrieval/“ \ --constrain True # 是否将检索范围限制在Ground Truth涉及的工具集内关键指标解读Hitk (k1, 3, 5)模型生成的前k个token中包含正确工具token的比例。这是最直接的指标。ToolGen在Hit1上提升显著说明其“直接命中”能力极强。MRR (Mean Reciprocal Rank)正确工具token在生成序列中排名的倒数的平均值。衡量排名质量。为什么强因为传统检索器如基于BERT的编码器是将查询和工具描述映射到同一个向量空间进行相似度计算。而ToolGen是将工具知识内化到了语言模型的生成概率中它“回忆”工具更像是在“回忆知识”对于训练中见过的工具模式其准确性有天然优势。5.2 端到端任务解决评估这是评估模型“用工具解决问题”的最终能力。流程更复杂推理生成轨迹使用训练好的完整ToolGen模型在测试查询集上运行生成包含工具调用和自然语言的完整轨迹文件.json。轨迹格式转换使用scripts/convert_answer/run_convert_answer.sh将生成的轨迹转换为ToolBench标准评估格式。可解通过率评估运行scripts/pass_rate/run_pass_rate.sh。这个脚本会调用ToolBench的评估器模拟执行轨迹中的工具调用判断任务是否被成功解决。核心指标是Solvable Pass Rate。胜率评估运行scripts/preference/run_preference.sh通常是将ToolGen的解决方案与基线模型如GPT-4 with ReAct的解决方案交由GPT-4作为裁判进行偏好比较得出Win Rate。避坑指南端到端评估是整个复现过程中最耗时的环节且严重依赖外部APIToolBench的RapidAPI的稳定性和额度。建议先在小规模测试集如10-20条数据上跑通全流程。仔细管理你的API调用密钥和额度避免因超额导致评估中断。保存好中间生成的轨迹文件便于出错时调试。5.3 结果分析与启示根据论文ToolGen在工具检索上大幅领先传统检索器在端到端任务解决上也达到了SOTA或接近SOTA的水平。这给我们带来几点启示生成式检索的潜力对于高度结构化、离散化的信息如工具、实体、代码将其编码为离散token并让生成模型直接输出可能比连续的向量检索更精准、更高效。课程学习的有效性三阶段训练是一个被验证有效的范式。对于其他需要让LLM掌握复杂技能的任务如编程、数据库查询也可以借鉴这种“记忆-应用-综合”的渐进式训练策略。系统设计的简洁性ToolGen将复杂的工具使用流程简化为了一个“生成特殊token”的动作。这大大降低了智能体系统的工程复杂度使得工具调用可以像文本生成一样流式进行。6. 常见问题与排查技巧实录在实际复现和实验过程中我踩过不少坑这里总结一下最常见的问题和解决方法。6.1 训练相关问题Q1: 在工具虚拟化后模型输出乱码或无法生成虚拟token。可能原因A分词器Tokenizer和模型的词汇表未同步保存和加载。检查加载模型时是否使用了与训练时完全相同的分词器文件务必使用model.save_pretrained(‘path‘)和tokenizer.save_pretrained(‘path‘)保存到同一目录并从这个目录加载。可能原因B虚拟token在输入文本中被错误地拆分。检查使用tokenizer(‘MyToolmy_api‘, add_special_tokensFalse).input_ids查看输出是否是一个单独的ID列表。如果不是说明你的分词器将这个token拆分了回顾虚拟化步骤是否使用了unidecode确保其为ASCII。可能原因C训练不充分或数据格式有误。检查在记忆训练阶段模型在验证集上对工具token的预测准确率是否接近100%如果没有增加训练轮数或检查数据中loss标签是否正确。Q2: 损失Loss不下降或下降非常慢。可能原因A学习率设置不当。尝试对于LoRA等参数高效微调方法学习率通常设在1e-4量级对于全参数微调设在5e-6到2e-5之间。使用学习率预热warmup和余弦衰减cosine decay调度器。可能原因B损失掩码Loss Mask未正确应用。调试打印一个batch的labels和attention_mask。确保只有assistant回复部分对应loss: true的token在labels中不是-100在PyTorch中-100的标签会被损失函数忽略。可能原因C批次大小Batch Size太小。尝试在GPU内存允许的情况下增大批次大小。如果使用梯度累积Gradient Accumulation确保有效批次大小足够大如32以上。6.2 推理与评估相关问题Q3: 模型在推理时陷入循环反复生成同一个工具token或无意义内容。可能原因Asingle_chain_max_step设置过小或过大。调整这个参数控制单次任务的最大推理步数。太小时任务可能未完成就被强制结束太大时模型可能在一个错误步骤上无限循环。根据任务复杂度调整一般8-20步是合理的起点。在agent.start()中设置。可能原因B温度Temperature和Top-p采样参数过于激进。调整在生成工具token这类需要高确定性的任务时应使用低温度如0.1和低Top-p如0.9甚至直接使用贪婪解码do_sampleFalse。你可以在初始化ToolGen时通过generation_config参数传递这些设置。可能原因C工具执行失败或返回异常结果导致模型上下文混乱。检查在RapidAPIWrapper中增加日志打印每次工具调用的请求和响应。确保API服务正常返回格式符合预期。ToolGen框架应能处理常见的API错误并将其作为“观察”反馈给模型。Q4: 端到端评估的通过率Pass Rate远低于论文报告值。可能原因A评估环境不一致。检查你使用的ToolBench API Key对应的后端服务StableToolBench是否与论文实验时一致API的版本、可用性可能影响结果。确保你运行的评估脚本版本和数据集划分与官方完全一致。可能原因B模型过拟合或欠拟合。检查在保留的验证集上检查模型的工具检索准确率和单步任务完成率。如果验证集性能尚可但测试集很差可能是过拟合需要更多样化的训练数据或更强的正则化。如果验证集性能就很差则是欠拟合需要更多训练或调整模型容量。可能原因C轨迹转换或评估脚本参数错误。仔细核对run_convert_answer.sh和run_pass_rate.sh脚本中的路径、模型名称、数据集路径等参数是否正确。一个常见的错误是评估了错误的数据分割split。6.3 工程化与扩展思考Q5: 如何为我自己的工具集定制ToolGen创建虚拟token列表为你每个工具的唯一标识如MySystem::getUserInfo创建一个虚拟token字符串。确保格式唯一且易于解析。生成训练数据记忆数据(工具描述/文档) - 虚拟token。检索数据(用户可能的问题/指令) - 虚拟token。这部分数据需要精心构造或通过少量样本大模型生成。轨迹数据编写或模拟用户使用你的工具完成任务的完整对话记录。重复训练流程从你的基础模型如Llama 3开始执行虚拟化、记忆训练、检索训练、轨迹训练。Q6: ToolGen方法有什么局限性词汇表膨胀每增加一个工具就永久占用一个词汇表位置。对于超大规模工具库百万级这会显著增加模型参数嵌入层并可能影响原有语言的生成能力。需要权衡。工具更新成本高新增或修改工具需要重新进行虚拟化和至少是记忆训练不如基于描述的检索方法灵活。对闭源模型不友好虚拟化需要修改模型的词汇表和嵌入层这对于仅提供API的闭源大模型如GPT-4是无法实现的。多模态工具当前方法主要针对API工具。对于需要图像输入或输出的工具如何将其“虚拟化”是一个开放问题。尽管有这些局限ToolGen无疑为工具学习领域提供了一个强大、简洁且可解释性高的新范式。它让我们看到让大模型“掌握”工具或许可以像教它一门新的方言一样自然。