Jupyter Agents: Training LLMs to Reason with Notebooks
Hugging Face has developed a pipeline to train small language models to function as data science agents capable of executing code within Jupyter Notebooks. By combining a curated synthetic dataset derived from Kaggle notebooks with simplified agent scaffolding, the team successfully improved the performance of Qwen3-4B models on data analysis benchmarks.
The DABStep Benchmark for Data Science Agents
To measure progress in agentic data science capabilities, Hugging Face utilizes the DABStep benchmark, which evaluates agents on realistic tasks by providing datasets and requiring the model to answer non-trivial data questions. This benchmark is notably challenging for current LLMs; for example, Claude 4 Sonnet achieves less than 20% accuracy on the benchmark's hard tasks.
Optimizing Agent Scaffolding
Agent performance is heavily dependent on the scaffolding—the code and prompting structures that steer model behavior. Hugging Face restructured the scaffolding for the Jupyter Agent to be lightweight and aligned with the model's needs:
- Simplified Architecture: The scaffolding was stripped down to approximately 200 lines of code with no external dependencies.
- Tool Integration: The agent operates in a while loop using two primary tools:
code executionandfinal_answer. - Performance Gain: This optimization alone increased accuracy on the DABStep easy split from 44.4% to 59.7%.
The Jupyter Agent Dataset Pipeline
To improve small models, Hugging Face created a multi-stage pipeline using Datatrove to transform ~2TB of raw Kaggle notebooks into a high-quality training set of 51k synthetic notebooks (approximately 0.2B tokens).
1. Deduplication and Data Acquisition
Raw Kaggle notebooks were deduplicated, reducing the volume from 2TB to 250GB. The team then used the kagglehub package to download approximately 5TB of linked datasets, filtering out model checkpoints, multimodal corpora, and datasets exceeding 10GB to ensure they could fit within E2B sandboxes.
2. Educational Scoring and Filtering
Using Qwen3-32B, the team implemented an "edu scoring" system to rate notebooks from 1–5 based on clarity and educational value, removing roughly 70% of the notebooks. An additional 20% were removed via LLM-based filtering to exclude notebooks unrelated to data analysis or those that did not utilize datasets.
3. Synthetic QA and Trace Generation
To create verifiable training data, the pipeline generated synthetic question-answer (QA) pairs grounded in real notebook traces.
- QA Generation: Qwen3-32B generated natural questions and answers, which were then verified by a second LLM to prevent hallucinations.
- Trace Generation: Qwen-3-Coder-480B-A35B-Instruct was used to generate clean, step-by-step code execution traces to answer the synthetic questions.
- Handling Missing Data: When datasets were unavailable, the model was prompted to act as a stateful Python code interpreter to simulate execution.
- Reasoning Emulation: To ensure the model provided commentary between code cells, the
commentfield was made a required part of the code execution tool call.
Training and Results
Fine-tuning was performed on Qwen3-4B using the TRL library. The team found that full-parameter fine-tuning outperformed PEFT (Parameter-Efficient Fine-Tuning) and that using assistant_loss_only=True provided a performance boost.
Performance Improvements
Training on the synthetic dataset significantly improved the model's performance on the DABStep easy split:
| Model | No. of Epochs | DABStep (Easy) |
|---|---|---|
| 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% |
The final tuned Qwen-4B model achieved a boost of up to 36% over the base model and 22% over the scaffolded base model on the easy score, while also showing improvements in the hard score.
Open Source Releases
Hugging Face has released the following resources to the community:
- Jupyter Agent Dataset: A collection of 51k synthetic notebooks.
- Tuned Models:
jupyter-agent-qwen3-4b-instructandjupyter-agent-qwen3-4b-thinking.
Future Directions
The team identified several paths for further improvement, including generating more challenging multi-step questions, scaling the volume of curated traces, investigating knowledge distillation, and building a Reinforcement Learning (RL) environment leveraging the existing verifiable QA setup.