PaperlayerView exact revision on arXiv

Source-linked AI summary

Natural Language to Code Generation in Interactive Data Science Notebooks

Pengcheng Yin, Wen-Ding Li, Kefan Xiao, Abhishek Rao, Yeming Wen, Kensen Shi, Joshua Howland, Paige Bailey, Michele Catasta, Henryk Michalewski, Alex Polozov, Charles Sutton

arXiv:2212.09248v1cs.CLcs.SE

TL;DR

Notebook code-generation benchmarks need to reflect concise user intents and the rich, multi-turn context of interactive notebooks. This paper introduces ARCADE and develops PACHINCO, which outperforms public code LMs on the benchmark; it also explores prompting strategies that improve prediction diversity and self-consistency reranking accuracy.

  • Problem

    Notebook code generation involves rich, multi-turn contexts, while existing benchmarks can use verbose tutorial prompts or require extra tests and examples.

    Source: §1 · p. 2

  • Method

    The paper presents ARCADE for notebook data-wrangling and EDA code generation and develops PACHINCO, a 62B language model tailored to data science.

    Source: §1 · p. 2, §9 · p. 15

  • Results

    PACHINCO outperforms public code LMs on ARCADE, and explanation-style few-shot prompting improves prediction diversity and self-consistency reranking accuracy.

    Source: §1 · p. 2, §9 · p. 15

  • Takeaways & Limitations

    Explanation-style prompting can improve self-consistency reranking accuracy by increasing the diversity of model predictions.

    Source: §1 · p. 2

    ARCADE focuses on data wrangling and EDA rather than broader problem types in existing datasets; expanding its task spectrum is future work.

    Source: §3.2 · p. 6

Abstract

from arXiv · show

Computational notebooks, such as Jupyter notebooks, are interactive computing environments that are ubiquitous among data scientists to perform data wrangling and analytic tasks. To measure the performance of AI pair programmers that automatically synthesize programs for those tasks given natural language (NL) intents from users, we build ARCADE, a benchmark of 1082 code generation problems using the pandas data analysis framework in data science notebooks. ARCADE features multiple rounds of NL-to-code problems from the same notebook. It requires a model to understand rich multi-modal contexts, such as existing notebook cells and their execution states as well as previous turns of interaction. To establish a strong baseline on this challenging task, we develop PaChiNCo, a 62B code language model (LM) for Python computational notebooks, which significantly outperforms public code LMs. Finally, we explore few-shot prompting strategies to elicit better code with step-by-step decomposition and NL explanation, showing the potential to improve the diversity and explainability of model predictions.

1 Introduction

Notebook code generation must account for multimodal, interrelated context, while existing benchmarks often lack realistic intents and complex dependencies. The paper introduces ARCADE and PACHINCO, then explores prompting strategies for more diverse and understandable predictions.

1 Introduction

  • Notebook tasks mix text, code, graphics, and execution results, with related intents and dependencies that can reach across distant cells.Problems may share execution context, form coherent turns, or depend on earlier cells beyond the immediately preceding turn.Source: §1 · p. 2
  • Prior benchmarks often use verbose tutorial questions, require user-provided tests or examples, or isolate tasks instead of capturing realistic interrelated notebook problems.The authors identify a need for realistic intents, rich notebook context, and multiple related problems.Source: §1 · p. 2
  • ARCADE contains 1,082 problems across 136 notebooks and 106 ML datasets, with professional pandas solutions and more than 67% of problems requiring at least five pandas API calls.65% of problems and notebooks were created from scratch to mitigate data-leakage risk.Source: §1 · p. 2
  • PACHINCO is a 62B code LM for Python computational notebooks that significantly outperforms public code LMs on ARCADE.It is trained on a mixture of general-domain natural language, source code, and Jupyter notebooks.Source: §1 · p. 2
  • Few-shot prompting with step-by-step solutions and inline natural-language explanations can diversify predictions and improve self-consistency reranking accuracy.The paper notes that this style may also be more understandable to novice data scientists.Source: §1 · p. 2

2 Problem: Code Generation in Computational Notebooks

