Jupyter Agents:训练 LLM 在笔记本中进行推理
Hugging Face 开发了一条流水线,用于训练小型语言模型,使其能够作为数据科学代理在 Jupyter 笔记本 中执行代码。通过结合从 Kaggle 笔记本中提取的精心策划的合成数据集和简化的代理脚手架,团队成功提升了 Qwen3-4B 模型在数据分析基准上的表现。
数据科学代理的 DABStep 基准
为了衡量代理在数据科学能力方面的进展,Hugging Face 使用 DABStep 基准,该基准通过提供数据集并要求模型回答非平凡的数据问题来评估代理在真实任务中的表现。该基准对当前 LLM 来说尤为困难;例如,Claude 4 Sonnet 在基准的困难任务上准确率不足 20%。
优化代理脚手架
代理的性能在很大程度上取决于脚手架——即引导模型行为的代码和提示结构。Hugging Face 对 Jupyter 代理的脚手架进行了重构,使其轻量化并贴合模型需求:
- 简化架构: 脚手架被精简至约 200 行代码,且没有外部依赖。
- 工具集成: 代理在
while循环中运行,使用两个主要工具:code execution和final_answer。 - 性能提升: 仅此优化就将 DABStep 简单子集的准确率从 44.4% 提升至 59.7%。
Jupyter 代理数据集流水线
为提升小模型的能力,Hugging Face 使用 Datatrove 构建了一个多阶段流水线,将约 2TB 的原始 Kaggle 笔记本转化为高质量的训练集,包含 51k 条合成笔记本(约 0.2B token)。
1. 去重与数据获取
原始 Kaggle 笔记本经过去重处理后,体积从 2TB 缩减至 250GB。随后团队使用 kagglehub 包下载约 5TB 的关联数据集,并过滤掉模型检查点、多模态语料库以及超过 10GB 的数据集,以确保它们能够适配 E2B 沙箱。
2. 教育评分与过滤
利用 Qwen3-32B,团队实现了一个“教育评分”系统,对笔记本的清晰度和教学价值进行 1–5 级评分,剔除约 70% 的笔记本。另有 20% 通过基于 LLM 的过滤被移除,以排除与数据分析无关或未使用数据集的笔记本。
3. 合成问答与轨迹生成
为了生成可验证的训练数据,流水线基于真实笔记本轨迹生成合成的问答(QA)对。
- QA 生成: Qwen3-32B 生成自然语言的问题和答案,随后由第二个 LLM 验证以防止幻觉。
- 轨迹生成: 使用 Qwen-3-Coder-480B-A35B-Instruct 生成干净的、逐步的代码执行轨迹来回答合成问题。
- 处理缺失数据: 当数据集不可用时,模型被提示充当有状态的 Python 代码解释器以模拟执行。
- 推理仿真: 为确保模型在代码单元之间提供注释,
comment字段被设为代码执行工具调用的必需部分。
训练与结果
在 Qwen3-4B 上使用 TRL 库进行微调。团队发现全参数微调优于 PEFT(参数高效微调),并且使用 assistant_loss_only=True 能带来性能提升。
性能提升
在合成数据集上训练显著提升了模型在 DABStep 简单子集上的表现:
| 模型 | 训练轮数 | DABStep(简易) |
|---|---|---|
| Qwen-3-4B-Instruct-2507 (Base) | 0 | 38.67% |
| Qwen-3-4B-Instruct-2507 (Our Scaffolding) | 0 | 52.78% |
| Qwen-3-4B-Instruct-2507 | 2 | 63.89% |
| Qwen-3-4B-Instruct-2507 | 3 | 73.61% |
| Qwen-3-4B-Instruct-2507 | 5 | 75% |
| Qwen-3-4B-Instruct-2507 | 7 | 70.83% |
最终微调的 Qwen-4B 模型在简易得分上相较基线提升最高达 36%,相较脚手架基线提升 22%,同时在困难得分上也表现出改进。
开源发布
Hugging Face 向社区发布了以下资源:
- Jupyter Agent 数据集: 包含 51k 条合成笔记本的集合。
- 已调优模型:
jupyter-agent-qwen3-4b-instruct和jupyter-agent-qwen3-4b-thinking。
未来方向
团队确定了多条进一步改进的路径,包括生成更具挑战性的多步骤问题、扩大精选轨迹的规模、研究知识蒸馏,以及构建利用现有可验证 QA 设置的强化学习(RL)环境。