PaperlayerView exact revision on arXiv

Source-linked AI summary

Training and Evaluating a Jupyter Notebook Data Science Assistant

Shubham Chandel, Colin B. Clement, Guillermo Serrato, Neel Sundaresan

arXiv:2201.12901v1cs.LGcs.SE

TL;DR

Complex code synthesis remains challenging, and CodeBLEU is not useful for determining program correctness. The paper introduces an executable notebook benchmark and trains JuPyT5 with cell infilling, reporting that it solves over 77% of DSP problems.

  • Problem

    Complex code synthesis remains challenging, and pass rate has only a weak correlation with CodeBLEU, which the authors find unsuitable for determining program correctness.

    Source: Discussion · p. 7, Conclusion · p. 7

  • Method

    The authors train JuPyT5 on Jupyter notebooks using cell infilling, predicting a target cell from preceding context and, in the best model, one following cell.

    Source: Datasets · p. 2, Code Infilling Pretraining · p. 4

  • Results

    JuPyT5 solves over 77% of DSP problems, while the authors report that CodeBLEU is not useful for determining hypothesis-program correctness.

    Source: Conclusion · p. 7, Discussion · p. 7

  • Takeaways & Limitations

    The authors conclude that DSP can help researchers address modeling challenges in complex code synthesis and that these models show feasibility for a data science assistant.

    Source: Conclusion · p. 7

    The over-77% result is an optimistic estimate because deployment may not tolerate 100 hypotheses.

    Source: Discussion · p. 6

Abstract

from arXiv · show

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.

Datasets

The paper constructs Data Science Problems (DSP) from pedagogical notebooks and trains on a broad GitHub notebook corpus, examining the corpus’s content and a Markdown-focused subset.

Datasets

  • The training corpus uses public GitHub repositories labeled primarily as Jupyter Notebook repositories, with DSP-source repositories excluded.The repositories were cloned and processed in April 2021; duplicate holdout examples were also removed to prevent leakage.Source: Datasets · p. 2, Pretraining dataset · p. 4

Data Science Problems

Pretraining dataset

  • In a 17K-notebook t-SNE sample, Markdown fraction separates notebook clusters; manually inspected low-Markdown clusters include scratchpads, research code, and personal projects, while high-Markdown notebooks include tutorials and university assignments.Cluster labels were assigned after manually inspecting about 10 notebooks per cluster.Source: Pretraining dataset · p. 4, Figure 3 · p. 4
  • The Markdown-focused subset requires at least one code cell and at least one-third Markdown cells, yielding 4.1M notebooks and 15.7B tokens.It was selected to test whether focusing on literate code improves DSP performance.Source: Training Subset: Markdown Focused · p. 4

Models

JuPyT5 uses a sequence-to-sequence transformer initialized from PyMT5 and trained to infill notebook cells from surrounding context, with control codes specifying the target type.

Models

  • JuPyT5 uses a BART architecture initialized from a PyMT5 checkpoint, with PyMT5’s training hyperparameters as its starting point.BART reconstructs masked spans, while PyMT5 reconstructs a syntactically selected part of a Python method.Source: Models · p. 4, Code Infilling Pretraining · p. 4

Code Infilling Pretraining

  • Cell infilling makes each notebook cell a target, with neighboring cells as source; experiments use one or three preceding context cells, and the best model also sees one following cell.The following cell can include the grading cell.Source: Code Infilling Pretraining · p. 4
  • Five control tokens specify whether the target is Markdown, code, a function, a class, or an import.The function, class, and import tokens were for studies not included in the manuscript.Source: Control Codes · p. 4

Training Details

  • Each model was trained for five epochs on either the full training set or the Markdown-focused subset using 80 32GB Tesla V100 GPUs.Batch size was adjusted for data parallelism; other hyperparameters followed PyMT5.Source: Training Details · p. 4

Experiments and Results

JuPyT5’s DSP performance improves with access to grading cells and additional notebook context, while Markdown-focused training brings a modest gain. Results on HumanEval and MBPP vary by benchmark, and the discussion notes limits from sampling cost, difficult compound tasks, and weak alignment between CodeBLEU and passing tests.

Evaluation Details

  • DSP counts a problem as passed only when generated code passes its grading-cell unit test, and evaluates success probability with pass@k over sampled hypotheses.Evaluation preserves the notebook’s other correct cells to avoid failures caused by propagating generated errors through dependent cells.Source: Evaluation Details · p. 5

DSP Results

  • On DSP, JuPyT5 pass@k across k=1–100 showed modest gains for the C=3 context model trained on a Markdown-focused subset versus the whole training set, most pronounced near k=100 despite a 2/5 data-size reduction.Source: DSP Results · pp. 5–6, Figure 4 · p. 5
  • Seeing the grading cell yields a dramatic pass@k gain, with cell-infilling pass@1 comparable to baseline pass@100; C=3 context adds about 10% across k.The experiments’ summary also reports that more context helps, Markdown-focused training gives a modest improvement, and seeing tests gives a large boost.Source: Figure 5 · p. 5, Figure 6 · p. 5, DSP Results · p. 6

HumanEval and MBPP Results

  • On HumanEval, Codex beats JuPyT5 except with the 85M-parameter model; on MBPP, JuPyT5 beats the Program Synthesis model except its largest model.The authors suggest that Markdown-versus-docstring formatting may explain the HumanEval gap; presenting docstrings as Markdown improved performance.Source: HumanEval and MBPP Results · p. 6, Table 5 · p. 6, Table 6 · p. 6

Discussion

  • The reported over-77% DSP success uses 100 hypotheses, which the authors call optimistic because deployment may not tolerate that many.Selecting one hypothesis by log-likelihood per token offered only a modest average improvement over a single sample.Source: Discussion · p. 6
  • An undefined L1-norm task requires learned knowledge of the norm, while compound operation chains such as embedded-model least squares are difficult for these models.The authors note that larger models might improve the undefined-norm case and that compound chaining was also reported as a challenge by Chen et al. (2021) and Austin et al. (2021).Source: Discussion · pp. 6–7
  • Pass rate and average CodeBLEU score have only a weak correlation across DSP problems, which the authors say makes CodeBLEU unhelpful for judging program correctness.Figure 7 plots problems in order of JuPyT5’s pass rate and overlays the average CodeBLEU score of 100 generated samples.Source: Figure 7 · p. 5, Discussion · p. 7

Conclusion

JuPyT5 solves over 77% of DSP problems, supporting the feasibility of a data science assistant, though complex code synthesis remains challenging.

Conclusion

  • Over 77% of DSP problems are solved by JuPyT5, trained on almost all publicly available Jupyter Notebooks.The paper describes this as an optimistic estimate of the feasibility of a data science assistant.Source: Conclusion · p. 7
  • DSP comprises over 1,000 executable problems, many with data dependencies, and each evaluated with unit tests.The benchmark is presented as a resource for addressing remaining modeling challenges in complex code synthesis.Source: Conclusion · p. 7
  • Seeing the subsequent grading cell improves cell infilling partly by preventing predictions of potentially erroneous assert statements.The baseline sometimes predicts assert statements alongside its hypothesis code.Source: Figure 8 · p. 7
  • JuPyT5 sometimes uses template-like behavior, adapting comments from code similar to the prompt.The paper illustrates this behavior with an example prediction.Source: Figure 9 · p. 7
  • JuPyT5 rarely solves examples requiring chained logic or relying on an unstated definition such as the L1 norm.The paper identifies compound operations as difficult for these models and suggests larger models might help with the undefined norm case.Source: Figure 10 · p. 7
Loading 2201.12901v1…