# Training and Evaluating a Jupyter Notebook Data Science Assistant

arXiv: https://arxiv.org/abs/2201.12901 · HTML: https://mlsanity.com/paper/2201.12901/

## Key facts

- Cited by 31 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 22 references that are in the map.
- First posted on arXiv 30 Jan 2022.

## Builds on (most cited first, 18 of 18)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension](https://mlsanity.com/paper/1910.13461/) (2019)
- [Program Synthesis with Large Language Models](https://mlsanity.com/paper/2108.07732/) (2021)
- [CodeBERT: A Pre-Trained Model for Programming and Natural Languages](https://mlsanity.com/paper/2002.08155/) (2020)
- [Measuring Coding Challenge Competence With APPS](https://mlsanity.com/paper/2105.09938/) (2021)
- [CTRL: A Conditional Transformer Language Model for Controllable Generation](https://mlsanity.com/paper/1909.05858/) (2019)
- [CodeSearchNet Challenge: Evaluating the State of Semantic Code Search](https://mlsanity.com/paper/1909.09436/) (2019)
- [CodeBLEU: a Method for Automatic Evaluation of Code Synthesis](https://mlsanity.com/paper/2009.10297/) (2020)
- [IntelliCode Compose: Code Generation Using Transformer](https://mlsanity.com/paper/2005.08025/) (2020)
- [Gmail Smart Compose: Real-Time Assisted Writing](https://mlsanity.com/paper/1906.00080/) (2019)
- [PyMT5: multi-mode translation of natural language and Python code with transformers](https://mlsanity.com/paper/2010.03150/) (2020)
- [Pythia: AI-assisted Code Completion System](https://mlsanity.com/paper/1912.00742/) (2019)
- [JuICe: A Large Scale Distantly Supervised Dataset for Open Domain Context-based Code Generation](https://mlsanity.com/paper/1910.02216/) (2019)
- [Generating Accurate Assert Statements for Unit Test Cases using Pretrained Transformers](https://mlsanity.com/paper/2009.05634/) (2020)
- [Generating Bug-Fixes Using Pretrained Transformers](https://mlsanity.com/paper/2104.07896/) (2021)

## Built on by (most cited first, 30 of 31)

- [A Survey on Large Language Models for Code Generation](https://mlsanity.com/paper/2406.00515/) (2024)
- [DS-1000: A Natural and Reliable Benchmark for Data Science Code Generation](https://mlsanity.com/paper/2211.11501/) (2022)
- [Large Language Models Meet NL2Code: A Survey](https://mlsanity.com/paper/2212.09420/) (2022)
- [Execution-Based Evaluation for Open-Domain Code Generation](https://mlsanity.com/paper/2212.10481/) (2022)
- [Natural Language to Code Generation in Interactive Data Science Notebooks](https://mlsanity.com/paper/2212.09248/) (2022)
- [WebGen-Bench: Evaluating LLMs on Generating Interactive and Functional Websites from Scratch](https://mlsanity.com/paper/2505.03733/) (2025)
- [PanGu-Coder2: Boosting Large Language Models for Code with Ranking Feedback](https://mlsanity.com/paper/2307.14936/) (2023)
- [Unifying the Perspectives of NLP and Software Engineering: A Survey on Language Models for Code](https://mlsanity.com/paper/2311.07989/) (2023)
- [Agents in Software Engineering: Survey, Landscape, and Vision](https://mlsanity.com/paper/2409.09030/) (2024)
- [xCodeEval: A Large Scale Multilingual Multitask Benchmark for Code Understanding, Generation, Translation and Retrieval](https://mlsanity.com/paper/2303.03004/) (2023)
- [SelfEvolve: A Code Evolution Framework via Large Language Models](https://mlsanity.com/paper/2306.02907/) (2023)
- [SALLM: Security Assessment of Generated Code](https://mlsanity.com/paper/2311.00889/) (2023)
- [CodeScope: An Execution-based Multilingual Multitask Multidimensional Benchmark for Evaluating LLMs on Code Understanding and Generation](https://mlsanity.com/paper/2311.08588/) (2023)
- [Fault-Aware Neural Code Rankers](https://mlsanity.com/paper/2206.03865/) (2022)
- [Execution-based Evaluation for Data Science Code Generation Models](https://mlsanity.com/paper/2211.09374/) (2022)
- [InsightBench: Evaluating Business Analytics Agents Through Multi-Step Insight Generation](https://mlsanity.com/paper/2407.06423/) (2024)
- [MHPP: Exploring the Capabilities and Limitations of Language Models Beyond Basic Code Generation](https://mlsanity.com/paper/2405.11430/) (2024)
- [LLM for SoC Security: A Paradigm Shift](https://mlsanity.com/paper/2310.06046/) (2023)
- [A Survey on Large Language Model-based Agents for Statistics and Data Science](https://mlsanity.com/paper/2412.14222/) (2024)
- [From Code Foundation Models to Agents and Applications: A Comprehensive Survey and Practical Guide to Code Intelligence](https://mlsanity.com/paper/2511.18538/) (2025)
- [JavaBench: A Benchmark of Object-Oriented Code Generation for Evaluating Large Language Models](https://mlsanity.com/paper/2406.12902/) (2024)
- [GitChameleon: Unmasking the Version-Switching Capabilities of Code Generation Models](https://mlsanity.com/paper/2411.05830/) (2024)
- [Capture the Flag: Uncovering Data Insights with Large Language Models](https://mlsanity.com/paper/2312.13876/) (2023)
- [CodeInsight: A Curated Dataset of Practical Coding Solutions from Stack Overflow](https://mlsanity.com/paper/2409.16819/) (2024)
- [Benchmarking AI Models in Software Engineering: A Review, Search Tool, and Unified Approach for Elevating Benchmark Quality](https://mlsanity.com/paper/2503.05860/) (2025)
- [Measuring Data Science Automation: A Survey of Evaluation Tools for AI Assistants and Agents](https://mlsanity.com/paper/2506.08800/) (2025)
- [Safurai 001: New Qualitative Approach for Code LLM Evaluation](https://mlsanity.com/paper/2309.11385/) (2023)
- [Safurai-Csharp: Harnessing Synthetic Data to improve language-specific Code LLM](https://mlsanity.com/paper/2311.03243/) (2023)
- [Linguacodus: A Synergistic Framework for Transformative Code Generation in Machine Learning Pipelines](https://mlsanity.com/paper/2403.11585/) (2024)
- [JupOtter: Cell-Level Bug Detection in Jupyter Notebooks](https://mlsanity.com/paper/2606.23877/) (2026)

## Abstract (arXiv)

We study the feasibility of a Data Science assistant powered by a sequence-to-sequence transformer by training a new model JuPyT5 on all publicly available Jupyter Notebook GitHub repositories and developing a new metric: Data Science Problems (DSP). DSP is a collection of 1119 problems curated from 306 pedagogical notebooks with 92 dataset dependencies, natural language and Markdown problem descriptions, and assert-based unit tests. These notebooks were designed to test university students' mastery of various Python implementations of Math and Data Science, and we now leverage them to study the ability of JuPyT5 to understand and pass the tests. We analyze the content of DSP, validate its quality, and we find that given 100 sampling attempts JuPyT5 is able to solve 77.5% of the DSP problems. We further present various ablation and statistical analyses and compare DSP to other recent natural language to code benchmarks.