Jupyter Agents: training LLMs to reason with notebooks

Jupyter Agentis to give the model the ultimate tool: code execution.

A natural way to display multi-step code execution together with reasoning is within a Jupyter Notebook, which consists of code and markdown cells. So we built Jupyter Agent to act as an agent that can execute code directly inside a Jupyter notebook and use this environment to solve data analysis and data science tasks. Think of it like Cursor, but living natively inside your data science workflow.

We built a demo of this vision with Qwen-3 Coder, currently one of the strongest coding models. This is a follow-up to our earlier work on jupyter-agent (v1).

While large models are starting to show useful behavior, the key question is how we can continue improving them. To this end, we focus on strengthening smaller models to perform well on agentic data science tasks as they currently struggle to compete with the large models.

The goal of this project is to build a pipeline to first generate high-quality training data, then fine-tune an existing small model, and finally evaluate whether the model's performance improves on relevant benchmarks.

Let’s begin with the last step: selecting a strong benchmark for evaluating models on data science tasks.

In order to understand if we are making progress towards better data science agents we need a benchmark to measure such capabilities. Last year, in partnership with Adyen, we introduced the DABStep benchmark: a way to evaluate data science agents on realistic tasks. The setup is simple: provide the LLM with datasets and ask it to answer non-trivial data questions.

Example tasks:

| Question | Answer | |---|---| | Which card scheme had the highest average fraud rate in 2023? | SwiftCharge | For the year 2023, focusing on the merchant Crossfit Hanna, if we incentivize users to switch to a different Authorization Characteristics Indicator, which option would be the most cost-effective? | E:346.49 |

This benchmark remains challenging for today’s LLMs — e.g. the best out-of-the-box model is Claude 4 Sonnet which reaches not even 20% accuracy on the hard tasks.

You can explore the live leaderboard here.

Now that we identified a good benchmark we can try to climb it! We set out to build a dataset for fine-tuning such that even a small data agent model could perform well on DABStep.

Our first choice was Qwen3-4B-Thinking-2507: extremely small (fast to iterate with, easy to run), yet strong enough to act in agentic scenarios.

Baseline results:

*Easy tasks:***44.4%***Hard tasks:*2.1%

Not great — but a promising starting point, since it left a lot of room for improvement. Let's see how we can improve it!

A core aspect of agents that sets it apart from a pure chat model is the scaffolding built around the model to steer its behaviour. The evaluation script in DABStep for example uses smolagents to execute code. Smolagents comes with predefined behaviors, prompting structures, and expected formats.

We also studied the Qwen-Agent codebase, where the authors tailoring scaffolding to the model. This makes sense: Claude Code, for example, works shockingly well with Claude Sonnet because their scaffolding is aligned.

So, we restructured our scaffolding:

• Stripped it down to ~200 lines of code.
• No external dependencies.
• Inspired by the spirit of tiny-agents.

👉 Check it out here: utils.py.

Results: accuracy jumped from 44.4% → 59.7% (easy split). 🚀

Our loop:

• While loop with two tools: code executionto run the code andfinal_answerto return the final answer. - We differ from Qwen-Agent by explicitly adding a final_answertool — which in our testing has improved performance. - Compared to smolagents, we simplified the scaffolding by removing a lot of prompts and tools. Smolagents also hardcodes a lot of assumptions into the model by using the ReACT framework.

With simplified scaffolding in place, we focused on fine-tuning Qwen3-4B for data science agentic tasks.

The recipe to improve a model on a certain task or behaviour is to train it on data that reflects the tasks as closely as possible. A natural starting point is to look at real Jupyter Notebooks and find notebooks that align closely with the task that we plan to tackle, namely data analysis.

Kaggle notebooks offer a wealth of high quality data analysis notebooks and are made available by Kaggle:

Datasets:

Kaggle Notebooks dataset: ~2TB of notebooks.Kaggle Datasets: 5TB of kaggle datasets that we manually downloaded and linked to the notebooks.- Rich metadata for each notebook (authors, datasets used, etc.).

Now that we have good results with a base model it's time to build a dataset that will help us improve it even further. We designed a multi-stage pipeline using Datatrove to clean and prepare Kaggle notebooks at scale.

Here’s how each step worked:

We started with ~2TB of Kaggle notebooks and reduced it to ~250GB reusing our work from the BigCode project. As part of the StarCoder2 training data processing the notebooks (without output cells) were already deduplicated. Most Kaggle notebooks are small variations or near-identical copies, so this step was essential.Key insight: ~90% of raw notebooks are duplicates, which would have skewed training if left unfiltered.

Most Kaggle notebooks reference external datasets via Kaggle metadata. To make sure the code inside notebooks could actually run, we built a pipeline that automatically fetched these linked datasets. This step was crucial, since many notebooks would otherwise be incomplete or non-executable.

Using the kagglehub package, we downloaded thousands of datasets — about 5TB in total. To keep things manageable and relevant:

• We filtered out datasets containing model checkpoints, large multimodal corpora, or LLM-related files.
• We also excluded very large datasets (10GB+) that couldn’t fit into the virtual E2B sandboxes we used for execution.

By the end, we had a rich collection of executable notebooks paired with their datasets, providing the foundation for training agents in realistic, runnable environments.

We scored notebooks based on educational quality using Qwen3-32B. We saw that using the whole notebook was not optimal, as many contained trivial or broken code. Our educational scoring approach is detailed in edu_scoring.py.

TL;DR: We assigned each notebook a score from 1–5 based on clarity, completeness, and educational value, and kept only those above a chosen threshold. This filtering removed about 70% of the notebooks.

This is similar to the insight from the BeyondWeb paper, which showed that using high-quality data is better for synthetic data generation — a step we relied on for QA (Question-Answer) generation. This helped the model learn from “high quality” notebooks instead of noisy ones.

We excluded notebooks about training LLMs or unrelated to data analysis. We also removed notebooks that didn’t actually use datasets through an automated LLM-based filtering process using Qwen3-32B. The implementation of filtering can be found in extract_packages_and_files.py.

TL;DR: We prompted Qwen3-32B to identify and remove notebooks that either (1) had nothing to do with data analysis, or (2) didn’t actually use datasets. This step removed about 20% of the notebooks.

This ensured we trained only on relevant data science tasks.

Using the cleaned notebooks, we generated question–answer pairs using Qwen3-32B. The questions and answer are grounded in the real notebook traces so the QA pairs are based on real code execution results. Prompt design: we asked the LLM to produce natural questions that could realistically be asked of the dataset, then validated whether the notebook provided a correct answer.

Challenge: We had to try many prompts to get higher-difficulty questions because LLMs tended to generate trivial ones like "what is the size of the dataset".Insight: We broke this into two steps because LLMs t