Jupyter Agent：训练小模型做数据分析

过去一年，大家都在给大语言模型（LLM）加工具、提自主性，好让它们能处理更复杂、更开放的任务。Jupyter Agent 的目标，就是给模型一个终极工具：代码执行。

展示多步代码执行和推理过程，最自然的方式就是 Jupyter Notebook——它由代码单元格和 Markdown 单元格组成。所以我们构建了 Jupyter Agent，让它能直接在 Jupyter Notebook 里执行代码，并用这个环境来解决数据分析和数据科学任务。你可以把它想象成 Cursor，但它原生就活在你的数据科学工作流里。

我们用目前最强的编程模型之一 Qwen-3 Coder 做了一个演示，这是之前 jupyter-agent（v1）工作的延续。

虽然大模型开始展现出有用的行为，但关键问题是我们如何继续改进它们。为此，我们专注于增强小模型在智能体（Agent）数据科学任务上的表现，因为它们目前还难以与大模型竞争。

这个项目的目标是构建一个管道：首先生成高质量的训练数据，然后微调一个现有小模型，最后评估模型在相关基准测试上的性能是否提升。

让我们从最后一步开始：为数据科学任务选择一个强大的基准测试。

🏁 入门：DABStep 基准测试

要判断我们是否在构建更好的数据科学智能体方面取得进展，我们需要一个基准来衡量这些能力。去年，我们与 Adyen 合作推出了 DABStep 基准测试：一种评估数据科学智能体在真实任务上表现的方法。设置很简单：给 LLM 提供数据集，让它回答非平凡的数据问题。

示例任务：

这个基准对今天的 LLM 来说仍然具有挑战性——例如，最好的开箱即用模型 Claude 4 Sonnet 在困难任务上的准确率甚至达不到 20%。你可以在这里探索实时排行榜。

🎯 第一个基线

既然我们找到了一个好的基准，就可以尝试攀登它了！我们着手构建一个用于微调的数据集，让即使是小型数据智能体模型也能在 DABStep 上表现良好。

我们的首选是 Qwen3-4B-Thinking-2507：体积极小（迭代快，易于运行），但又足够强大，能在智能体场景中行动。

基线结果：

• 简单任务：44.4%
• 困难任务：2.1%

不算好——但这是一个有希望的起点，因为它留下了很大的改进空间。让我们看看如何改进它！

🔧 执行框架（Harness）入门

智能体（Agent）与纯聊天模型的一个核心区别，是围绕模型构建的执行框架（Harness），用于引导其行为。例如，DABStep 中的评估脚本使用 smolagents 来执行代码。Smolagents 带有预定义的行为、提示结构和预期格式。

我们还研究了 Qwen-Agent 代码库，作者在那里为模型量身定制了执行框架（Harness）。这很有道理：例如，Claude Code 与 Claude Sonnet 配合得惊人地好，因为它们的执行框架（Harness）是对齐的。

所以，我们重构了我们的执行框架（Harness）：

• 将其精简到约 200 行代码。
• 没有外部依赖。
• 灵感来源于 tiny-agents 的精神。

👉 在这里查看：utils.py。

结果： 准确率从 44.4% 跃升至 59.7%（简单任务分割）。🚀

我们的循环：

• 带有两个工具的 While 循环：code_execution 用于运行代码，final_answer 用于返回最终答案。
• 我们与 Qwen-Agent 的不同之处在于，明确添加了一个 final_answer 工具——在我们的测试中，这提高了性能。
• 与 smolagents 相比，我们通过移除大量提示和工具简化了执行框架（Harness）。Smolagents 还通过使用 ReACT 框架将许多假设硬编码到模型中。

🏃‍♂️ 训练管道

有了简化的执行框架（Harness），我们专注于微调 Qwen3-4B，用于数据科学智能体（Agent）任务。

⚙️ 数据集管道

改进模型在特定任务或行为上的表现，秘诀是训练它使用尽可能接近任务的数据。一个自然的起点是查看真实的 Jupyter Notebook，并找到与我们计划处理的任务（即数据分析）密切相关的笔记本。

Kaggle 笔记本提供了大量高质量的数据分析笔记本，并由 Kaggle 提供：

数据集：

• Kaggle Notebooks 数据集： 约 2TB 的笔记本。
• Kaggle Datasets： 5TB 的 Kaggle 数据集，我们手动下载并将其链接到笔记本。
• 每个笔记本的丰富元数据（作者、使用的数据集等）。

