Training and Evaluating a Jupyter Notebook Data Science Assistant

TL;DR

This paper introduces JuPyT5, a transformer-based model trained on a large dataset of Jupyter Notebooks, capable of understanding and solving data science problems with high accuracy, and presents a new evaluation metric for such tasks.

Contribution

The paper develops JuPyT5, a novel model trained on a comprehensive dataset of pedagogical notebooks, and introduces the DSP metric for evaluating data science problem-solving capabilities.

Findings

JuPyT5 solves 77.5% of DSP problems with 100 sampling attempts.

The DSP dataset effectively tests data science understanding and coding skills.

JuPyT5 outperforms existing benchmarks on natural language to code tasks.

Abstract

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.