Jupiter: Enhancing LLM Data Analysis Capabilities via Notebook and Inference-Time Value-Guided Search

TL;DR

This work presents NbQA, a large-scale dataset of authentic, multi-step data-analysis tasks extracted from real Jupyter notebooks, and Jupiter, a value-guided search framework that treats notebook solving as an inference-time MCTS problem guided by a learned value model. By combining supervised fine-tuning on NbQA with value-guided search, the approach significantly boosts the multi-step data-analysis capabilities of open-source LLMs, matching or surpassing strong commercial baselines on InfiAgent-DABench and showing strong generalization to DSBench and AIME. The results demonstrate substantial improvements in tool-use reasoning and reasoning efficiency, illustrating a scalable path to more capable notebook-based data analysis with LLMs. The work also provides detailed dataset construction methods and experimental settings, enabling replication and extension by the community.

Abstract

Large language models (LLMs) have shown great promise in automating data science workflows, but existing models still struggle with multi-step reasoning and tool use, which limits their effectiveness on complex data analysis tasks. To address this, we propose a scalable pipeline that extracts high-quality, tool-based data analysis tasks and their executable multi-step solutions from real-world Jupyter notebooks and associated data files. Using this pipeline, we introduce NbQA, a large-scale dataset of standardized task-solution pairs that reflect authentic tool-use patterns in practical data science scenarios. To further enhance multi-step reasoning, we present Jupiter, a framework that formulates data analysis as a search problem and applies Monte Carlo Tree Search (MCTS) to generate diverse solution trajectories for value model learning. During inference, Jupiter combines the value model and node visit counts to efficiently collect executable multi-step plans with minimal search steps. Experimental results show that Qwen2.5-7B and 14B-Instruct models on NbQA solve 77.82% and 86.38% of tasks on InfiAgent-DABench, respectively-matching or surpassing GPT-4o and advanced agent frameworks. Further evaluations demonstrate improved generalization and stronger tool-use reasoning across diverse multi-step reasoning tasks. Code and data are available at https://github.com/microsoft/Jupiter.

Figures (13)

• Figure 1: Accuracy by Questions (%) across different models on InfiAgent-DABench under three settings: base model, SFT on NbQA, and SFT on NbQA + Jupiter search, demonstrating the base models suffer from accurate multi-step data analysis.
• Figure 2: Construction of the NbQA dataset.
• Figure 3: Overview of Jupiter, including (a) trajectory collection and (b) inference strategy. The collected trajectories are used for training the value model (VM), and VM is used for value estimation during inference.
• Figure 4: Task completion rates of QWen2.5-7B-Instruct and QWen2.5-14B-Instruct on the DSBench data modeling task as the number of iterations increases. “w/o VM” and “with VM” denote whether a trained value model is used. For comparison, when the maximum number of iterations is limited to 40, the task completion rates of QWen2.5-7B-Instruct and QWen2.5-14B-Instruct using ReAct are 63.51% and 66.22%, respectively.
• Figure 5: The distribution of the number of data files processed by tasks for tasks that have complete data dependency files and exhibit low randomness