The task is to generate code for a requested next notebook cell from the partial notebook and user intent. The paper focuses on data-wrangling and exploratory-analysis programs using pandas.

2 Problem: Code Generation in Computational Notebooks

  • Given a partial notebook and an intent for the next cell, the system generates code for that cell; later turns use the updated notebook context.The context includes preceding cells and earlier user intents.Source: §2 · pp. 2–3
  • The target tasks are data wrangling, which cleans and transforms data, and exploratory data analysis, which queries data for insights.The stated focus is code using pandas, whose DataFrames and Series represent tables and columns.Source: §2 · p. 3

3 ARCADE: A Benchmark of pandas Data Science Code Generation

ARCADE brings together concise pandas code-generation tasks from existing and newly created notebooks, including related prompts in shared notebook contexts. Its analysis characterizes task ambiguity and complexity, while its evaluation metric accommodates alternative output structures.

3.1.1 Mining Examples from Existing Notebooks

  • Existing-notebook tasks were mined using pandas-call or short-question heuristics, then annotated with concise intents and added output specifications where needed.About 45% of intents remained under-specified despite those specifications.Source: §3.1.1 · pp. 3–4

3.1.2 Creating Notebooks with Examples from Scratch

  • 660 scratch-built problems target 70 recent Kaggle datasets, with annotators creating notebook task sequences and peer review checking each notebook.Tasks were encouraged to require at least five pandas API calls.Source: §3.1.2 · p. 4

3.2 Dataset Analysis

  • 67% of New Tasks require at least five pandas API calls, while about 45% of intents are under-specified and notebook context disambiguates about 30% of those cases.The dataset contains 1,082 problems across 133 notebooks and 106 unique ML datasets.Source: §3.2 · pp. 4–5, Figure 2 · p. 6
  • ARCADE combines concise intents, interrelated notebook problems, grounded language understanding, and recent held-out datasets; unlike tutorial benchmarks, it avoids extra unit-test input, but covers wrangling and EDA rather than broader tasks.The authors identify expanding the task spectrum as future work.Source: §3.2 · pp. 5–6

3.3 Evaluation by Fuzzy Output Matching

  • The evaluation heuristics canonicalize container types and accept DataFrames containing every reference column and its cell values, accommodating alternative output structures.The authors report a false-negative rate under 10% for identifying solutions with alternative output structures.Source: §3.3 · pp. 6–7

4 PACHINCO: Adapting Code LMs to Computational Notebooks

PACHINCO adapts a 62B PALM model to Python and computational notebooks through successive fine-tuning stages. Notebook training uses cell-level near-deduplication against evaluation notebooks to reduce overlap.

4 PACHINCO: Adapting Code LMs to Computational Notebooks

  • PACHINCO starts from PALM 62B and fine-tunes first on 64B tokens of deduplicated Python source code.The initial PALM model was trained on 1.3T tokens across conversational, webpage, and code data.Source: §4 · p. 7
  • A second fine-tuning stage uses about 3.8M Jupyter notebooks totaling about 9.6B tokens, with notebooks linearized into Python code and markdown converted to comments.Notebooks with fewer than four cells are filtered out.Source: §4 · p. 7
  • Cell-level fuzzy near-deduplication against Existing Tasks removes training notebooks sharing a cell cluster with evaluation notebooks, eliminating about 350K notebooks.This preprocessing is intended to mitigate evaluation-notebook leakage into training data.Source: §4 · p. 7

5 Experiments

The experiments evaluate code language models on ARCADE and test how notebook context and few-shot prompting affect code generation. PACHINCO leads the evaluated models on both splits, while step-by-step prompting improves code style and supports stronger self-consistency reranking.

5.1 Models and Setup

  • The evaluation compares PACHINCO with CODEGEN, INCODER, base PaLM, and a Python-fine-tuned PaLM model, estimating pass@k from 50 sampled completions per problem.Pass@k is the fraction of problems with at least one correct sample among k samples.Source: §5.1 · pp. 7–8

5.2 LM Prompting Strategies

  • The experiments contrast notebook-context-only prompts with few-shot exemplars added before the context to steer predicted code style.Few-shot prompts are intended to impose control over the style of generated code.Source: §5.2 · p. 8, §5.2.2 · p. 8

