Loading…
Jupyter Agents: training LLMs to reason with notebooks
Hugging FaceBaptiste Colle, Hanna Yukhymenko, Leandro von Werra
Summary
Small language models often struggle to compete with large frontier models on complex, agentic data science tasks. To improve notebook-based reasoning, researchers simplified agent scaffolding down to roughly two hundred lines of code with dedicated execution and final answer tools, boosting baseline easy accuracy on the DABStep benchmark from 44.4 percent to 59.7 percent. They constructed a curated training dataset by deduplicating two terabytes of Kaggle notebooks, automatically fetching five terabytes of linked datasets, and scoring educational value and relevance using Qwen3-32B. Question-answer pairs grounded in verified execution traces were generated to fine-tune compact Qwen3-4B thinking and instruct models. The team released the trained models alongside the dataset and execution sandboxes, establishing a foundation for reinforcement learning and distillation on notebook workflows.
Context
Small language models struggle to compete with large models on agentic data science tasks and multi-step notebook code execution, as seen in low baseline accuracy on the DABStep benchmark.
Approach / What changed
The authors created lightweight agent scaffolding of around 200 lines with code execution and final answer tools. They built a data pipeline with Datatrove to deduplicate Kaggle notebooks, download linked datasets, filter notebooks via Qwen3-32B educational scoring and relevance checks, and generate QA pairs grounded in real execution traces to fine-tune Qwen3-4B.
Takeaways
- Streamlining agent scaffolding to roughly 200 lines without external framework dependencies raised Qwen3-4B baseline accuracy on DABStep easy tasks from 44.4% to 59.7%.
- Kaggle notebook data curation required aggressive filtering, where deduplication eliminated approximately 90% of raw notebooks and Qwen3-32B educational scoring removed another 70%.
- Training data was created by generating question-answer pairs grounded in real code execution traces across executable Kaggle notebooks paired with their downloaded datasets.
Related reading
Accelerating Qwen3-8B Agent on Intel® Core™ Ultra with Depth-Pruned Draft Models
Local agentic AI workflows generate multi-step reasoning traces that increase token counts and strain real-time inference on personal computers. To accelerate execution, researchers implemented speculative decoding using OpenVINO.GenAI on an Intel Lunar Lake integrated GPU, pairing a 4-bit Qwen3-8B target model with a smaller Qwen3-0.6B draft model. By measuring angular distance across layers, the team pruned six of the draft model's twenty-eight layers to minimize draft latency and fine-tuned it on synthetic data generated from 500,000 prompts. This depth-pruned draft elevated speculative decoding throughput from a 1.3x speedup to a 1.4x speedup over the standalone baseline. Finally, deploying the optimized model alongside the Hugging Face smolagents framework allowed local agents to perform tool invocation and code generation tasks efficiently.
Igor Margulis, Ofir Zafrir, Shira Guskin, Guy Boudoukh, Pedro CuencaRed Hat ·