既然我们在基础模型上取得了不错的结果，是时候构建一个数据集来帮助我们进一步改进它了。我们设计了一个多阶段管道，使用 Datatrove 大规模清理和准备 Kaggle 笔记本。

以下是每个步骤的工作原理：

1. 大规模去重

我们从约 2TB 的 Kaggle 笔记本开始，通过重用我们在 BigCode 项目中的工作，将其减少到约 250GB。作为 StarCoder2 训练数据处理的一部分，笔记本（不含输出单元格）已经去重。大多数 Kaggle 笔记本是微小的变体或近乎相同的副本，因此这一步至关重要。

关键洞察： 约 90% 的原始笔记本是重复的，如果不过滤，会扭曲训练。

2. 下载链接的数据集

大多数 Kaggle 笔记本通过 Kaggle 元数据引用外部数据集。为了确保笔记本内的代码能够实际运行，我们构建了一个管道，自动获取这些链接的数据集。这一步至关重要，因为许多笔记本否则将不完整或无法执行。

使用 kagglehub 包，我们下载了数千个数据集——总计约 5TB。为了保持可管理性和相关性：

• 我们过滤掉了包含模型检查点、大型多模态语料库或 LLM 相关文件的数据集。
• 我们还排除了无法放入我们用于执行的虚拟 E2B 沙箱的非常大的数据集（10GB+）。

最终，我们拥有一个丰富的可执行笔记本集合，并配有其数据集，为在真实、可运行的环境中训练智能体（Agent）提供了基础。

3. 教育评分

我们使用 Qwen3-32B 根据教育质量对笔记本进行评分。我们发现使用整个笔记本并不理想，因为许多包含琐碎或损坏的代码。我们的教育评分方法详见 edu_scoring.py。

TL;DR： 我们根据清晰度、完整性和教育价值，为每个笔记本分配 1-5 的分数，并只保留高于选定阈值的笔记本。这种过滤移除了大约 70% 的笔记本。

这与 BeyondWeb 论文的见解类似，该论文表明使用高质量数据对于合成数据生成更好——这是我们依赖 QA（问答）生成的一步。这有助于模型从“高质量”笔记本中学习，而不是从嘈杂的笔记本中学习。

4. 过滤不相关的笔记本

我们排除了关于训练 LLM 或与数据分析无关的笔记本。我们还通过使用 Qwen3-32B 的自动化基于 LLM 的过滤过程，移除了实际上未使用数据集的笔记本。过滤的实现可以在 extract_packages_and_files.py 中找到。

TL;DR： 我们提示 Qwen3-32B 识别并移除那些（1）与数据分析无关，或（2）实际上未使用数据集的笔记本。这一步移除了大约 20% 的笔记本。

这确保了我们只在相关的数据科学任务上进行训练。

5. QA 生成

使用清理后的笔记本，我们使用 Qwen3-32B 生成问答对。问题和答案基于真实的笔记本执行轨迹，因此 QA 对基于真实的代码执行结果。

提示设计： 我们要求 LLM 生成可能对数据集提出的自然问题，然后验证笔记本是否提供了正确答案。

挑战： 我们必须尝试许多提示来获得更高难度的问题，因为 LLM 倾向于生成像“数据集的大小是多少”这样的琐碎问题。

洞察： 我们将其分为两步，因为 LLM 倾向于产生模型幻觉（Hallucination）：

1. 生成问题和答案。
2. 要求另一个 LLM（可以访问笔记本）检查答案是否正确。

完整的提示策略和实现可在 generate_qa.py 中找到。

最后一步是生成干净的代码执行轨迹（Trace）。原始 notebook 经过处理后，往往内容冗长、结构开放，包含大量无关部分。但我们希望 Jupyter Agent 能高效地得出结果。为此，我们基于原始 notebook 合成了更简洁的轨迹用于训练。

我们提示 Qwen-3-Coder-480B 模型，让它生成一个 Jupyter notebook 代码来回答之前合成的 QA 对中的问题。这些轨迹捕获了逐步的代码执行过程，包括中间输出，这对智能体训练至关重要。

我们使用 E2B 来让智能体解决这些合成 QA 对，这需要获取 Kaggle 数据集，以便代码能在 E2B 中实际运行。

挑战 1： 许多数据集无法获取。 技巧： 由于 LLM 擅长代码且具备不错的世界模型，当数据集缺失时，我们提示它们扮演代码解释器的角色。

提示的开头部分：