5.2.1 Prompting using Notebook Context

  • Default prompts include preceding notebook cells, the current natural-language intent, DataFrame schema descriptions, and prior reference solutions for later problems.Schema descriptions include columns and example DataFrame values.Source: §5.2.1 · p. 8

5.2.2 Few-shot Prompting

  • The four few-shot strategies range from chained one-line pandas calls to step-by-step code, optionally with a preamble and inline natural-language explanations.Each strategy uses four exemplars in its prompt prefix.Source: §5.2.2 · pp. 8–9

5.3 Results

  • PACHINCO achieves the best performance on both ARCADE splits; comparisons also show that model size and domain-specific fine-tuning distinguish results across the evaluated models.The authors note that New Tasks are harder and may be less exposed to training-data leakage than Existing Tasks.Source: Table 2 · p. 9, §5.3 · pp. 9–10
  • New Tasks remain harder than Existing Tasks even at similar problem complexity, and the authors note possible training-data leakage for some Existing Tasks.Removing schema descriptions substantially worsens results, especially on New Tasks.Source: §5.3 · p. 10
  • More notebook context improves pass@30: PACHINCO rises from 44% to 72% on Existing Tasks and from 17% to 36% on New Tasks when context grows from one to three cells.Gains begin to plateau after roughly five to seven context cells; more context also helps reduce schema-understanding errors.Source: §5.3 · p. 10

5.4 Few-shot Prompting Beyond Test Notebook Contexts

  • Step-by-step prompting improves pass rate and code decomposition; adding explanations yields still more decomposition with similar accuracy and more comment lines.Step-by-step prompting also produces more diverse solutions and improves self-consistency reranking; its one-sample accuracy with explanations is on par with PACHINCO pass@5.Source: §5.4 · pp. 11–13

6 Error Analysis

The error analysis finds that incorrect predictions include both execution failures and executable but functionally incorrect code. Errors reflect missing intermediate-data context, complex tasks, intent interpretation, underspecification, and limits in fuzzy output matching.

6 Error Analysis

  • Roughly 35% of incorrect PACHINCO predictions under fuzzy output matching result in execution errors, while 65% execute but are functionally incorrect.The analysis also reports a significant drop in execution errors after two-stage code fine-tuning.Source: §6 · p. 13
  • Missing schemas for intermediate DataFrames contribute to KeyError and AttributeError; pd.groupby accounts for more than 50% of these errors.The prompts describe loaded DataFrame schemas, but prompt-length limits leave later derived DataFrames undescribed.Source: §6 · p. 13
  • In 50 sampled incorrect predictions, complex problems accounted for 43%, intent-interpretation errors 26%, and under-specified intents 19%.False negatives from fuzzy matching and annotation errors each accounted for 6%.Source: §6 · p. 13
  • The authors identify complex problems as the primary error source and intent misunderstandings as room to improve instruction following.They also note that under-specified intents are common with short instructions, while fuzzy matching can miss valid output representations such as nested columns or different orientations.Source: §6 · pp. 13–14

7 Case Study: How Useful is Predicted Code with Step-wise Explanations?

The qualitative study finds that step-wise natural-language explanations can make generated pandas code easier to follow, compare, and verify. The authors present these observations as potential directions for improving code LMs’ utility to developers.

7 Case Study: How Useful is Predicted Code with Step-wise Explanations?

  • Step-wise explanations make chains of pandas transformations easier to follow by clarifying dataflow, especially for complex computations or less intelligible schemas.Inline descriptions of intermediate outputs can also clarify advanced functions such as pd.unstack.Source: §7 · p. 14
  • High-level procedural explanations let users browse different solution approaches without being distracted by implementation details.The authors also report that explanations help users spot potentially incorrect steps and make corrections.Source: §7 · p. 14
  • The authors identify step-by-step explanations as a potential avenue for improving code LMs’ utility to developers.Their preceding experiments found that step-wise explanations improved solution diversity and self-consistency reranking accuracy.Source: §7 · p. 14

8 Related Works

The related work spans code language models, automation of data-science tasks, and context-driven code generation. This paper focuses on multi-turn notebook tasks involving pandas and interdependent context.

LLMs for Code

Automating Data Science

  • Prior data-science automation includes AutoML for feature engineering and model selection, while this paper targets data wrangling and exploratory data analysis.Prior pandas synthesis approaches use input/output examples, natural-language intents with examples, or unit tests [Bavishi et al., 2019; Jain et al., 2021; Chandel et al., 2022].Source: §8 · p. 14
  • ARCADE includes 57 plotting examples, but the paper does not use them and leaves systematic evaluation of plotting tasks for future work.Source: §8 · p. 14

Context-driven Code Generation

  • Unlike context-driven work on structured-database queries, ARCADE includes more complex query patterns and additional data-wrangling tasks using pandas and Python.Evaluation uses ground-truth solutions for previous turns; conditioning on model-predicted solutions is left for future work.Source: §8 · p. 15

9 Conclusion

The paper presents ARCADE, a notebook benchmark for data-wrangling and exploratory-data-analysis code generation, and PACHINCO, a 62B language model tailored for data science. It reports that PACHINCO outperforms public code LMs on ARCADE and that few-shot learning can improve code style and solution diversity.

9 Conclusion

  • ARCADE benchmarks code generation for data-wrangling and exploratory-data-analysis tasks in computational notebooks.Source: §9 · p. 15
  • PACHINCO is a 62B language model tailored for data science that outperforms public code LMs on ARCADE.Source: §9 · p. 15
  • Few-shot learning with PACHINCO improves code style and solution diversity.Source: §9 · p. 15

Supplementary Materials

The supplementary materials describe how ARCADE tasks and intents are annotated, then detail PACHINCO’s two-stage fine-tuning and notebook linearization.

A Outline of ARCADE Annotation Guideline

  • For existing tasks, annotators identify instructive code cells, clean notebooks, then write and verify natural-language intents and code solutions while removing hints.The workflow also removes redundant notebook text that could reveal reference solutions.Source: A Outline of ARCADE Annotation Guideline · p. 20
  • Intent-writing guidance favors concise, natural instructions that clearly describe desired outputs without redundant explanation or ambiguity.Annotators are asked to imagine realistic inputs to an AI programmer.Source: A Outline of ARCADE Annotation Guideline · p. 20
  • For new tasks, annotators create notebook snippets for supplied ML datasets, write the intent first, and then provide a multiline solution using different pandas APIs.Source: A Outline of ARCADE Annotation Guideline · p. 20
  • Tasks cover data wrangling and EDA, with guidance to favor multistep problems using multiple pandas APIs; a heuristic targets at least 50 tokens or four API uses.Examples include preprocessing needed before answering an EDA question.Source: A Outline of ARCADE Annotation Guideline · p. 21

B Details of Fine-tuning PACHINCO

  • Fine-tuning starts from PaLM 62B and proceeds through Python-only training followed by Jupyter-notebook training.The notebook stage uses 10B tokens for three epochs, approximately 572K steps.Source: B Details of Fine-tuning PACHINCO · p. 21
  • Notebook markdown and code cells are concatenated as Python source using “# In[]:”, with markdown text commented out.The authors state this format is common in GitHub repositories and helps mitigate domain transfer from general Python code to notebooks.Source: B Details of Fine-tuning PACHINCO · pp. 21–22

C Performance of Python-finetuned PALM 62B on Code Benchmarks

The supplementary analyses compare model performance across benchmarks, inference settings, task complexity, context sizes, and prompting conditions. They also describe scaling behavior and a comparison with CODEX.

C Performance of Python-finetuned PALM 62B on Code Benchmarks

D Inference Setup

  • Inference truncates notebook context at 900 tokens, always includes NL schema descriptions, and permits up to 2,100 prompt tokens for few-shot experiments.Few-shot prompting is performed only on PACHINCO because its rotary positional embedding can generalize to longer inference contexts.Source: D Inference Setup · p. 22

E pass@1 for PACHINCO