You are a stateful Python code interpreter that executes code in a persistent environment. Your role is to execute Python code while maintaining state across multiple code cells, similar to a Jupyter notebook environment.
[REST OF THE PROMPT]

挑战 2： Qwen3-Coder-480B-A35B 模型不支持思考模式——我们如何提取代码注释？默认情况下，它通常只输出简短注释，然后就是几行代码执行。但我们希望每个单元格之间都有一些推理或注释。 技巧： 当我们从 Qwen3-32B 切换到 Qwen3-Coder-480B-A35B 时，发现输出消息内容经常为空。这其实是 Qwen3-Coder 模型的一个已知特性：在使用工具调用时，模型不会返回空的助手响应。我们通过工具强制要求一些文本注释，在代码执行工具调用中传递 'comment' 作为必填字段。这样，当使用非推理模型生成代码单元格时，它会默认从第一人称视角输出一些动作描述，模拟思考轨迹的结构。

注意： notebook 中生成的最终答案可能与 QA 对中指定的答案不同。这是因为智能体模型可能使用与原始 Kaggle notebook 不同的数据预处理方法和步骤，而合成问题通常不会指定这些细节。这种差异是正常的，也为一个新的研究方向奠定了基础：语言模型如何处理数据分析，以及它们是否与人类做法不同。为了完全透明，我们同时保留了 LLM 生成的最终答案和真实 Kaggle notebook 的原始答案，作为模型性能的信号。我们鼓励社区尝试不同的数据集组合，看看如何能进一步提升性能。

7. 最终筛选

我们截断了过长的输出，并过滤掉琐碎的轨迹，以防止内容长度问题，只保留高质量的轨迹。 我们保留了与 DABStep 风格任务一致的非琐碎、多轮交互轨迹。 最终的 Jupyter Agent Dataset 包含了 5.1 万个合成 notebook 和近 2 亿个 token，成为在 Qwen3-4B 模型上进行监督微调（SFT）的基础。

有了这个数据集，下一步自然是要看看它是否真的能帮助我们的模型成为更强的数据科学智能体。接下来，我们进入训练流程并评估其影响。

🏃‍♂️ 训练流程

数据集准备就绪后，我们转向核心问题：这些数据真的能帮助模型更好地解决数据分析任务吗？ 为了找到答案，我们建立了一个简单的微调流程，并通过实验来测量训练对合成 notebook 的影响。

一些训练步骤特别有趣，并给了我们有用的见解：

• 对于轨迹生成，我们使用 LLM 生成 QA 对，这提供了一个可验证的环境。
• 最后，我们使用 TRL 对 Qwen3-4B 进行了微调。
  • 设置 assistant_loss_only=True → 带来了小幅性能提升。
  • 在全参数多轮训练中添加了 neftune 噪声 → 避免过拟合。

挑战：

• 提示模型进行工具调用很棘手：并非所有提示都能带来相同的性能（参考 Qwen 文档）。
• 我们必须手动测试每一个提示，以找到效果最好的。
• 工具调用的响应格式没有标准化，这使得在不同模型间切换变得困难。
• 原生的 Qwen 生成提示不适应 TRL 中的 assistant_loss_only=True 训练模式，该模式默认需要生成 token。因此，我们通过将助手响应部分包裹在生成标签中来调整原始聊天模板。
• 在短推理文本上训练思考模型可能会破坏模型能力 → 在这种情况下，全参数训练比 PEFT 效果更好。

我们完整的训练实现，包括超参数配置和模板调整，都可以在仓库的 finetuning 目录 中找到。

📊 结果

首先，我们使用 Qwen3-Coder-480B-A35B 生成了最终数据集，其中包含高质量的代码和简短的类推理轨迹。之后，我们开始了训练，并尝试了各种配置，如 PEFT/适配器与全参数调优、学习率、轮数、添加噪声等。我们发现，全参数微调能让模型更好地学习和复制 Qwen3-Coder-480B-A35B 的行为响应质量，其支持性注释更简短，更贴合数据分析任务，避免了不必要的冗长推理。

我们进行了一个小型的消融研究，探讨训练轮数的影响：

我们观察到，对于 SFT，使用较低的学习率和较高的 neftune 噪声（7），进行比通常稍多的训练轮数是有益的。最后，我们将训练好的模型与已实现的脚手架进行比较，以确定训练数据集的纯粹影响。总结来说，与基础模型/带脚手架模型相比，我们在 DABStep Easy 分数上可以看到高达 36%/22% 的提升：