F Scaling Curve of CODEGEN on ARCADE

  • CODEGEN’s ARCADE pass rate scales nearly linearly with model size and has not saturated, especially on New Tasks, where its scaling curve has a smaller slope.The authors interpret this as evidence that ARCADE can study scaling behavior and that New Tasks are more challenging for CODEGEN.Source: F Scaling Curve of CODEGEN on ARCADE · pp. 22–23, Figure 13 · p. 23

G Comparing with CODEX on ARCADE

  • PACHINCO significantly outperforms CODEX-cushman-001, matches CODEX-davinci-002 on Existing Tasks, but trails it on New Tasks.Error analysis finds CODEX-davinci-002 better at following instructions involving complex DataFrame schema descriptions.Source: G Comparing with CODEX on ARCADE · p. 23, Table 5 · p. 23

H Performance on Problems with Different Complexity

I Distribution of Execution Error Types w.r.t Context Size

J Few-shot Prompting Results on Existing Tasks Split

  • On Existing Tasks, few-shot prompting is less effective than on New Tasks; explanations seem to hurt simpler problems but help those with more than seven pandas API calls.The authors suggest memorization may make extra exemplars less effective on Existing Tasks.Source: J Few-shot Prompting Results on… · p. 24

K Further Analysis of Solution Diversity

The analysis finds that step-by-step prompting and natural-language explanations produce more diverse solution patterns, including among correct predictions.

  • Step-by-step prompting increases the diversity of pandas API sequences across all, executable, and correct predictions; NL explanations increase it further.The figure’s cumulative distributions use curves farther to the right to represent greater diversity.Source: K Further Analysis of Solution Diversity · p. 24
  • Greater diversity can help find a correct solution at larger sample sizes and give users varied correct approaches for learning or personal preference.Source: K Further Analysis of Solution Diversity · p. 25

L Error Analysis Examples

The error analysis finds that semantic errors account for most incorrect predictions, while execution failures often involve dataframe-schema mistakes. Examples also show complex reasoning demands and evaluation mismatches.

M Data Card for the Training Data of PACHINCO

This section documents PACHINCO’s source-code training data, including its preprocessing and limitations, and presents prompting examples. The examples cover pandas exercises and notebook-context and few-shot setups, including prompts that combine existing code, dataframe schema, and a new question.

M Data Card for the Training Data of PACHINCO

  • The card notes licensing-classifier errors, no explicit screening for potentially offensive text, and that the static dataset will become stale.It says the dataset should not be used for unacceptable code or language modeling, such as generating toxic or biased content.Source: M Data Card for the Training Data of… · pp. 31–33
  • Training data was a small random sample of publicly available GitHub Python files and Jupyter notebooks, with restrictive-license code filtered out.The data was collected by Google researchers using software to extract and clean source files.Source: M Data Card for the Training Data of… · pp. 31–32
  • Preprocessing applied license and quality filters and deduplication; notebooks were also deduplicated at the cell level against the evaluation dataset.Python files were deduplicated at file level using a custom minhash variant; lower-level redundancies may remain.Source: M Data Card for the Training Data of… · pp. 31–32

N Detailed Prompting Examples

  • Prompt examples combine notebook context with prior problems and schema descriptions; few-shot variants prepend exemplar prefixes, with four exemplars per prefix.Three prefixes per style were evaluated across four styles, with results averaged over three restarts; prompt engineering was limited and prefixes predated collection of 70% of the dataset.Source: N Detailed Prompting Examples · p. 34
  • Exercise prompts cover employee counts by gender and birth year, science-score averages, California population, and each brand’s most expensive phone.The examples pair dataframe schemas and task statements with pandas solutions.Source: N Detailed Prompting Examples · pp. 41–44
  • Another prompt setup demonstrates counting countries hosting at least two Olympics, plotting letter-grade counts, and selecting the most expensive phone per brand.The examples include loading data and showing schema, problem, and solution cells.Source: N Detailed Prompting Examples · pp. 44–45
  • A Game Pass notebook context includes code to parse playtime ranges and dates before asking which year the most-played game was added.The prompt shows the loaded dataset and schema, then feature extraction for minimum and maximum hours and date parsing for ADDED.Source: N Detailed Prompting Examples · p. 46
Loading 2212.09248v1…