我们还可以看到，Hard 分数也能提升，尽管我们的数据集主要关注较简单的问题：

从上面的图表中可以看出，新的脚手架和基于合成 notebook 的微调都产生了显著影响。这使得 Qwen-4B（搭配我们的流程和脚手架）成为 DABStep 上最先进的小模型智能体。

在实践中，该模型现在能够一致地执行并解决广泛的现实 Kaggle 风格数据分析任务。 它对于最难的查询还不够强，但我们已经证明，即使是小模型，在搭配正确的数据和脚手架时，也能成为强大的智能体。

亲自试试 Jupyter Agent

这些结果表明，即使小模型，采用正确的训练方法也能成为强大的数据科学智能体。准备好亲自尝试了吗？我们已经将所有内容开源，你可以用我们微调好的模型和数据集进行实验。

我们开源了性能最佳的微调 Qwen3-4B-Instruct-2507 和 Qwen3-4B-Thinking-2507 检查点，以及训练数据集，你可以试用和实验：

• Jupyter Agent Dataset
• Jupyter-Agent-Qwen3-4B-Instruct
• Jupyter-Agent-Qwen3-4B-Thinking

你可以用以下几行代码加载 Jupyter Agent Dataset：

from datasets import load_dataset
# To load the train split of a specific subset, such as non-thinking, you can do
ds = load_dataset("jupyter-agent/jupyter-agent-dataset", split="non-thinking")
# apply chat template
tokenizer.apply_chat_template(ds[0]["text"])

你也可以使用以下代码，通过 E2B 代码执行直接使用来源 Kaggle 数据集：

import kagglehub
import e2b_code_interpreter as e2b
from datasets import load_dataset

# load the Jupyter Agent Dataset
ds = load_dataset("jupyter-agent/jupyter-agent-dataset", split="thinking")
# get the kaggle dataset name
dataset_name = ds[0]["kaggle_dataset_name"]
# load the dataset locally from Kaggle Hub
path = kagglehub.dataset_download(dataset_name)
print(path) # this is the folder path where the dataset is downloaded
# initialize sandbox
sandbox_init = e2b.Sandbox(timeout=240)
# write used file to E2B sandbox
file_name = ds[0]["files_used"][0]
file_name = file_name.split('/')[-1] if '/' in file_name else file_name
with open(f"{path}/{file_name}", "rb") as file:
    sandbox_init.files.write(f"/home/user/input/{file_name}", file)
# execute code with E2B
execution = sandbox_init.run_code("<some code>")

你可以按照 Qwen 文档代码使用微调好的 Jupyter Agent Qwen 模型：

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "jupyter-agent/jupyter-agent-qwen3-4b-instruct"

# load the tokenizer and the model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto"
)

# prepare the model input
prompt = "Give me a short introduction to large language model."
messages = [
    {"role": "user", "content": prompt}
]
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
)
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

# conduct text completion
generated_ids = model.generate(
    **model_inputs,
    max_new_tokens=16384
)
output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist() 

content = tokenizer.decode(output_ids, skip_special_tokens=True)

print("content:", content)

对于 Thinking 模型，你可以用以下代码解码思考响应和内容：

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "jupyter-agent/jupyter-agent-qwen3-4b-thinking"

# ...use same processing code from above...
try:
    # index finding 151668 (</think>)
    index = len(output_ids) - output_ids[::-1].index(151668)
except ValueError:
    index = 0

thinking_content = tokenizer.decode(output_ids[:index], skip_special_tokens=True).strip("\n")
content = tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n")

🔮 下一步

• 更困难的任务： 生成更具挑战性、多步骤的问题，更好地反映现实世界分析。
• 扩大规模： 在更大规模的精选轨迹上训练，以突破当前在 Hard 分集上 3.4% 的性能。
• 知识蒸馏： 研究知识蒸馏，这在改进小模型方面已显示出强大效果。
• 强化学习（RL）： 构建 RL 环境，这在智能体任务上已显示出实现最先进性能的潜力。由于我们的 QA 设置已经提供了一个可验证的环境，我们可以直接将其用于 RL 训练。

也许这会带来…… Jupyter-Agent 3。 😉

我们希望我们的发现能激励他人在开发更强大的 notebook 编码智能体方面继续取得进展，我们很期待看到社区接下来会构建什么。深入探索 🤗 Hub 上的 jupyter-agent dataset，并在 https://github.com/huggingface/jupyter-agent 探索代码库，开始你在 Jupyter notebook 智能体上的实验吧。