# 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

Google Inc  Correspondence to<EMAIL_ADDRESS>
###### Abstract

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
$1,082$ 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.111Dataset will
be released shortly.

⬇

}

`\minted@fvset}`

Natural Language to Code Generation in Interactive
Data Science Notebooks

Pengcheng Yin††thanks:  Correspondence to<EMAIL_ADDRESS>Wen-Ding Li, Kefan
Xiao, Abhishek Rao, Yeming Wen, Kensen Shi, Joshua Howland, Paige Bailey,
Michele Catasta, Henryk Michalewski, Alex Polozov, Charles Sutton Google Inc.

##  1 Introduction

Figure 1: An example of a computational notebook adapted from our dataset,
with examples of reading and preprocessing data (cell $c_{1}$), data wrangling
(cell $c_{2},c_{3}$), and data analysis (cells $c_{3}-c_{7}$). Annotated NL
intents are shown in green.

Data science is the process of extracting insights from data Wang et al.
(2021a), and has become an integral part of decision making and knowledge
discovery Donoho (2017). Data scientists and machine learning (ML)
practitioners often use _computational notebooks_ , which are interactive
environments such as Jupyter notebooks Kluyver et al. (2016) and Google
Colab,222 https://colab.research.google.com/ in their work. As illustrated in
Fig. 1, data scientists spend significant amount of time on data wrangling and
exploratory data analysis (EDA) Agashe et al. (2019); Wang et al. (2022a).
This has motivated research on automating and accelerating the data science
workflow in general Aggarwal et al. (2019); Wang et al. (2021a, b), with
particular interest in data wrangling and EDA tasks Bavishi et al. (2019);
Jain et al. (2021); Karmaker et al. (2020); Nazabal et al. (2020); Kandel et
al. (2011).

On the other hand, large language models (LLMs) trained on code can assist
developers by translating natural language (NL) intents into executable
programs (Chen et al., 2021a; Austin et al., 2021; Chowdhery et al., 2022;
Nijkamp et al., 2022; Fried et al., 2022, _inter alia_), with promising
applications in synthesizing code for data wrangling and EDA tasks Jain et al.
(2021); Rajkumar et al. (2022); Cheng et al. (2022). Computational notebooks
also present unique challenges to LLMs, because notebooks freely mix NL, code,
graphics, and execution results Perkel (2021), and because of their
interactivity, notebooks feature multiple turns of related NL-to-code problems
Heyman et al. (2021). As illustrated in Fig. 1, those problems often have
interesting dependency structures. They may share common execution context
(_e.g._ DataFrame df), form semantically coherent turns (_e.g._
$c_{4},c_{5}$), or exhibit non-trivial long range data dependency (_e.g._ from
$c_{6}$ to $c_{2}$, or $c_{7}$ to $c_{3}$). This dependency structure is more
complex than existing multi-turn text to code tasks with sequentially
dependent problems Nijkamp et al. (2022).

Several benchmarks have been proposed to evaluate program synthesis of data
science programs from NL intents, but these datasets have several limitations.
First, some datasets derive from data science tutorial notebooks Agashe et al.
(2019); Chandel et al. (2022), but this data tends to contain NL text (_e.g._
exercise questions) which is more verbose and elaborate than the concise,
ephemeral style that developers write when interacting with code LMs (Barke et
al., 2022, more in § 3). Other datasets assume that the developer provides
extra information, such as unit tests or input/output examples Chandel et al.
(2022); Jain et al. (2022), but such systems pose an extra burden to users who
might not normally write such tests or examples during their workflow Pimentel
et al. (2019). Finally, existing datasets usually contain independent tasks
with isolated contexts (Lai et al., 2022), or limited number of contextually
dependent problems Huang et al. (2022), rather than having multiple, related
tasks with complex dependencies such as in Fig. 1. Therefore, there is a need
for a benchmark that provides _realistic NL intents, rich notebook context,
and multiple interrelated problems_ , so as to better reflect real-world usage
by data scientists.

To fill this gap, we present ARCADE,333Answer Repository for Computational
Analysis and Data Engineering a new benchmark for code generation for data
wrangling and EDA tasks in computational notebooks (§ 3). ARCADE consists of
$1,082$ problems spanning across 136 notebooks for $106$ ML datasets. ARCADE
features a series of NL utterances written by professional data scientists
with the intention of interacting with an AI pair programmer when working in a
notebook (_e.g._ , green texts in Fig. 1), with high-quality code solutions
using the pandas library. To mitigate the risk of data leakage when evaluating
LLMs, $65\%$ of the problems and their notebooks are created from scratch,
based on recent ML datasets on Kaggle.444https://www.kaggle.com/ ARCADE also
challenges LLMs with grounded language understanding, where a model needs to
ground relevant semantic concepts in intents (_e.g._ “min and max” in
$\bm{u}_{1}$) to variable states (_e.g._ the column df[’TIME’]) to understand
the intent. Besides featuring concise NL intents and contextually rich
problems, more than $67\%$ of problems in ARCADE also require complex
solutions using $5$ or more pandas API calls, making it more challenging than
most existing benchmarks.

To demonstrate how ARCADE can motivate new research, we evaluate public code
LMs, finding that they struggle on this task, especially newly created
problems from scratch. To build an LM tailored for data science, we develop
PaChiNCo, a $62\mathrm{B}$ code language model for Python computational
notebooks, trained on a mixture of general-domain NL, source code, and Jupyter
notebooks (§ 4). PaChiNCo significantly outperforms public code LMs on ARCADE.
Further, we explore few-shot prompting strategies to alter the style of model
predictions, such as decomposing code into step-by-step solutions and adding
inline NL explanations (§ 5.2). Not only is code in this style potentially
more understandable to novice data scientists (§ 7), we also demonstrate that
this improves the accuracy of self-consistency reranking Wang et al. (2022b),
because prompting the model to explain its solutions improves the diversity of
the model’s predictions.

##  2 Problem: Code Generation in Computational Notebooks

##### Problem Definition

A computational notebook is an interactive computing environment that allows
mixing code, text, and graphics. The notebook itself consists of a sequence of
markdown or source code cells. Formally, given a partial notebook context with
$k$ cells $\\{c_{i}\\}_{i=1}^{k}$ and a user-specified intent $\bm{u}$ for the
next cell $c_{k+1}$ (_e.g._ , $\bm{u}_{1}$ in Fig. 1 for $k=1$), we aim to
generate code for $c_{k+1}$ that fulfills the user’s intent Agashe et al.
(2019). This process could proceed sequentially with multiple rounds between
the user and a system Heyman et al. (2021). To answer subsequent intents
(_e.g._ , $\bm{u}_{4}$), a system will leverage the updated notebook context
(_e.g._ , $\\{c_{i}\\}_{i=1}^{5}$) which includes previous problems (_e.g._ ,
$\langle\bm{u}_{1}\sim\bm{u}_{3}\rangle$).

##### Tasks and Programming Library

We focus on synthesizing programs for data wrangling and exploratory data
analysis tasks. Specifically, data wrangling refers to the process of cleaning
and transforming data to allow for downstream analysis and machine learning,
while exploratory data analysis involves querying the data for insights to
support later statistical analysis and decision making. We focus on generating
code using pandas, a popular Python library for performing wrangling and EDA
tasks in computational notebooks, with built-in objects such as DataFrames and
Series to model tabular data types (tables and columns, respectively) .

##  3 ARCADE: A Benchmark of pandas Data Science Code Generation

###  3.1 Constructing ARCADE

ARCADE consists of $1,082$ NL-to-code problems from both existing data science
notebooks on GitHub (§ 3.1.1) and new ones curated from scratch (§ 3.1.2). In
this section we elaborate on details to create annotated examples from both
sources.

####  3.1.1 Mining Examples from Existing Notebooks

We identify candidate code cells performing data wrangling and EDA tasks from
existing high-quality data science notebooks, and then manually annotate these
cells with NL intents.

##### Collecting Notebooks for Annotation

To form a pool of candidate notebooks, we use JuICe Agashe et al. (2019), a
collection of Jupyter notebooks from GitHub, together with additional
notebooks from BigQuery, yielding over $1.5\mathrm{M}$ notebooks in
total.555https://cloud.google.com/bigquery/public-data/ These notebooks are
first filtered and near-deduplicated, similar to PaChiNCo’s training data
preprocessing step in § 4. We then identify candidate code cells from the
remaining notebooks for annotation. Specifically, we select code cells that
are either (1) contain pandas programs with at least 3 pandas API calls, or
(2) preceded by a markdown cell with a short question as its content (_e.g._ ,
What are the top 10 producers?). The first heuristic is useful to identify
complex wrangling tasks, while the second one is particularly effective in
finding interesting dataset-specific EDA tasks, and the existing markdown
texts also provide reference for labeling intents later. Next, we group the
notebooks with at least one candidate cell based on their underlying ML
datasets (_e.g._ imported using pd.read_csv()), and then select the top 5
notebooks with the most number of candidate cells from a curated set of 36
dataset groups for annotation.666This set contains ML datasets from a variety
of domains and schema. We favor notebooks with more candidate cells so that we
could extract multiple NL-to-code problems within the same notebook.

##### Annotating Cells with NL Intents

We hired a group of data scientists to annotate the notebooks selected
above.777Eight freelancers reported skill in pandas are hired from Upwork,
with an average of 3 years of experience. Annotation primarily consists of
judging the quality of candidate code cells, fixing any errors, and creating
NL intents summarizing the code. Specifically, the annotators are instructed
to frame their intents the way they prefer when interacting with an AI system
to help them implement the existing code solution, while keeping the intents
natural, concise without redundant elaboration, such as line-by-line
explanation. In addition, to make the intents more challenging, we encourage
annotators to refer to entities and variables in intents using semantic
rewrites without introducing ambiguity (_e.g._ , use “convert all binary
columns to bool” instead of listing columns verbatim), reminiscent of synonym
substitution for labeling utterances in text-to-SQL Gan et al. (2021).

##### Mitigate Ambiguity in NL Intents

Creating succinct NL intents without ambiguity could be non-trivial in this
open-domain code generation setting, especially when there could be multiple
plausible interpretations of an intent. For example, without the underlined
part of $\bm{u}_{5}$ (Fig. 1), a programmer or a system may propose
alternative solutions using different table schema. Therefore, for such open-
ended problems where there could be multiple alternative ways to present the
answer, we ask annotators to provide extra specification in their intents
about the desired output (_e.g._ schema of output DataFrame, such as the
underlined part in $\bm{u}_{5}$). Even with these additional semantic
constraints, empirically we observe around $45\%$ intents are still under-
specified, making ARCADE a challenging benchmark for handling realistic NL
intents with uncertainty. We present more analysis in § 3.2, while introducing
a robust evaluation metric that mitigates this issue in § 3.3.

Dataset | Source | Execu- | Evaluation | No. | No. | Problems | Intents Type | Intent | AST Size⋆ | No.
---|---|---|---|---|---|---|---|---|---|---
table? | Method | Notebooks | Problems | per N.B. | Length | All / pandas | API△
Existing Datasets
JuICE (Eval) | GH |  | EM$+$BLEU | 1,457 | 3,946 | 2.7 | Markdown | 60.2 | 21.2 | / 24.3‡ | 2.5
MS DSP | GH |  | Unit Tests | 305 | 1,096 | 3.6 | Markdown$+$Unit Tests | 54.3 | 28.7 | / 34.8‡ | 3.1
ExeDS | GH |  | Output Match | 277 | 534 | 1.9† | Annotated NL | 20.0 | 9.0 | / 10.7 | 2.4
NLGP | GH |  | BLEU$+$Manual Eval | 150 | 201 | 1.3 | Annotated NL | 7.7 | 13.5 | / 15.1 | 2.1
DS-1000⋄ | SO |  | Constraints$+$Unit Tests | N/A | 1,000 | N/A | Annotated NL | 166.5 | 27.3 | / 41.6‡ | 5.0
This Work: ARCADE
$\llcorner\,$Existing Tasks | GH |  | Fuzzy | 63 | 422 | 7.7 | Annotated NL | 14.1 | 19.8 | 4.2
$\llcorner\,$New Tasks | New |  | Output Match | 70 | 660 | 9.4 | 18.4 | 27.2 | 5.8

Table 1: Summary statistics of ARCADE compared with existing work. Source:
GitHub (GH) or StackOverflow (SO). △Number of pands API calls calculated on
the subset of examples using the library. ⋆Metrics are averaged over all and
(/) the subset of examples using pandas. ‡AST sizes for existing datasets are
for reference only because some code snippets also contain extra boilerplate
code (_e.g._ function/class sketches) which is not part of solution. †Most
problems in ExeDS actually do not use preceding examples in the same notebook
as contexts for evaluation. ⋄DS-1000 derived StackOverflow problems without
multi-cell notebook context, and is not directly comparable with other work.

##### Other Considerations

Indeed, re-purposing notebooks in the wild for our benchmark is not an easy
task. As an example, many notebooks in JuICe are data science tutorials, which
often contains documentation that includes background knowledge, reference
materials, and even solution hints. Those extra information makes the code
generation task easier, and may not reflect the style of ordinary notebooks
authored by data scientists in their day-to-day work. We therefore ask the
annotators to clean the notebook and remove such extra information whenever
possible. We compile a 35-page annotation guideline to cover corner cases like
this. Refer to Appendix A for more details about the guideline.

####  3.1.2 Creating Notebooks with Examples from Scratch

The problems derived from high-quality GitHub notebooks in § 3.1.1 capture
realistic tasks and notebook contexts, but may result in artificially high
evaluation accuracies due to potential leakage of evaluation notebooks to the
training data of LLMs, which is a common issue in LLM evaluation Brown et al.
(2020).888Since JuICe and BigQuery primarily contain source files before or in
2019, this issue could be more problematic. To defend against this data
contamination, we additionally annotated 660 problems by creating notebooks
from scratch.

##### Sourcing Novel ML Datasets

To ensure that those newly-created examples can be used to evaluate the
generalization ability of code LMs on unseen ML datasets, we create notebooks
targeting data wrangling and EDA tasks for $70$ tabular ML datasets that have
been recently uploaded to the Kaggle data science platform since February
2022. Those short-listed datasets are manually selected from a pool of $600$
datasets with reasonably complex schema (_e.g._ , having columns with diverse
data type), and are verified by our annotators that no older-versioned
datasets with similar schema appeared before.

##### Creating Notebooks from Scratch

For each ML dataset, the annotators were asked to create one notebook with a
series of wrangling and EDA tasks annotated with NL intents.999We only invite
the top-3 performers for this task since it is harder than labeling existing
notebooks. Specifically, we ask annotators to come up with tasks that they
would like to perform in order to gain insights into these recently appeared
ML datasets in order to build models for them. We follow the same standard to
create intents as in § 3.1.1. To make the problems more challenging,
annotators are encouraged to create harder tasks whose code solutions require
at least $5$ pandas API calls. For quality assurance, each notebook is peer
reviewed by another annotator, before a final round of review by the first
author. Since the annotators have already worked on the prior task of creating
examples from existing notebooks, they are fairly familiar with the
requirement, and are able to create each problem in $13$ minutes on average.

###  3.2 Dataset Analysis

ARCADE consists of $1,082$ NL-to-code problems from $133$ notebooks based on
$106$ unique ML datasets. We first present some analysis on our dataset,
before comparing the statistics of ARCADE with existing datasets in Tab. 1.

##### NL Intents are often Under-specified

Since ARCADE aims to evaluate code LMs in the real-world scenario where data
scientists provide succinct NL intents without extra specification (_e.g._ I/O
examples), the intents we collected are often under-specified and may not
contain sufficient information to generate a solution that executes to the
exact reference output. To understand the patterns of semantic ambiguity in
user-issued intents, we conduct a case study, in which we manually examine a
small group of random samples. Among those samples, around $55\%$ of them are
unambiguous, _i.e._ , they are precise and sufficient to infer the target
outputs. Those intents are often numerical queries with limited variety in
output type (_e.g._ , How many customers receive an income of more than 95
percentile?), or contain sufficient descriptions of output specification (§
3.1.1, _e.g._ , List the 3 largest towns, return the name and area). The
remaining 45% of intents are under-specified. Of these, $30\%$ lack a
description of target columns in output DataFrames (_e.g._ , $c_{4}$, Fig. 1).
An additional $25\%$ of intents imply that the desired DataFrames should have
a complex schema, such as a nested row index or table header (_e.g._ , Show
the time of the day and the fare price for each airline), which is difficult
to predict without additional information. Another $25\%$ of ambiguous intents
have lists of entities as their answers, and do not specified the desired
container type (_e.g._ List or pandas.Series). Finally, the remaining $20\%$
cases require outputs with more complex structures (_e.g._ , multiple
variables), or imply additional post-processing steps such as data imputation.

##### Notebook Context help Disambiguate Intents

Interestingly, while nearly half of the intents are under-specified, we
observe that $\sim 30\%$ of those cases can be disambiguated by referring to
intents or code snippets from prior rounds of problems with similar query
patterns in the notebook context. A common scenario is follow-up EDA queries
(_e.g._ How many of them are $\ldots$) that refine the previous query step
(_e.g._ Show me all $\ldots$) by narrowing the down search condition, while
the output DataFrame structure is the same as previous turns, reminiscent of
similar thematic relation patterns in contextual semantic parsing Iyyer et al.
(2017); Yu et al. (2019b).

##### Comparing Existing and New Tasks

Looking closer into the problems in the _Existing Tasks_ (§ 3.1.1) and _New
Tasks_ (§ 3.1.2) splits, the _New Tasks_ split has more challenging problems
than _Existing Tasks_ , as measured by the number of pandas API invocations
and the AST size in reference code solutions (Tab. 1, _Bottom_). Fig. 2 plots
the histogram of pandas API usage, which shows $67\%$ problems in _New Tasks_
require at least 5 API calls to solve them. As we show in § 5, with more
complex held-out problems targeting for recent ML datasets, the _New Tasks_
split serves both as a more robust and challenging benchmark for state-of-the-
art code LLMs.

##### Comparing with Existing Datasets

Tab. 1 compares ARCADE with existing datasets for data science code generation
in computational notebooks. Specifically, JuICe Agashe et al. (2019) contains
exercise problems in assignment notebooks from data science tutorials or
coures, where the NL intents are usually elaborative assignment problem
definitions. Notebooks in JuICe are not executable so evaluation is performed
by surface-level matching (exact match or Bleu) between reference and
predicted programs. DSP Chandel et al. (2022) contains problems from a
filtered set of JuICe notebooks that are executable and also associated with
unit tests for auto-grading. Hence the intents in DSP follow similar patterns
as those in JuICe. To ensure that the free-form model-predicted code is
compatible with unit tests, DSP uses the unit test code iteself as extra model
input besides NL intents to constrain the model to generate code that could be
directly consumed by the tests. Next, ExeDS Huang et al. (2022), a concurrent
work to this paper, is another set of filtered problems from JuICe. Similar to
this work, ExeDS uses hand-annotated intents, and compares the execution
output between reference and predicted code for evaluation instead of relying
on unit tests (§ 3.3). Moreover, NLGP Heyman et al. (2021) is another
collection of the NL-to-code problems in Jupyter notebooks with short
annotated intents for simple data manipulation tasks, where most notebooks
have one associated problem. Finally, for the sake of completeness, we also
list another concurrent work DS-1000 Lai et al. (2022), a collection of data
science problems derived from StackOverflow questions. It primarily features
problems using synthetic contexts with minimal working examples, and therefore
does _not_ concerns with code generation in notebooks with interrelated
problems on general ML datasets.

We remark that ARCADE is the only benchmark that satisfies _all_ the following
criteria: _First_ , our dataset consists of manually annotated, concise NL
intents (“Intents Type” column in Tab. 1). Those utterances are significantly
shorter (“Intent Length” column) than verbose markdown texts often found in
tutorial notebooks (c.f. JuICe, DSP), and do not require a model to rely on
extra specification such as unit tests (c.f. DSP). Those succinct intents are
also under-specified, requiring a more robust evaluation metric to consider
alternative answers (later in § 3.3). This may also motivate future research
on generating more diverse predictions to cover possible interpretations of a
problem — which we will explore in § 5 — or even models to explicitly capture
such uncertainty in NL intents Lin et al. (2022). _Second_ , ARCADE contains
more interrelated problems in a single notebook (“Problems per N.B.”) with
complex dependencies (Fig. 1), capturing the essence of interactive computing.
This makes our dataset useful in testing an LLM’s skill at understanding rich
contexts, including existing user-written code and markdown cells, as well as
preceding problems and their solutions (§ 2). _Third_ , ARCADE challenges LLMs
with grounded language understanding, where the model needs to link relevant
concepts (_e.g._ “max and min” in $\bm{u}_{1}$) in an intent to the
corresponding variable execution states in the context (_e.g._ content of the
TIME column in df). Such necessity of understanding semi-structured tabular
data contents Pasupat and Liang (2015) and performing data transformations
using an open domain programming language (Python) potentially makes NL
grounding in ARCADE more challenging than other applications like database
semantic parsing Yu et al. (2019b). _Fourth_ , ARCADE has problems with
complex data transformations and richer usage of real-world data science APIs.
As evidence for this, note that the number of pandas APIs used in each problem
(“No. API” in Tab. 1)101010Calculated by counting function names in a
predefined list of pandas functions, including advanced array indexing
operations. Functions with similar names from other libraries (_e.g._ numpy)
may also be included. is on par with DS-1000 and significantly higher than
ExeDS, DSP and NLGP. _Finally,_ nearly $70\%$ of problems and their notebooks
(_New Tasks_) are curated from scratch targeting for recent ML datasets, which
mitigates the risk of training data contamination of LLMs, and also makes
ARCADE a more reliable benchmark to test the generalization ability of LMs in
understanding held-out tabular datasets Lee et al. (2021).

To summarize, ARCADE features _more realistic NL intents_ and _multiple
related problems in the same notebook context_ to better reflect the real-
world scenario where data scientists prompt LLMs using ephemeral NL comments
for assistance with complex data wrangling or EDA tasks Barke et al. (2022),
therefore making it ideal to evaluate LLMs in this realistic setting.
Nevertheless, we remark that existing datasets in Tab. 1 usually contain
broader types of problems other than the wrangling and EDA tasks (_e.g._
fitting ML models) considered in this paper. We leave expanding the task
spectrum as important future work.

Figure 2: Histogram of the number of pandas API calls in ARCADE.

###  3.3 Evaluation by Fuzzy Output Matching

Recent code generation datasets have shifted from surface form evaluation
metrics like Bleu Yin and Neubig (2017) to functional correctness, where the
execution results of predicted programs on some input examples are matched
with the reference output verbatim Hendrycks et al. (2021); Chen et al.
(2021a); Austin et al. (2021), or approximately (for complex output like
plots, see Chandel et al., 2022) using unit tests. Such evaluation methods
assume a model generates code that strictly follows the signature of unit
tests, which is impractical in data science notebooks with free-form code.

In this paper we aim to evaluate accuracy of code LLMs in synthesizing code
for data science notebooks in the wild, _i.e._ , the model is only conditioned
on preceding notebook context besides a short user-issued intent (§ 2). As
discussed in § 3.2, those intents are often under-specified and have multiple
alternative solutions. We therefore use a set of heuristics to approximately
match the execution output111111For code that in-place modifies a variable
(_e.g._ df in $c_{2}$), we treat the modified variable as the output. of a
predicted program with the annotated reference to determine if they are
functionally equivalent. These heuristics primarily cover two scenarios.
First, if the output variable is a container type,121212List, Tuple, Set,
numpy.ndarray, pandas.Series and single-column pandas.DataFrame. we
canonicalize variables to the same type. Second, we allow for partial matching
between complex DataFrame variables. Formally, consider a reference DataFrame
$\bm{v}$ with a set of columns $\\{v_{i}\\}$, where each $v_{i}$ is a vector
of cell values for the $i$-th column. We define that $\bm{v}$ is equivalent
with the output DataFrame of a predicted program $\hat{\bm{v}}$, iff for any
$v_{i}\in\bm{v}$, we have $v_{i}\in\hat{\bm{v}}$. Intuitively, we consider a
predicted program as correct if its output DataFrame contains all the columns
(and cell entries) in the reference frame. Intuitively, a predicted program is
already useful as long as it covers all the necessary information in the
reference. In cases of DataFrames, a user could easily create a more compact
view of the frame by selecting a subset of target columns.

Empirically, we find our evaluation metric is reliable in identifying
solutions with alternative output structures, with a relatively low false-
negative rate (under $10\%$). We present an error analysis in § 6.

##  4 PaChiNCo: Adapting Code LMs to Computational Notebooks

In this section we elaborate on PaChiNCo, our code LM for Python notebooks.

##### Base LM

PaChiNCo is based on PaLM, a family of Transformer-based LMs developed for
few-shot learning for general natural language tasks Chowdhery et al. (2022).
Specifically, we use the 62B PaLM model trained on $1.3$T tokens with a
mixture of conversational, webpages and code data (Section F, Chowdhery et al.
(2022)). Starting with 62B PaLM as the base model, we perform two stages of
fine-tuning, first on Python source code data, followed by another fine-tuning
run on a collection of Jupyter notebooks.

##### Fine-tuning on Python Code

We first fine-tune PaLM 62B on a large corpus of Python source code with 64B
tokens. This corpus consists of deduplicated Python source files with
permissive license collected from GitHub. We finetune PaLM for 1 epoch with a
batch size of 256 while keeping the other hyper parameters consistent with
Chowdhery et al. (2022). As we show in § 5 and also in Appendix C, this
Python-finetuned PaLM 62B model is already a strong code LM, registering
significant improvements on ARCADE over the base model, while even
outperforming the larger code LM PaLM-Coder 540B on existing benchmarks.

##### Continued Fine-tuning on Notebooks

After fine-tuning on Python code, we perform a second stage of fine-tuning on
a large collection of Jupyter notebooks. Our Jupyter notebook dataset is
created similarly as the Python code corpus from GitHub, with additional
domain-specific pre-processing steps, such as filtering out notebooks with
fewer than 4 cells. In addition, to mitigate the risk of having notebooks
similar to the evaluation notebooks from GitHub in the _Existing Tasks_ split
leaked into the training data, we perform near de-duplication against
notebooks in _Existing Tasks_ at the cell level. Specifically, we cluster the
cells of notebooks in both the evaluation and training sets based on a fuzzy-
matching similarity metric, and remove any training notebooks that has one
cell that falls into the same cluster as a cell from one of the evaluation
notebooks. This process eliminates ${\sim}350\mathrm{K}$ notebooks from the
fine-tuning data. Our final training set consist of ${\sim}3.8\mathrm{M}$
notebooks and ${\sim}9.6\mathrm{B}$ tokens in total. We use nbconverter to
linearize Jupyter notebooks to Python code, where cells in a notebook are
concatenated using the special delimiter ‘# In[]:’, with markdown texts
converted to comments. Refer to Appendix B for details about fine-tuning, and
Appendix M for a data card.

##  5 Experiments

###  5.1 Models and Setup

To demonstrate the challenges of ARCADE, we compare the performance of
PaChiNCo with existing large code language models.

CodeGen Nijkamp et al. (2022) is a family of code LMs ($350\mathrm{M}$,
$2.7\mathrm{B}$, $6.1\mathrm{B}$ and $16.1\mathrm{B}$ parameters) trained
using a combination of NL and code data scraped from GitHub (ThePile, Gao et
al., 2020). We evaluated on both the multilingual version trained on BigQuery
and the monolingual version trained on $71\mathrm{B}$ Python tokens.

InCoder Fried et al. (2022) is a code LM ($1.3\mathrm{B}$ and $6.7\mathrm{B}$
parameters) for both left-to-right generation and infilling tasks. The LM is
trained on a mixture of StackOverflow (SO) posts and $159\mathrm{GB}$ of
multilingual code data. In particular, the corpus contains a dedicated portion
of Jupyter notebooks ($\sim 5\mathrm{GB}$), making it appealing to evaluate on
ARCADE.

PaLM Besides PaChiNCo, we also evaluate the base PaLM $62\mathrm{B}$ LM, as
well as the Python code LM obtained from the first-stage of code fine-tuning
(§ 4).

xx

##### Inference

For each problem, we convert the problem into a prompt (§ 5.2) and draw random
samples from the LMs using nucleus sampling with a top probability of $0.95$
and a temperature $0.8$. Refer to Appendix D for details.

##### Metrics

Following prior work Chen et al. (2021a); Austin et al. (2021), we measure
model performance using the pass@$k$ metric, defined as the fraction of
problems with at least one correct sample given a sample size $k$. To reduce
variance, we estimate pass@$k$ ($k\leq 30$) by drawing $50$ samples for each
problem Chen et al. (2021a).

Figure 3: An example problem. Cells 1-3 encode the notebook context. Cell 4
contains the NL intent as problem statement. Figure 4: Example completions of
Cell 4 from Fig. 3 which illustrate the desired coding style for different
few-shot prompting strategies.

###  5.2 LM Prompting Strategies

For a given problem, we conducted two types prompting experiments: prompting
using only the notebook context of the problem (§ 5.3), and few-shot prompting
with extra exemplars as prompt prefix before notebook context (§ 5.4), in
order to impose more control to the style of predicted code.

####  5.2.1 Prompting using Notebook Context

In our default setting (§ 5.3), for a given problem we turn its notebook
context to a prompt to query LLMs. Fig. 3 gives an example. The context
consists of preceding cells before the problem, which may include previous
rounds of interactions (§ 2), and is followed by the current NL intent. To
help LLMs understand intents using variable state information (§ 3.2), we also
add NL descriptions of the schema of the dataset imported to the notebook
(schema description, _e.g._ cell 3, Fig. 3), which consist of columns and
example values in the DataFrame. Representing structured schema information in
NL is a common strategy when prompting LLMs to solve tasks requiring
understanding structured data Xie et al. (2022). Next, for the following
problems after cell 4, we include reference solutions to previous problems in
their prompts, following the multi-turn task-oriented dialogue setting in
Andreas et al. (2020). See Appendix N for a complete example.

####  5.2.2 Few-shot Prompting

Besides the basic setting using the notebook context of a problem, we also
explored prompting using additional NL-to-code exemplars as prompt prefix
before the notebook context. The motivation is to nudge the LM to generate
code solutions following different coding and commenting styles, such as
documenting the “reasoning path” used by the code. We develop four prompting
strategies:

  1. 1.

Vanilla code strategy as in Fig. 4, completion 4a. The predicted code follows
the common practice of chaining multiple API calls in a single line.

  2. 2.

Step-by-step (SbS) code strategy as in Fig. 4, completion 4b, which leads to
code with fine-grained decomposition structure.

  3. 3.

SbS code with preamble, as in Fig. 4, completion 4c, which could further
elicit decomposition in predictions.

  4. 4.

SbS code with preamble and explanation, as in Fig. 4, completion 4d, with
inline NL explanations per step.

pass@$k$ | Existing Tasks | New Tasks
---|---|---
5 | 10 | 20 | 30 | 5 | 10 | 20 | 30
InCoder $1\mathrm{B}$ | $30.1$ | $36.4$ | $42.6$ | $46.0$ | $3.8$ | $5.8$ | $8.4$ | $10.0$
InCoder $6\mathrm{B}$ | $41.3$ | $48.0$ | $54.0$ | $57.2$ | $7.0$ | $10.0$ | $13.5$ | $15.8$
CodeGenmulti $350\mathrm{M}$ | $13.3$ | $16.0$ | $18.8$ | $20.4$ | $1.0$ | $1.6$ | $2.5$ | $3.1$
CodeGenmulti $2\mathrm{B}$ | $25.0$ | $29.9$ | $34.8$ | $37.7$ | $2.7$ | $4.0$ | $5.5$ | $6.5$
CodeGenmulti $6\mathrm{B}$ | $28.0$ | $34.0$ | $39.9$ | $43.3$ | $3.0$ | $4.6$ | $6.7$ | $8.0$
CodeGenmulti $16\mathrm{B}$ | $31.2$ | $37.6$ | $43.8$ | $47.2$ | $4.6$ | $6.8$ | $9.5$ | $11.3$
CodeGenmono $350\mathrm{M}$ | $18.9$ | $24.3$ | $30.1$ | $33.4$ | $1.9$ | $2.8$ | $4.0$ | $4.9$
CodeGenmono $2\mathrm{B}$ | $35.8$ | $43.3$ | $50.7$ | $54.9$ | $6.5$ | $9.8$ | $13.8$ | $16.3$
CodeGenmono $6\mathrm{B}$ | $42.1$ | $49.4$ | $56.3$ | $59.7$ | $8.9$ | $13.1$ | $18.3$ | $21.6$
CodeGenmono $16\mathrm{B}$ | $46.7$ | $54.5$ | $60.9$ | $64.0$ | $12.0$ | $17.1$ | $22.7$ | $26.3$
PaLM $62\mathrm{B}$ ($1.3\mathrm{T}$ Tokens) | $49.7$ | $57.5$ | $64.4$ | $67.8$ | $12.5$ | $17.4$ | $22.8$ | $26.0$
$+$ Python Code |  $58.8$ $\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color<EMAIL_ADDRESS>|  $65.7$ $\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color<EMAIL_ADDRESS>|  $71.7$ $\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color<EMAIL_ADDRESS>|  $74.7$ $\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color<EMAIL_ADDRESS>|  $21.4$ $\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color<EMAIL_ADDRESS>|  $28.4$ $\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color<EMAIL_ADDRESS>|  $35.6$ $\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color<EMAIL_ADDRESS>|  $39.8$ $\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color<EMAIL_ADDRESS>
$+$ Notebooks (PaChiNCo) |  $64.6$ $\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color<EMAIL_ADDRESS>|  $71.0$ $\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color<EMAIL_ADDRESS>|  $76.0$ $\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color<EMAIL_ADDRESS>|  $78.3$ $\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color<EMAIL_ADDRESS>|  $30.6$ $\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color<EMAIL_ADDRESS>|  $38.0$ $\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color<EMAIL_ADDRESS>|  $45.0$ $\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color<EMAIL_ADDRESS>|  $48.6$ $\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color<EMAIL_ADDRESS>
$-$ Schema Description |  $60.5$ $\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color<EMAIL_ADDRESS>|  $67.1$ $\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color<EMAIL_ADDRESS>|  $72.8$ $\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color<EMAIL_ADDRESS>|  $75.5$ $\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color<EMAIL_ADDRESS>|  $22.7$ $\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color<EMAIL_ADDRESS>|  $28.5$ $\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color<EMAIL_ADDRESS>|  $34.1$ $\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color<EMAIL_ADDRESS>|  $37.2$ $\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color<EMAIL_ADDRESS>

Table 2: pass@$k$ evaluation on ARCADE using notebook context as prompts. Grey
figures indicate difference to previous line.

Each strategy has four exemplars in its prompt prefix, as listed in Appendix
N. Those extra exemplars before notebook context could help the model better
understand the task Brown et al. (2020). More importantly, we could also use
additional exemplars to impose more control on the inference process of LMs,
with each strategy imposing a slightly different level of granularity of the
expected completion. Specifically, we focus on prompting LMs to generate code
that follows a multi-line, step-by-step decomposition structure
($\mathrm{SbS}$), with NL explanations of the underlying data transformation
for each step. Prompting LMs to explain its own solution in such a step-by-
step fashion has been found effective in improving accuracy of natural
language reasoning Wei et al. (2022); Gao et al. (2022) and program induction
Nye et al. (2021) tasks. In § 5.4, we show it effectively improves diversity
of predicted code solutions, which is beneficial for self-consistency
reranking.

###  5.3 Results

In our first set of experiments, we evaluate the performance of state-of-the-
art code LMs on ARCADE. Results are listed in Tab. 2. As in § 5.2, inputs to
those models are only the notebook context (no few-shot prompting). We
truncate the prompt up to $900$ tokens using the tokenizer of PaLM.

##### Comparing PaLM Models

Our PaChiNCo model registers the best performance on both the _Existing Tasks_
split based on high-quality GitHub notebooks as well as on the _New Tasks_
split with held-out ones, thanks to its larger size and domain-specific fine-
tuning on notebooks.131313See Appendix E for pass@$1$ results. It is worth
noting that the base NL-focused PaLM 62B model already outperforms most public
code LMs and is on par with CodeGenmono $16\mathrm{B}$. As the LM is fine-
tuned on Python ($+$ Python Code) and notebooks ($+$ Notebooks) data, the
domain gap is closing, and pass@$k$ continues improving. Indeed, the Python
fine-tuned PaLM is already a strong model for Python code generation,
achieving state-of-the-art results on benchmarks such as HumanEval and MBPP
(Appendix C). Meanwhile, the absolute gain after fine-tuning on Python code is
higher than continued training on notebooks, which is not surprising since the
semantic gap between NL data and Python code is larger than that between code
in Python and notebooks.

##### Comparing Public Code LMs

Next, we discuss results on public code LMs.141414We also evaluate Codex.
Results are in Appendix G. First, among models at similar size and trained
with roughly similar amount of Python code data, namely InCoder 6B and
CodeGen${}_{\textrm{multi}}$ 6B, InCoder 6B performs better. This is likely
due to that InCoder contains dedicated portion of Jupyter notebooks (8GB) in
its training data, which makes the model more suitable for ARCADE. With
$4\times$ more Python data, CodeGen${}_{\textrm{mono}}$ 6B takes over, which
is in line with the the performance pattern of these two model families on
existing code generation datasets (Fried et al., 2022). Moreover, comparing
CodeGen models at different size, we observe that pass@$k$ scales linearly
with model size. As we discuss more in Appendix F, the scaling curve for
ARCADE (_New Tasks_) is also more flatten than that for other code generation
datasets, which might suggest that it is more challenging for CodeGen.

##### Comparing _Existing Tasks_ and _New Tasks_

It is interesting to compare results between the _Existing Tasks_ and _New
Tasks_ splits, where we observe two patterns. First, results on _Existing
Tasks_ are significantly higher across all models. Second, comparing the
improvements after Python and notebook-specific fine-tuning of PaLM 62B, the
gain on _New Tasks_ is higher. These results are likely due to two factors.
First, an obvious reason is that problems in _Existing Tasks_ are overall
simpler than _New Tasks_ (§ 3.2). Second, it is likely that some code data
similar to our evaluation notebooks in _Existing Tasks_ split is leaked into
the training data of those LMs. Specifically for PaLM, its training mixture
contains general Web documents and sampled Python source files. Despite the
fact that we went through significant amount of effort to deduplicate the
training data against ARCADE during the second-stage notebook fine-tuning (§
4), given that some of the notebooks used to create _Existing Tasks_ are
popular data science tutorials on GitHub, it is still possible PaLM has
already seen similar code data from the Web or notebook-converted Python
source files. This suggests the importance of evaluating code LMs on held-out
data, which is the purpose of _New Tasks_. We leave investigating training
data contamination as important future work.

##### Impact of Schema Descriptions

Our prompts contain NL descriptions of schema for DataFrames imported to a
notebook (§ 5.2.2). As in § 3.2, variable state information like DataFrame
schema is crucial for understanding NL intents. The last row in Tab. 2
($-\bm{\mathrm{Schema~{}Description}}$) reports an ablation where we remove
such NL schema descriptions from prompts, which leads up to significantly
worse results, especially on the _New Tasks_ split. Encoding the states for
more intermediate variables could potentially further improve performance,
which we leave as interesting future work.

Figure 5: pass@$k$ of PaChiNCo w.r.t problem complexity

##### Accuracy with Problem Complexity

To better understand PaChiNCo’s performance on problems at different levels of
complexity, we plot pass@$30$ with respect to the number of pandas function
calls in the annotated reference solutions, as shown in Fig. 5. For problems
with similar complexity, PaChiNCo generally achieves higher pass rate on
_Existing Tasks_ , again suggesting that the _New Tasks_ split is still more
challenging even after controlling problem complexity. Refer to Appendix H for
a similar plot w.r.t AST size.

##### How Much Notebook Context is Useful?

ARCADE requires a model to leverage rich programmatic and NL context in test
notebooks to generate code solutions for the current cell. To study PaChiNCo’s
performance with varying amount of available notebook context, we control the
number $d$ of context cells $\\{c_{i}\\}_{i=k-d}^{k-1}$ (§ 2) when generating
code for each problem (at cell $c_{k}$) in our dataset. Fig. 6 depicts
pass@$30$ as a function of the context size $d$. Since we use the first
preceding cell $c_{k-1}$ to store the NL intent $\bm{u}_{k}$ for $c_{k}$
(Appendix N), having only one context cell is equivalent to the “cold-start”
setting of only using $\bm{u}_{k}$ (besides schema description) to predict
$c_{k}$. PaChiNCo achieves a pass rate of $44\%$ (existing tasks) and $17\%$
(new tasks) in this challenging setting ($d=1$), with errors mostly due to
failure in referring to variables that the solution relies on, whose
information is not present in the short context. Indeed, including additional
context cells is crucial for good performance. In particular, having 3 context
cells could already lift the pass@$30$ to $72\%$ and $36\%$ on the two splits
— $1.6\sim 2\times$ higher than $d=1$. The results also start to plateau after
including $5\sim 7$ context cells, with diminishing returns after including
more cells, which is in line with findings in Agashe et al. (2019).
151515Prior work suggests that the plateau point is around three neighboring
cells, while in our case, the number if approximately doubled since we need
extra cells to include intents in previous turns (Appendix N). Empirically, we
observe that using more context helps to reduce schema understanding errors
(_e.g._ using undefined columns in DataFrames). Refer to Appendix I for more
details.

Figure 6: pass@$30$ as w.r.t the number of context cells from the notebook.
Figure 7: pass@$30$ for few-shot prompting on _New Tasks_. Results averaged
over three runs with different prompt prefixes. Figure 8: Few-shot prompting
code style metrics Figure 9: Pass rate on examples grouped by context size. We
also report the number of examples in each group.

##### Does Problem Location Impact Performance?

Another interesting angle to study the effect of context is through the lens
of model accuracy when solving problems $c_{k}$ at different locations.
Intuitively, problems located later in a notebook ($k$ is larger) would have
more context available, therefore they could be easier to answer Wang and Cho
(2016). Fig. 9 shows pass@$30$ on problems grouped by their preceding context
size, which shows increased task success rate when solving problems with more
context, confirming the prior intuition.161616We remark that our setup is
different from multi-turn semantic parsing where later turns are conditioned
on the _predictions_ of prior turns, while we use reference solutions (§ 5.2).
See § 8 for discussion.

###  5.4 Few-shot Prompting Beyond Test Notebook Contexts

Our previous set of experiments demonstrate diminishing returns after
including more notebook context cells. As motivated in § 5.2, besides test
notebook context, we use few-shot prompting with additional exemplars to teach
the model to perform this challenging task, while generating code in different
styles. Here, we report both functional correctness (Fig. 7) as well as
metrics evaluating the style of model-predicted code (Fig. 8) for problems in
_New Tasks_. We only evaluate PaChiNCo due to that the prompt length (maximal
$2,100$ sub-tokens) exceeds the limit of public code LMs.

With Step-by-Step Prompting ($\bm{\mathrm{SbS}}$), we observe an improvement
on pass@$30$ over the baseline using only notebook contexts (Tab. 2). This is
consistent with the general results on other tasks when comparing few- vs
zero-shot learning with LLMs Brown et al. (2020), while we remark that in our
case, the “zero”-shot prompting setting (Tab. 2) still contains notebook cells
as rich contexts. Nevertheless, as Fig. 7 suggests, $\mathrm{SbS}$ prompting
is quite helpful for improving pass rate. Empirically, we observe
$\mathrm{SbS}$ prompting is especially helpful for problems (at code cells
$c_{k}$) without adequate preceding notebook context ($k$ is small), For
instance, $\mathrm{SbS}$ prompting results in 6$\%$ absolute improvement for
the first two rounds of problems. Interestingly, even if we include more test
notebook context in the non-exemplar baseline such that the prompt length
match with $\mathrm{SbS}$ prompting using extra exemplars
($\bm{\mathrm{Baseline}+\mathrm{More~{}Context}}$), $\mathrm{SbS}$ prompting
is still better, which again suggests the value of extra exemplars to
compliment the information in notebook context.

In line with our intuition, $\mathrm{SbS}$ prompting yields code predictions
that are better decomposed into more lines of code (LoC↑), with each line
being simpler (Tokens/API per Line↓). It also appears to generate slightly
more complex solutions, as indicated by the increased pandas API usage. Our
finding demonstrates the effectiveness of $\mathrm{SbS}$ prompting in
improving accuracy and also code style. In contrast, we also report results
from prompting the model using exemplars with “vanilla”-styled code following
the common practice of chaining multiple pandas API calls in a single line
($\bm{\mathrm{Vanilla~{}Code}}$), which only leads to marginal improvements on
pass@$k$ over the non-exemplar baseline, while the code style remains
consistent.

(a)

(b)

Figure 10: Cumulative distributions of the number of unique API sequences
(top) and output clusters (bottom) extracted from PaChiNCo’s $50$ predictions
on the New Tasks split. Curves that appear more to the right represent
prompting methods with greater diversity in the samples. Step-by-step
prompting leads to much greater diversity than the baselines.

Next, on top of $\mathrm{SbS}$ prompting, using a preamble at the beginning of
a target cell to further encourage the model to solve the problem in a step-
by-step fashion ($\bm{\mathrm{SbS}+\mathrm{Preamble}}$) improves pass@$k$, as
well as the level of decomposition in predicted code (LoC↑, Tokens/API per
Line↓). This finding is inline with Kojima et al. (2022) where using preambles
to nudge the model to generate reasoning steps can still be useful together
with few-shot step-by-step demonstrations. More surprisingly, with additional
inline NL explanations for each step
($\bm{+\mathrm{Preamble}+\mathrm{Explanation}}$), PaChiNCo produces even more
decomposed code while maintaining similar accuracy. Those predictions also
come with rich NL comments as explanations, with the number of comment lines
nearly equal to the number of code lines. Intuitively, step-wise NL
explanations could be potentially helpful for developers to better understand
predictions from code LMs, as we demonstrate later in § 7.

We also report breakdown results of pass rate on problems with varying level
of complexity. $\mathrm{SbS}$ prompting and its variants are helpful across
the board, especially for harder tasks with more than 7 pandas function calls.
This might suggest the value of step-by-step decomposition when synthesizing
complex programs. Next, we provide more insight into the effectiveness of
$\mathrm{SbS}$ prompting.

Figure 11: Self-consistency decoding accuracy (using 1 reranked sample) on New
Tasks

##### Step-by-Step Prompting Produces More Diverse Solutions

Given the effectiveness of $\mathrm{SbS}$ prompting in improving accuracy,
code style, and explainability (§ 7), next we explore why it is helpful by
analyzing patterns in model-predicted solutions. One interesting question is
whether $\mathrm{SbS}$ prompting could produce more diverse solution
approaches to solve a problem. Intuitively, greater diversity in the output
space could improve the odds of finding a solution at higher sample size $k$.
Determining whether two solutions are “different” is difficult and subjective,
but we approximate this from two angles. First, we use the sequence of pandas
API calls as a signature of the high-level solution pattern. Second, since
different solutions might have the same functionality (_i.e._ , they execute
to the same output), we also measure the diversity in _functionality_ by
clustering predictions based on their outputs.171717Here, each cluster
corresponds to programs that executes to the same exact results (_i.e._ no
fuzzy matching as in § 3.3). Fig. 10 plots the cumulative distributions of the
number of unique solution patterns and output clusters on the _New Tasks_
split. Step-by-step prompting results in greater diversity on both metrics
compared to the baselines. Interestingly, prompting with NL explanations
yields even more solution patterns. Having diverse predictions could be
potentially useful to handle under-specified intents (§ 3.2), as those
predictions could correspond to different interpretations of the ambiguous
intents. Empirically, we find diverse predictions are particularly effective
for post-hoc reranking such as self-consistency decoding, where we show the
user one prediction from the largest output cluster instead of showing all $k$
predictions. As shown in Fig. 11, $\mathrm{SbS}$ prompting significantly
improves over baselines when using self-consistency decoding Wang et al.
(2022b). Notably, the $1$-sample accuracy of $\mathrm{SbS}$ with NL
explanations is on par with PaChiNCo’s pass@$5$ accuracy (last row, Tab. 2).
We present further analysis of solution diversity from PaChiNCo in Appendix K.

##  6 Error Analysis

Figure 12: Frequency of execution errors on _New Tasks_.

To understand the types of errors that LLMs make on ARCADE, especially on
challenging problems, we conduct an error analysis on PaLM’s predictions on
the _New Tasks_ split (Tab. 2). Overall, we notice a significant drop in
execution errors after two-stage code fine-tuning (PaLM$\mapsto$PaChiNCo, §
4). Out of all the incorrect predictions from PaChiNCo under the fuzzy output
matching evaluation metric (§ 3.3), roughly $35\%$ of the samples result in
execution errors, while the remaining $65\%$ predictions have executable
programs but are functionally incorrect. First, for un-executable samples, we
present an analysis of the distribution of different execution error types, as
illustrated in Fig. 12. The primary source of execution error is KeyError and
AttributeError due to reference to non-existing index or columns in
DataFrames. While in the prompts we provide NL descriptions of schema for
DataFrames loaded to notebooks (§ 5.2), such descriptions for intermediate
DataFrames that are later derived in the context are still missing due to
limited prompt length, and the model may not be able infer their schema
information solely from the source code. This could be especially problematic
for APIs that create compound intermediate DataFrames with complex schema,
such as pd.groupby, which accounts for that more than $50\%$ of those
KeyErrors and AttributeErrors. Similarly, other execution errors such as
ValueError and TypeError are often caused by the insufficient knowledge about
the DataFrame contents. For example, ValueError occurs when a model tries to
calculate the mean of a column which has NaN values. This finding suggests the
importance of developing LLMs that could handle longer context Transformers et
al. (2022) in order to include more DataFrame information in prompts. We gave
a detailed case study on these types of execution errors in Appendix L.

Next, we conduct a manual analysis on $50$ randomly sampled incorrect
predictions (with a focus on executable but incorrect predictions). The cause
of these errors can be grouped into the following categories: 1\. Complex
problems requiring non-trivial reasoning or data transformation steps
($43\%$); 2. Errors in interpreting NL intents, such as missing a requirement
specified in the intent (_e.g._ round to two decimal places) in the code
solution ($26\%$); 3\. Errors caused by under-specified intents (§ 3.2,
$19\%$); 4\. False-negatives due to limited coverage of the fuzzy-matching
evaluation metric (§ 3.3, $6\%$); 5\. Annotation errors ($6\%$).

The primary source of errors is due to complex problems, which reiterates the
motivation of ARCADE — evaluating code LLMs on challenging data wrangling and
EDA tasks. The second majority type of errors (misunderstanding intents)
suggests room to improve PaChiNCo’s skill in instruction following. Next, a
non-trivial amount of errors are caused by under-specified intents, which are
common in the setting of prompting LLMs using short instructions (§ 3.2),
calling for future research to specifically address this issue. Finally, our
evaluation metric based on fuzzy output matching seems effective in
identifying plausible alternative solutions. Still, there are non-trivial
cases where there are multiple ways of presenting the outputs (_e.g._
DataFrames with nested columns or different orientations, Fig. 32). We present
more detailed examples of these errors in Appendix I.

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

In § 5 we show how prompting PaChiNCo to generate code with step-wise
explanations is helpful in improving solution diversity (Fig. 10) and also
accuracy of self-consistency reranking (Fig. 11). In this section we conduct a
qualitative analysis on model-generated code snippets with explanations to
explore how they could be useful for novice data scientists.

First, we observe that NL explanations could help users follow the flow of
complex data transformations for programs involving a chain of pandas
operations. By decomposing and explaining how data is manipulated after
individual transformation steps, it is easier for users to understand the
solution and track its dataflow behind the scene, especially when some steps
involve complex computation (Fig. 19), or the underlying schema is less
intelligible (_e.g._ , column names with abbreviations, Fig. 20).
Additionally, some inline explanations also describe the output of
intermediate steps, which is particularly helpful when these steps involve
advanced pandas functions whose output may not be obvious, such as pd.unstack
(Fig. 21)

Meanwhile, step-wise NL explanations serve as high-level procedural
descriptions of code, which enables users to easily browse through and
understand different solution approaches without being distracted by nuances
in the actual code implementation (Fig. 22). Moreover, explanations also help
users verify the code solutions by identifying potentially incorrect steps and
make necessary corrections (Fig. 23). The observations presented here offer
insight into potential future avenues to improve the utility of code LMs for
developers through the use of step-by-step explanations, which we leave as
important future work.

##  8 Related Works

##### LLMs for Code

In recent years there has been a burgeoning of LMs trained on code Chen et al.
(2021a); Austin et al. (2021); Nijkamp et al. (2022); Fried et al. (2022); Li
et al. (2022); Tunstall et al. (2022); Xu et al. (2022). Many generalist LMs
also have source code as part of their training data mixture Thoppilan et al.
(2022); Chowdhery et al. (2022); Wang and Komatsuzaki (2021); Black et al.
(2021). Those LLMs have registered impressive performance on a variety of
benchmarks for code, such as code translation Lachaux et al. (2020a), program
repair Gupta et al. (2017), and natural language to code generation Chen et
al. (2021a); Hendrycks et al. (2021); Austin et al. (2021); Li et al. (2022).

##### Automating Data Science

Given the sheer amount of tasks involved in the lifecycle of data science and
the expertise required, it calls for development of systems to automate this
lifecycle Aggarwal et al. (2019); Wang et al. (2021a, b). As an example,
automating the process of feature engineering and finding the best-performing
models has been the central theme of AutoML research He et al. (2021);
Karmaker et al. (2020), with well-established systems Feurer et al. (2015) and
evaluation protocols Zöller and Huber (2021). In this paper we focus on
automating data wrangling and EDA tasks, which account for nearly the same
amount of code and documentations in notebooks as that for feature engineering
and building ML models Agashe et al. (2019); Wang et al. (2022a). Along this
line of research, AutoPandas Bavishi et al. (2019) synthesizes pandas programs
for data transformation given examples of input/output DataFrames. Other
approaches aim to generate pandas programs using LLMs using both NL intents
and I/O samples Jain et al. (2021) or unit tests Chandel et al. (2022). In
this paper we consider code generation in notebook environments with multiple
rounds of problems and interdependent contexts (refer to § 3.2 for detailed
discussions of recent work). Related to this line, another thread of research
focuses on synthesizing problems for visualization plots Amar et al. (2005);
Narechania et al. (2020); Fu et al. (2020); Chen et al. (2021b); Wu et al.
(2022). ARCADE also includes 57 plotting examples which are not used in this
paper. We leave collection and systematic evaluation of plotting tasks as
important feature work.

##### Context-driven Code Generation

Our work is another application of context-driven code generation, which
concerns with mapping a series of contextually dependent utterances to
executable programs, such as domain-specific logical forms Zettlemoyer and
Collins (2009); Long et al. (2016); Iyyer et al. (2017), SQL queries over
databases Suhr et al. (2018); Yu et al. (2019a, b), or general-purpose
languages (PLs) like Python Nijkamp et al. (2022). Some of those existing
works also feature EDA problems, mainly for querying structured databases,
while the problems in ARCADE exhibit more complex query patterns (_e.g._
string manipulation over semi-structured data, as $u_{1}$ in Fig. 1), in
addition to extra data wrangling tasks, thanks to the expressiveness of pandas
API and the general-purpose PL (Python) compared to domain-specific
formalisms. As the first step forward, in this paper we perform evaluation
using ground-truth solutions for previous turns (§ 5.2), similar to multi-turn
program synthesis for task-oriented dialogue in Andreas et al. (2020). We
leave conditioning on model-predicted solutions as context for evaluation
Nijkamp et al. (2022) as an important future avenue.

##  9 Conclusion

In this paper we present ARCADE, a code generation benchmark for data
wrangling and EDA tasks in computational notebooks. ARCADE features problems
with realistic NL intents and rich notebook contexts. We also develop
PaChiNCo, a $62\mathrm{B}$ LM tailed for data science, and show that PaChiNCo
outperforms public code LMs on ARCADE, while being effective in few-shot
learning to improve code style and solution diversity.

## Acknowledgements

We are grateful to Meg Risdal and Goeff Thomas from Kaggle for help with
dataset collection, and Miltos Allamanis for research discussion. We thank Jo
Chick from the research partnership team, and Rebecca Watson, Ashley Dawe and
Kimberly Herrera from Upwork to help with managing the annotation project. We
thank Aroma Mahendru for writing the data card section. We also thank
Cheriskumar Patel, Preet Patel, and Jayendra Parmar for general assistance
with the project.

## References

  * Agashe et al. (2019) Rajas Agashe, Srinivasan Iyer, and Luke Zettlemoyer. 2019.  JuICe: A large scale distantly supervised dataset for open domain context-based code generation.  In _Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP)_ , pages 5436–5446, Hong Kong, China. Association for Computational Linguistics.
  * Aggarwal et al. (2019) Charu Aggarwal, Djallel Bouneffouf, Horst Samulowitz, Beat Buesser, Thanh Hoang, Udayan Khurana, Sijia Liu, Tejaswini Pedapati, Parikshit Ram, Ambrish Rawat, et al. 2019.  How can AI automate end-to-end data science?  _arXiv preprint arXiv:1910.14436_.
  * Amar et al. (2005) Robert A. Amar, James R. Eagan, and John T. Stasko. 2005.  Low-level components of analytic activity in information visualization.  _IEEE Symposium on Information Visualization, 2005. INFOVIS 2005._ , pages 111–117.
  * Andreas et al. (2020) Jacob Andreas, Johannes Bufe, David Burkett, Charles C. Chen, Joshua Clausman, Jean Crawford, Kate Crim, Jordan DeLoach, Leah Dorner, Jason Eisner, Hao Fang, Alan Guo, David Leo Wright Hall, Kristin Delia Hayes, Kellie Hill, Diana Ho, Wendy Iwaszuk, Smriti Jha, Dan Klein, Jayant Krishnamurthy, Theo Lanman, Percy Liang, C. H. Lin, Ilya Lintsbakh, Andy McGovern, Aleksandr Nisnevich, Adam Pauls, Dmitrij Petters, Brent Read, Dan Roth, Subhro Roy, Jesse Rusak, Beth Ann Short, Div Slomin, B Snyder, Stephon Striplin, Yu Su, Zachary Tellman, Sam Thomson, A. A. Vorobev, Izabela Witoszko, Jason Wolfe, A. G. Wray, Yuchen Zhang, and Alexander Zotov. 2020.  Task-oriented dialogue as dataflow synthesis.  _Transactions of the Association for Computational Linguistics_ , 8:556–571.
  * Austin et al. (2021) Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoc Le, et al. 2021.  Program synthesis with large language models.  _arXiv preprint arXiv:2108.07732_.
  * Barke et al. (2022) Shraddha Barke, Michael B James, and Nadia Polikarpova. 2022.  Grounded copilot: How programmers interact with code-generating models.  _arXiv preprint arXiv:2206.15000_.
  * Bavishi et al. (2019) Rohan Bavishi, Caroline Lemieux, Roy Fox, Koushik Sen, and Ion Stoica. 2019.  Autopandas: neural-backed generators for program synthesis.  _Proceedings of the ACM on Programming Languages_ , 3:1 – 27.
  * Black et al. (2021) Sid Black, Leo Gao, Phil Wang, Connor Leahy, and Stella Biderman. 2021.  GPT-Neo: Large Scale Autoregressive Language Modeling with Mesh-Tensorflow.
  * Brown et al. (2020) Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. 2020.  Language models are few-shot learners.  _Advances in neural information processing systems_ , 33:1877–1901.
  * Chandel et al. (2022) Shubham Chandel, Colin B Clement, Guillermo Serrato, and Neel Sundaresan. 2022.  Training and evaluating a Jupyter notebook data science assistant.  _arXiv preprint arXiv:2201.12901_.
  * Chen et al. (2021a) Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde, Jared Kaplan, Harrison Edwards, Yura Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, David W. Cummings, Matthias Plappert, Fotios Chantzis, Elizabeth Barnes, Ariel Herbert-Voss, William H. Guss, Alex Nichol, Igor Babuschkin, S. Arun Balaji, Shantanu Jain, Andrew Carr, Jan Leike, Joshua Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Matthew M. Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob McGrew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and Wojciech Zaremba. 2021a.  Evaluating large language models trained on code.  _ArXiv_ , abs/2107.03374.
  * Chen et al. (2021b) Xinyun Chen, Linyuan Gong, Alvin Cheung, and Dawn Xiaodong Song. 2021b.  Plotcoder: Hierarchical decoding for synthesizing visualization code in programmatic context.  In _Annual Meeting of the Association for Computational Linguistics_.
  * Cheng et al. (2022) Zhoujun Cheng, Tianbao Xie, Peng Shi, Chengzu Li, Rahul Nadkarni, Yushi Hu, Caiming Xiong, Dragomir Radev, Mari Ostendorf, Luke Zettlemoyer, et al. 2022.  Binding language models in symbolic languages.  _arXiv preprint arXiv:2210.02875_.
  * Chowdhery et al. (2022) Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. 2022.  Palm: Scaling language modeling with pathways.  _arXiv preprint arXiv:2204.02311_.
  * Donoho (2017) David Donoho. 2017.  50 years of data science.  _Journal of Computational and Graphical Statistics_ , 26(4):745–766.
  * Feurer et al. (2015) Matthias Feurer, Aaron Klein, Katharina Eggensperger, Jost Springenberg, Manuel Blum, and Frank Hutter. 2015.  Efficient and robust automated machine learning.  _Advances in neural information processing systems_ , 28.
  * Fried et al. (2022) Daniel Fried, Armen Aghajanyan, Jessy Lin, Sida Wang, Eric Wallace, Freda Shi, Ruiqi Zhong, Wen-tau Yih, Luke Zettlemoyer, and Mike Lewis. 2022.  Incoder: A generative model for code infilling and synthesis.  _arXiv preprint arXiv:2204.05999_.
  * Fu et al. (2020) Siwei Fu, Kai Xiong, Xiaodong Ge, Yingcai Wu, Siliang Tang, and Wei Chen. 2020.  Quda: Natural language queries for visual data analytics.  _ArXiv_ , abs/2005.03257.
  * Gan et al. (2021) Yujian Gan, Xinyun Chen, Qiuping Huang, Matthew Purver, John R. Woodward, Jinxia Xie, and Pengsheng Huang. 2021.  Towards robustness of text-to-SQL models against synonym substitution.  pages 2505–2515, Online. Association for Computational Linguistics.
  * Gao et al. (2020) Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. 2020.  The Pile: An 800GB dataset of diverse text for language modeling.  _arXiv preprint arXiv:2101.00027_.
  * Gao et al. (2022) Luyu Gao, Aman Madaan, Shuyan Zhou, Uri Alon, Pengfei Liu, Yiming Yang, Jamie Callan, and Graham Neubig. 2022.  PAL: Program-aided language models.  _ArXiv_ , abs/2211.10435.
  * Gupta et al. (2017) Rahul Gupta, Soham Pal, Aditya Kanade, and Shirish K. Shevade. 2017.  DeepFix: fixing common C language errors by deep learning.  In _AAAI Conference on Artificial Intelligence_.
  * He et al. (2021) Xin He, Kaiyong Zhao, and Xiaowen Chu. 2021.  AutoML: a survey of the state-of-the-art.  _Knowledge-Based Systems_ , 212:106622.
  * Hendrycks et al. (2021) Dan Hendrycks, Steven Basart, Saurav Kadavath, Mantas Mazeika, Akul Arora, Ethan Guo, Collin Burns, Samir Puranik, Horace He, Dawn Song, and Jacob Steinhardt. 2021.  Measuring coding challenge competence with apps.  _NeurIPS_.
  * Heyman et al. (2021) Geert Heyman, Rafael Huysegems, Pascal Justen, and Tom Van Cutsem. 2021.  Natural language-guided programming.  _Proceedings of the 2021 ACM SIGPLAN International Symposium on New Ideas, New Paradigms, and Reflections on Programming and Software_.
  * Huang et al. (2022) Junjie Huang, Chenglong Wang, Jipeng Zhang, Cong Yan, Haotian Cui, Jeevana Priya Inala, Colin Clement, Nan Duan, and Jianfeng Gao. 2022.  Execution-based evaluation for data science code generation models.  _arXiv preprint arXiv:2211.09374_.
  * Iyyer et al. (2017) Mohit Iyyer, Wen tau Yih, and Ming-Wei Chang. 2017.  Search-based neural structured learning for sequential question answering.  In _Annual Meeting of the Association for Computational Linguistics_.
  * Jain et al. (2022) Naman Jain, Skanda Vaidyanath, Arun Iyer, Nagarajan Natarajan, Suresh Parthasarathy, Sriram Rajamani, and Rahul Sharma. 2022.  Jigsaw: Large language models meet program synthesis.  In _Proceedings of the 44th International Conference on Software Engineering_ , pages 1219–1231.
  * Jain et al. (2021) Naman Jain, Skanda Vaidyanath, Arun Shankar Iyer, Nagarajan Natarajan, Suresh Parthasarathy, Sriram K. Rajamani, and Rahul Sharma. 2021.  Jigsaw: Large language models meet program synthesis.  _2022 IEEE/ACM 44th International Conference on Software Engineering (ICSE)_ , pages 1219–1231.
  * Kandel et al. (2011) Sean Kandel, Andreas Paepcke, Joseph Hellerstein, and Jeffrey Heer. 2011.  Wrangler: interactive visual specification of data transformation scripts.  In _Proceedings of the SIGCHI Conference on Human Factors in Computing Systems_ , CHI ’11, pages 3363–3372, New York, NY, USA. Association for Computing Machinery.
  * Karmaker et al. (2020) Shubhra (Santu) Karmaker, Md. Mahadi Hassan, Micah J. Smith, Lei Xu, ChengXiang Zhai, and Kalyan Veeramachaneni. 2020.  Automl to date and beyond: Challenges and opportunities.  _ACM Computing Surveys (CSUR)_ , 54:1–36.
  * Kluyver et al. (2016) Thomas Kluyver, Benjamin Ragan-Kelley, Fernando Pérez, Brian Granger, Matthias Bussonnier, Jonathan Frederic, Kyle Kelley, Jessica Hamrick, Jason Grout, Sylvain Corlay, Paul Ivanov, Damián Avila, Safia Abdalla, Carol Willing, and Jupyter development team. 2016.  Jupyter notebooks - a publishing format for reproducible computational workflows.  In _Positioning and Power in Academic Publishing: Players, Agents and Agendas_ , pages 87–90, Netherlands. IOS Press.
  * Kojima et al. (2022) Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa. 2022.  Large language models are zero-shot reasoners.  _ArXiv_ , abs/2205.11916.
  * Lachaux et al. (2020a) Marie-Anne Lachaux, Baptiste Rozière, Lowik Chanussot, and Guillaume Lample. 2020a.  Unsupervised translation of programming languages.  _ArXiv_ , abs/2006.03511.
  * Lachaux et al. (2020b) Marie-Anne Lachaux, Baptiste Roziere, Lowik Chanussot, and Guillaume Lample. 2020b.  Unsupervised translation of programming languages.  _arXiv preprint arXiv:2006.03511_.
  * Lai et al. (2022) Yuhang Lai, Chengxi Li, Yiming Wang, Tianyi Zhang, Ruiqi Zhong, Luke Zettlemoyer, Scott Wen tau Yih, Daniel Fried, Sida Wang, and Tao Yu. 2022.  Ds-1000: A natural and reliable benchmark for data science code generation.  _ArXiv_ , abs/2211.11501.
  * Lee et al. (2021) Chia-Hsuan Lee, Oleksandr Polozov, and Matthew Richardson. 2021.  Kaggledbqa: Realistic evaluation of text-to-sql parsers.  In _ACL_.
  * Li et al. (2022) Yujia Li, David H. Choi, Junyoung Chung, Nate Kushman, Julian Schrittwieser, Rémi Leblond, Tom, Eccles, James Keeling, Felix Gimeno, Agustin Dal Lago, Thomas Hubert, Peter Choy, Cyprien de, Masson d’Autume, Igor Babuschkin, Xinyun Chen, Po-Sen Huang, Johannes Welbl, Sven Gowal, Alexey, Cherepanov, James Molloy, Daniel Jaymin Mankowitz, Esme Sutherland Robson, Pushmeet Kohli, Nando de, Freitas, Koray Kavukcuoglu, and Oriol Vinyals. 2022.  Competition-level code generation with alphacode.  _ArXiv_ , abs/2203.07814.
  * Lin et al. (2022) Zi Lin, Jeremiah Liu, and Jingbo Shang. 2022.  Neural-symbolic inference for robust autoregressive graph parsing via compositional uncertainty quantification.  In _Proceedings of EMNLP_.
  * Long et al. (2016) Reginald Long, Panupong Pasupat, and Percy Liang. 2016.  Simpler context-dependent logical forms via model projections.  _ArXiv_ , abs/1606.05378.
  * Narechania et al. (2020) Arpit Narechania, Arjun Srinivasan, and John T. Stasko. 2020.  Nl4dv: A toolkit for generating analytic specifications for data visualization from natural language queries.  _IEEE Transactions on Visualization and Computer Graphics_ , 27:369–379.
  * Nazabal et al. (2020) Alfredo Nazabal, Christopher K I Williams, Giovanni Colavizza, Camila Rangel Smith, and Angus Williams. 2020.  Data engineering for data analytics: A classification of the issues, and case studies.  _arXiv_.
  * Nijkamp et al. (2022) Erik Nijkamp, Bo Pang, Hiroaki Hayashi, Lifu Tu, Huan Wang, Yingbo Zhou, Silvio Savarese, and Caiming Xiong. 2022.  A conversational paradigm for program synthesis.  _arXiv preprint arXiv:2203.13474_.
  * Nye et al. (2021) Maxwell Nye, Anders Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, Charles Sutton, and Augustus Odena. 2021.  Show your work: Scratchpads for intermediate computation with language models.  _ArXiv_ , abs/2112.00114.
  * Pasupat and Liang (2015) Panupong Pasupat and Percy Liang. 2015.  Compositional semantic parsing on semi-structured tables.  In _Annual Meeting of the Association for Computational Linguistics_.
  * Perkel (2021) Jeffrey Perkel. 2021.  Reactive, reproducible, collaborative: computational notebooks evolve.  _Nature_ , 593.
  * Pimentel et al. (2019) João Felipe Pimentel, Leonardo Gresta Paulino Murta, Vanessa Braganholo, and Juliana Freire. 2019.  A large-scale study about quality and reproducibility of Jupyter notebooks.  _2019 IEEE/ACM 16th International Conference on Mining Software Repositories (MSR)_ , pages 507–517.
  * Rajkumar et al. (2022) Nitarshan Rajkumar, Raymond Li, and Dzmitry Bahdanau. 2022.  Evaluating the text-to-sql capabilities of large language models.  _arXiv preprint arXiv:2204.00498_.
  * Su et al. (2021) Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. 2021.  Roformer: Enhanced transformer with rotary position embedding.  _ArXiv_ , abs/2104.09864.
  * Suhr et al. (2018) Alane Suhr, Srini Iyer, and Yoav Artzi. 2018.  Learning to map context-dependent sentences to executable formal queries.  _ArXiv_ , abs/1804.06868.
  * Thoppilan et al. (2022) Romal Thoppilan, Daniel De Freitas, Jamie Hall, Noam M. Shazeer, Apoorv Kulshreshtha, Heng-Tze Cheng, Alicia Jin, Taylor Bos, Leslie Baker, Yu Du, Yaguang Li, Hongrae Lee, Huaixiu Zheng, Amin Ghafouri, Marcelo Menegali, Yanping Huang, Maxim Krikun, Dmitry Lepikhin, James Qin, Dehao Chen, Yuanzhong Xu, Zhifeng Chen, Adam Roberts, Maarten Bosma, Yanqi Zhou, Chung-Ching Chang, I. A. Krivokon, Willard James Rusch, Marc Pickett, Kathleen S. Meier-Hellstern, Meredith Ringel Morris, Tulsee Doshi, Renelito Delos Santos, Toju Duke, Johnny Hartz Søraker, Ben Zevenbergen, Vinodkumar Prabhakaran, Mark Díaz, Ben Hutchinson, Kristen Olson, Alejandra Molina, Erin Hoffman-John, Josh Lee, Lora Aroyo, Ravindran Rajakumar, Alena Butryna, Matthew Lamm, V. O. Kuzmina, Joseph Fenton, Aaron Cohen, Rachel Bernstein, Ray Kurzweil, Blaise Aguera-Arcas, Claire Cui, Marian Croak, Ed Chi, and Quoc Le. 2022.  Lamda: Language models for dialog applications.  _ArXiv_ , abs/2201.08239.
  * Transformers et al. (2022) Memorizing Transformers, Yuhuai Wu, Markus N. Rabe, DeLesley S. Hutchins, and Christian Szegedy. 2022.  M emorizing t ransformers.
  * Tunstall et al. (2022) Lewis Tunstall, Leandro von Werra, and Thomas Wolf. 2022.  _Natural language processing with transformers_.  " O’Reilly Media, Inc.".
  * Wang et al. (2022a) April Yi Wang, Dakuo Wang, Jaimie Drozdal, Michael Muller, Soya Park, Justin D Weisz, Xuye Liu, Lingfei Wu, and Casey Dugan. 2022a.  Documentation matters: Human-centered ai system to assist data science code documentation in computational notebooks.  _ACM Transactions on Computer-Human Interaction_ , 29(2):1–33.
  * Wang and Komatsuzaki (2021) Ben Wang and Aran Komatsuzaki. 2021.  GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model.  https://github.com/kingoflolz/mesh-transformer-jax.
  * Wang et al. (2021a) Dakuo Wang, Josh Andres, Justin D Weisz, Erick Oduor, and Casey Dugan. 2021a.  Autods: Towards human-centered automation of data science.  In _Proceedings of the 2021 CHI Conference on Human Factors in Computing Systems_ , pages 1–12.
  * Wang et al. (2021b) Dakuo Wang, Q Vera Liao, Yunfeng Zhang, Udayan Khurana, Horst Samulowitz, Soya Park, Michael Muller, and Lisa Amini. 2021b.  How much automation does a data scientist want?  _arXiv preprint arXiv:2101.03970_.
  * Wang and Cho (2016) Tian Wang and Kyunghyun Cho. 2016.  Larger-context language modelling with recurrent neural network.  In _Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 1319–1329, Berlin, Germany. Association for Computational Linguistics.
  * Wang et al. (2022b) Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc Le, Ed Chi, and Denny Zhou. 2022b.  Self-consistency improves chain of thought reasoning in language models.  _arXiv preprint arXiv:2203.11171_.
  * Wei et al. (2022) Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Ed Chi, Quoc Le, and Denny Zhou. 2022.  Chain of thought prompting elicits reasoning in large language models.  _ArXiv_ , abs/2201.11903.
  * Wu et al. (2022) Zhengkai Wu, Vu Le, Ashish Tiwari, Sumit Gulwani, Arjun Radhakrishna, Ivan Radicek, Gustavo Soares, Xinyu Wang, Zhenwen Li, and Tao Xie. 2022.  NL2Viz: Natural language to visualization via constrained syntax-guided synthesis.  _Proceedings of the 30th ACM Joint European Software Engineering Conference and Symposium on the Foundations of Software Engineering_.
  * Xie et al. (2022) Tianbao Xie, Chen Henry Wu, Peng Shi, Ruiqi Zhong, Torsten Scholak, Michihiro Yasunaga, Chien-Sheng Wu, Ming Zhong, Pengcheng Yin, Sida I. Wang, Victor Zhong, Bailin Wang, Chengzu Li, Connor Boyle, Ansong Ni, Ziyu Yao, Dragomir R. Radev, Caiming Xiong, Lingpeng Kong, Rui Zhang, Noah A. Smith, Luke Zettlemoyer, and Tao Yu. 2022.  Unifiedskg: Unifying and multi-tasking structured knowledge grounding with text-to-text language models.  _ArXiv_ , abs/2201.05966.
  * Xu et al. (2022) Frank F. Xu, Uri Alon, Graham Neubig, and Vincent J. Hellendoorn. 2022.  A systematic evaluation of large language models of code.  _Proceedings of the 6th ACM SIGPLAN International Symposium on Machine Programming_.
  * Yin and Neubig (2017) Pengcheng Yin and Graham Neubig. 2017.  A syntactic neural model for general-purpose code generation.  In _Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 440–450, Vancouver, Canada. Association for Computational Linguistics.
  * Yu et al. (2019a) Tao Yu, Rui Zhang, He Yang Er, Suyi Li, Eric Xue, Bo Pang, Xi Victoria Lin, Yi Chern Tan, Tianze Shi, Zihan Li, Youxuan Jiang, Michihiro Yasunaga, Sungrok Shim, Tao Chen, Alexander R. Fabbri, Zifan Li, Luyao Chen, Yuwen Zhang, Shreya Dixit, Vincent Zhang, Caiming Xiong, Richard Socher, Walter S. Lasecki, and Dragomir R. Radev. 2019a.  Cosql: A conversational text-to-sql challenge towards cross-domain natural language interfaces to databases.  In _Conference on Empirical Methods in Natural Language Processing_.
  * Yu et al. (2019b) Tao Yu, Rui Zhang, Michihiro Yasunaga, Yi Chern Tan, Xi Victoria Lin, Suyi Li, He Yang Er, Irene Z Li, Bo Pang, Tao Chen, Emily Ji, Shreya Dixit, David Proctor, Sungrok Shim, Jonathan Kraft, Vincent Zhang, Caiming Xiong, Richard Socher, and Dragomir R. Radev. 2019b.  Sparc: Cross-domain semantic parsing in context.  _ArXiv_ , abs/1906.02285.
  * Zettlemoyer and Collins (2009) Luke Zettlemoyer and Michael Collins. 2009.  Learning context-dependent mappings from sentences to logical form.  In _Annual Meeting of the Association for Computational Linguistics_.
  * Zöller and Huber (2021) Marc-André Zöller and Marco F Huber. 2021.  Benchmark and survey of automated machine learning frameworks.  _Journal of artificial intelligence research_ , 70:409–472.

Supplementary Materials

##  Appendix A Outline of ARCADE Annotation Guideline

In this section we provide a brief summary of the outline in our annotation
guideline.

##### _Existing Tasks_

The annotators are given a list of Jupyter notebooks. Each notebook uses
pandas to perform certain data analysis tasks. For each notebook, an annotator
is asked to:

  1. 1.

Identify code cells that contain instructive code snippets that perform data
wrangling or exploratory data analysis tasks.

  2. 2.

Fix the notebook and make them clean and executable.

  3. 3.

For each code snippet identified in Step 1, create natural language
descriptions of the task. Also verify the code solution and fix them as
appropriate. Finally, remove any redundant text in the notebook (_e.g._
solution outline or hints for tutorial notebooks) that could give away to the
refernce solution.

##### Instruction on Creating Natural Intents

Specifically, for step 3, in order to collect realistic NL intents, the
annotators are given the following high-level description, followed by
detailed instructions and examples.

Below we share some suggestions to write good intents.

Keep it natural without redundant explanations. Imagine an AI programmer can
help you accomplish simple data wrangling and EDA tasks, what kind of intents
will you send to the system? Our goal is to collect real inputs to such a
system from data scientists like you.

One idea to write good intents is to keep it concise such that another
programmer could quickly understand and implement a solution that executes to
the same outputs. You are encouraged to create simple, short intents while
describing the desired outputs without much ambiguity.

##### _New Tasks_

For each ML dataset we provided, an annotator creates a Colab notebook with
code snippets for some interesting data wrangling and exploratory data
analysis tasks using this dataset. Each code snippet is paired with its
natural language intent, simliar to the process of annotating _Existing
Tasks_. We ask annotators to feel free to work on any tasks that they may find
interesting for the given dataset, as long as the code solution for the task
should consist of multiple lines and use different pandas API functions.
Different from annotating _Existing Tasks_ , we ask them to first create a
natural language intent for their task, and then write a code solution in the
next cell.

Below is an excerpt from the annotation guideline describing the types of data
wranling and EDA tasks to create.

What Tasks to Create

In general, you may create whatever exploratory data analysis tasks that you
find interesting for the given datasets. To come up with interesting tasks,
you can think in this way: before training your ML models for the dataset,
what kind of data wrangling or EDA tasks would you like to perform on the
dataset? Below are some more concrete descriptions of such wrangling or EDA
tasks:

Data Preprocessing/Wrangling Tasks which involves modifying existing
dataframes or creating new ones. Such as normalizing column names, adding new
columns, modifying existing columns (e.g., converting string values to date
times), generating new dataframes using ops like group_by, and so on. Some
datasets we shared are just raw data without any preprocessing or cleaning.
Feel free to . Please also refer to Section: Identify Code Snippets to
Annotate in our previous annotation guideline for more examples.

Exploratory Data Analysis Tasks that Require Some Wrangling and Preprocessing
Answering interesting EDA questions using the given dataset, but some data
wrangling steps are required in order to derive the answer. For example, given
a dataframe df of user shopping history and credit card expiration dates in
the format of df.loc[0][‘cc_exp’] = ‘08/26’. To answer the EDA question “How
many users have a credit card expiring in 2024?”, we need to first convert the
expiration year from the string-formatted cc_exp column.

To encourage the annotators to create more complex tasks, we also provide the
following high-level instruction:

Complexity of Tasks

You should create relatively complex tasks that require multiple steps and
also a combination of different pandas APIs to solve them. Avoid problems that
can be solved using one-liner code such as df.group_by(...).sort_values(...).
An ideal task should be reasonably complex and needs to be broken down into
multiple smaller steps to solve, and each step may require using one or
multiple pandas functions.

As a general rule of thumb, you should aim at creating tasks that either have
at least 50 tokens or use at least 4 pandas APIs (dataframe/series indexing,
like df[df[‘continent’] == ‘NA’] is also counted as one API usage). You can
find more concrete example tasks at the end of this doc.

##### Full Guideline

Our annotation guideline is 35-pages long in total, which we provide on a per
request basis. Please contact<EMAIL_ADDRESS>to request access.

##  Appendix B Details of Fine-tuning PaChiNCo

We take the fully trained PaLM 62B model that has been trained on
$1,325\mathrm{B}$ tokens Chowdhery et al. (2022) as the base model. To achieve
the best performance on , we finetune the base model in two stages: a python
only and a ARCADE only finetuning. For the two-stage code fine-tuning (§ 4),
we use the similar training recipe of the base PaLM 62B model. More
specifically, we apply the learning rate decay scheduling of $0.2/\sqrt{t}$
where $t$ is the number of steps. We train the model for 124K steps (1 epoch)
with a batch size of 256 on 512 TPU v4 chips. Afterwards, we reload the
optimizer state and continue training on the Jupyter notebooks data
($10\mathrm{B}$ tokens) using the same hyper parameter for 3 epochs ($\sim
572\mathrm{K}$ steps).

##### Linearize Notebooks to Python Source Code

We convert computational notebooks for finetuning (§ 4) and evaluation (§ 5.2)
to Python source code using nbconverter. Specifically, markdown and code cells
in a notebook are concatenated using the special delimiter ‘# In[]:’, and text
in markdown cells is commented out using the ‘# ’ prefix. See Fig. 39 for an
example of the linearized notebook for Fig. 1 (up to $c_{3}$). Jupyter
notebooks that are converted to Python files in such format are common in
GitHub repositories, which mitigates the domain transfer gap between general
Python code and notebook-specific data, and also allows us to prompt public
code LLMs, which might have seen similar data during pre-training.

##  Appendix C Performance of Python-finetuned PaLM 62B on Code Benchmarks

We evaluate the PaLM $62\mathrm{B}$ model after the first-stage fine-tuning on
Python data (§ 4) on other code generation benchmarks, namely HumanEval Chen
et al. (2021a), MBPP Austin et al. (2021) and Transcoder Lachaux et al.
(2020b), following the evaluation setup of PaLM-Coder in Chowdhery et al.
(2022). Tab. 3 lists the results. With $7\times$ more Python code tokens, our
Python fine-tuned PaLM $62\mathrm{B}$ model outperforms the $8\times$ larger
PaLM-Coder $540\mathrm{B}$ model on all the benchmarks, including the
multilingual code translation task.

Dataset | Human Eval | MBPP | Transcoder
---|---|---|---
Metric |  pass@$100$ |  pass@$80$ |  pass@$25$
PaLM-Coder $540\mathrm{B}$ Chowdhery et al. (2022) | $88.4\%$ | $80.8\%$ | $82.5\%$
PaLM $62\mathrm{B}$ Python (§ 4) | $\mathbf{91.5\%}$ | $\mathbf{86.0\%}$ | $\mathbf{86.4\%}$
Table 3: Python fine-tuned PaLM 62B model evaluation results on public code
benchmarks.

##  Appendix D Inference Setup

For CodeGen, we use the inference script from the official GitHub repository.
For InCoder, we follow the official inference example script and use the
release on Huggingface model hub. We convert each example in our dataset to
Python source code to a prompt, as outlined in § 5.2. Notebooks are linearized
using nbconverter similar as generating fine-tuning data (Appendix B). One
exception is InCoder, for which we follow Fried et al. (2022) and use the
Jupyter notebook linearization template used in its pre-training.

At inference time, we left-truncate notebook context up to 900 tokens
(measured by PaLM’s vocabulary), which fit in the context window size of all
LLMs we evaluated. We also make sure to always include NL schema descriptions
in prompts given their importance in understanding NL intents. In addition,
for few-shot experiments in § 5.4, we use additional 1,200 tokens to
accommodate the prompt prefix, making the total maximal prompt length to be
2,100. Due to its excessive length, we only perform few-shot prompting
experiments on PaChiNCo since its rotatory positional embedding Su et al.
(2021) could generalize to encode longer contexts at inference time. During
sampling, we set the maximum target length to be 512 tokens.

##  Appendix E pass@$1$ for PaChiNCo

_Existing Tasks_ | _New Tasks_
---|---
48.8 | 17.5
Table 4: pass@$1$ for PaChiNCo (temperature=0.2)

For reference, we also report pass@$1$ of PaChiNCo on ARCADE. Samples are
generated using the same hyper parameters as in § 5.3, except for the
temperature, which is 0.2.

##  Appendix F Scaling Curve of CodeGen on ARCADE

Figure 13: Scaling curve of CodeGen on ARCADE and existing code generation
benchmarks. Results on HumanEval are selected from the best temperature
$t\in\\{0.2,0.6,0.8\\}$ Nijkamp et al. (2022).

Fig. 13 depicts the scaling curve of ARCADE with respect to the number of
parameters for CodeGen${}_{\textrm{mono}}$ models. The pass rate scales nearly
linearly as a function of model size, and the performance has not saturated,
especially on the _New Tasks_ split. This shows ARCADE is a reliable dataset
to study the scaling behavior of code LLMs. The slope of the curve on _New
Tasks_ is also smaller than on other datasets, suggesting that this problem
set is more challenging for CodeGen models.

##  Appendix G Comparing with Codex on ARCADE

pass@$k$ | Existing Tasks | New Tasks
---|---|---
5 | 10 | 20 | 30 | 5 | 10 | 20 | 30
PaChiNCo ($62\mathrm{B}$) | $64.6$ | $71.0$ | $76.0$ | $78.3$ | $30.6$ | $38.0$ | $45.0$ | $48.6$
Codex-cushman-001 (12B) | $51.0$ | $59.0$ | $65.9$ | $69.4$ | $14.2$ | $20.1$ | $27.0$ | $31.3$
Codex-davinci-002 | $66.1$ | $72.6$ | $78.3$ | $81.2$ | $36.0$ | $43.8$ | $50.8$ | $54.8$
Table 5: _Pass_ $@k$ evaluation on ARCADE using Codex API.

For reference, we also report the results on ARCADE using the public Codex
API, following the same evaluation setting as in Tab. 2. PaChiNCo
significantly outperforms the smaller Codex-cushman-001 API, on par with
Codex-davinci-002 on _Existing Tasks_ , while Codex-davinci-002 is stronger on
_New Tasks_. While we cannot gain much insight from the results due to limited
knowledge about Codex-davinci-002, through error analysis, we find that Codex-
davinci-002 is better at instruction following, especially in understanding NL
descriptions of complex DataFrame schema (§ 5.2). We leave improving the
instruction following skills of PaChiNCo as interesting future work.

##  Appendix H Performance on Problems with Different Complexity

Figure 14: pass@$k$ of PaChiNCo w.r.t AST size of reference programs.

Fig. 14 plots pass@$30$ with respect to the AST size of reference programs.
Similar to Fig. 5, results on _New Tasks_ are generally lower. Meanwhile, it
seems that AST size correlates better with pass@$k$ compared to the number of
API usage, while the latter metric offers more intuitive information about the
data transformation steps involved.

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

Figure 15: Frequency of runtime errors from PaChiNCo’s predictions aggreated
over the two splits with varying amount of notebook context cells.

In Fig. 9 we observe that pass@$k$ improves when more notebook context is
available. Fig. 15 illustrates the distribution of execution error types on
failed predictions. Notably, using more notebook context cells significantly
reduces the chance of NameErrors caused by using undefined variables in
context. The number of KeyErrors is also reduced, indicating that the model
makes fewer schema understanding errors when referring to columns in
DataFrames.

##  Appendix J Few-shot Prompting Results on _Existing Tasks_ Split

Figure 16: pass@$30$ for few-shot prompting on _Existing Tasks_

Figure 17: Code style metrics for few-shot prompting on _Existing Tasks_

We report results of prompting PaChiNCo using few-shot exemplars on the
_Existing Tasks_ split in Fig. 17. Compared to the results obtained on _New
Tasks_ (Fig. 7), while few-shot prompting, especially step-by-step prompting,
is still effective compared to the baseline, the gap is not as profound as the
results on _New Tasks_. The difference between different prompting methods is
also less significant, and using natural language explanations
($\mathrm{SbS}+\mathrm{Explanations}$) seems to hurt performance on simpler
tasks, while still beneficial for complex problems with more than 7 pandas API
calls. This is likely due to that the model relies on memorization to solve
problems in _Existing Tasks_ , rendering using extra exemplars to control
generation less effective. Nevertheless, these results reiterate the value of
_New Tasks_ as a more reliable benchmark to better differentiate different
prompting strategies.

##  Appendix K Further Analysis of Solution Diversity

(a)

(b)

(c)

Figure 18: Cumulative distributions of the number of unique API sequences
extracted from PaChiNCo’s $50$ predictions on the New Tasks split, for all
predictions (top left), only executable predictions (top right), or only
correct predictions (bottom). Curves that appear more to the right represent
prompting methods with greater diversity in the samples.

Here we provide further analysis of the diversity in solution patterns,
measured by the number of distinct pandas API call sequences used in the
samples. Fig. 18 shows a cumulative distribution of the number of solution
patterns for different subsets of the predictions: all predictions, only those
that execute successfully, and only the correct predictions. In each case, we
see that step-by-step prompting leads to increased diversity compared to the
baselines, and prompting with NL explanations further increases the diversity.
While increased diversity is helpful in finding a correct solution at higher
sample size $k$, it is also helpful when considering only correct solutions
because a user might want to see a variety of solution approaches, whether for
educational purposes or to choose the one they like best (which is partially
subjective). Refer to § 7 for such examples.

##  Appendix L Error Analysis Examples

In this section, we provided a preliminary study on selected example model
errors. We analyzed two types of errors: execution error and semantic error
(executable but incorrect). Execution error has error message from the
notebook environment. We can classify these errors into more fine-grained
categories in Fig. 12. As the result shows, KeyError is the top error mode in
the execution errors. Over 50% of the KeyError are associated with the
pd.groupby API call. pd.groupby API call changes the dataframe schema as the
model generates the code. For example, pd.groupby(*).mean() will remove non-
numeric columns in the dataframe. This requires the model to have a profound
understanding of the dataframe schema. We gave an example in Fig. 24. The
column shipping_fee is string value which will be removed after
df.groupby(ship_state).sum(). The secondary source of execution error is
AttributeError, which shares a similar cause to the KeyError. This is because
AttributeError is often triggered by calling a non-existing column as an
attribute of a pd.dataframe. An example is given in Fig. 25, where the model
tries to call the non-existing column signupdate as an attribute of df_users,
leading to the AttributeError. These two error modes suggest that building a
better schema aware language model is a promising future research direction.
We also present Fig. 26 and Fig. 27 as the example for TypeError and
ValueError respectively. These two error modes are often caused by
insufficient knowledge of the column values. For example, the model tried to
compare a string-value column to a integer in Fig. 26, which causes TypeError.
Fig. 27 showcased that the model tries to apply numeric operation
pd.DataFrame.mean() on a column with NaN values, leading to ValueError.

Semantic errors account for more than 65% of the incorrect predictions. We
cannot get any error messages from these executable but incorrect programs.
Therefore, we conducted a manual analysis on 50% randomly sampled incorrect
predictions. We classified the semantic errors into 5 categories: 1. Complex
task requiring non-trivial computation ($43\%$); 2. NL misunderstanding
($26\%$); 3. Under-specified intents ($19\%$); 4. False-negatives of the
fuzzy-matching evaluation metric ($6\%$); 5. Annotation Errors ($6\%$). To
complement the discussion in § 6, we showcase examples of these types of
errors. The primary source of semantic error is complex reasoning. Examples
are given in Fig. 28 and Fig. 29. In Fig. 28, the model need to infer that
last 10 year can be computed by dt.datetime.today().year - 10 then apply
string operation str.contains(accident) to select the required contents. Fig.
29 is another example of complex data wrangling steps. To generate the correct
program, the model need to compare the current rank to the past rank, ensure
the 2021 rank column exists and then aggregate the information. NL
misunderstanding is the secondary source of semantic errors. In Fig. 30, the
generated output does not reflect the institute constraint in the intent.
Another source of semantic errors is under-specified intent, which happens
occasionally in natural language instructions. In Fig. 31, the intent does not
specify the output format. The program generated by the model calculates the
sum of front and left facing trees while the reference calculates them
respectively. The evaluation fails but the model output is not necessarily
incorrect in this case. Fig. 32 illustrates another type of semantic error:
evaluation coverage. In this example, the model gives an acceptable answer
which only differs from the reference in the column order. It does not pass
the evaluation due to its limited coverage.

$\bm{u}$: What are the five most common genres for shows released during the
year having the
highest percent increase in votes?

Figure 19: An example prediction from PaChiNCo with a complex single step.

$\bm{u}$: Convert crash dates to datetime and show the total number of
vehicles involved in crashes over the years.

Figure 20: An example prediction from PaChiNCo that explains the semantics of
a column.

$\bm{u}$: What is the distribution of student adaptivity level across each age
group?
Return a DataFrame with age groups as an index and adaptivity levels as
columns.

Figure 21: An example prediction from PaChiNCo that explains the return value
from a complex API.

$\bm{u}$: Drop columns with more than 70 percent null values

Figure 22: Two predictions from PaChiNCo with different solution approaches.

$\bm{u}$: In which year, within the last ten years, did Israel receive the
highest amount of financial aid, in constant amount? Show the year and amount
received.

Figure 23: Two predictions from PaChiNCo. NL explanations help users identify
incorrect steps.

$\bm{u}$: What are the average shipping fees for each state, starting from
highest to lowest? (rounded to 2 decimal places)?

Figure 24: An example of KeyError: the model calls a column which is removed
after the pd.groupby().mean() API call.

$\bm{u}$: Show how many new users signed up for every year since 2000

Figure 25: An example of AttributeError: the model tries to call a non-
existing column.

$\bm{u}$: What are the top five models with most number of bikes having
mileage less than 5000 kilometers?

Figure 26: An example of TypeError: the model is intent to compare a string-
value column to an integer.

$\bm{u}$: What is the average number of filed charges for drug related cases?

Figure 27: An example program of ValueError: the model tries to calculate the
mean of a column containing NaN values.

$\bm{u}$: What is the number of deaths by accident as a percentage of total
deaths in the last ten years?

Figure 28: An example of complex reasoning: the model has to infer the API
call (dt.datetime.today().year - 10) from the last-ten-years constraint in the
intent.

$\bm{u}$: Which hotels had a worse ranking this year than in 2021? Show the
hotel name, location and the difference in ranking from last year.

Figure 29: An example of complex reasoning: the model needs to compare the
current rank and the past rank while making sure the rank in 2021 exists.

$\bm{u}$: What was the mean opening rank across all IIT institutes for each
program over the years? Show the mean opening rank for each year in columns
with program as index and consider only general students.

Figure 30: An example of NL misunderstanding: the model does not filter the
institute type according to the intent.

$\bm{u}$: How many front and left facing trees were planted in that park?

Figure 31: An example of under-specified intent. It does not specify the
output should not sum the number of front facing and left facing trees.

$\bm{u}$: Return a matrix with the average ticket prices to and from all the
cities for each ticket class.

(a)

| destination_city | Bangalore | Chennai | Delhi | Hyderabad | Kolkata | Mumbai
---|---|---|---|---|---|---|---
class | source_city |  |  |  |  |  |
Business | Bangalore | NaN | 52436.915395 | 48144.337108 | 50395.796948 | 58854.693091 | 58024.618208
| Chennai | 53113.008692 | NaN | 52443.367242 | 51559.874283 | 57078.895872 | 56223.838086
| Delhi | 48576.027921 | 52031.778099 | NaN | 44457.376775 | 56239.853659 | 44364.442811
| Hyderabad | 50358.290706 | 51132.155288 | 44250.700281 | NaN | 53729.157762 | 52184.424666
| Kolkata | 58681.104437 | 56502.775035 | 55047.492193 | 54732.447908 | NaN | 57422.551724
| Mumbai | 57970.544389 | 55703.326197 | 43846.329273 | 51593.643678 | 57106.526385 | NaN
Economy | Bangalore | NaN | 7105.953850 | 6124.897982 | 6360.141698 | 7375.638594 | 6381.093332
| Chennai | 7175.020192 | NaN | 6075.961190 | 5960.788831 | 7547.295815 | 6529.119453
| Delhi | 6175.622535 | 6102.317245 | NaN | 6031.164261 | 7045.621678 | 6059.826087
| Hyderabad | 6234.882649 | 6049.884930 | 6072.296659 | NaN | 6881.680392 | 5969.259906
| Kolkata | 7471.621990 | 8011.745229 | 7161.400077 | 7489.144374 | NaN | 7405.787239
| Mumbai | 6432.511946 | 6420.917984 | 5889.281400 | 5774.891130 | 7227.971735 | NaN

(b)

class | Business | Economy
---|---|---
destination_city | Chennai | Delhi | Hyderabad | Kolkata | Mumbai | Bangalore | Chennai | Delhi | Hyderabad | Kolkata | Mumbai | Bangalore
source_city |  |  |  |  |  |  |  |  |  |  |  |
Bangalore | 52437.0 | 48144.0 | 50396.0 | 58855.0 | 58025.0 | NaN | 7106.0 | 6125.0 | 6360.0 | 7376.0 | 6381.0 | NaN
Chennai | NaN | 52443.0 | 51560.0 | 57079.0 | 56224.0 | 53113.0 | NaN | 6076.0 | 5961.0 | 7547.0 | 6529.0 | 7175.0
Delhi | 52032.0 | NaN | 44457.0 | 56240.0 | 44364.0 | 48576.0 | 6102.0 | NaN | 6031.0 | 7046.0 | 6060.0 | 6176.0
Hyderabad | 51132.0 | 44251.0 | NaN | 53729.0 | 52184.0 | 50358.0 | 6050.0 | 6072.0 | NaN | 6882.0 | 5969.0 | 6235.0
Kolkata | 56503.0 | 55047.0 | 54732.0 | NaN | 57423.0 | 58681.0 | 8012.0 | 7161.0 | 7489.0 | NaN | 7406.0 | 7472.0
Mumbai | 55703.0 | 43846.0 | 51594.0 | 57107.0 | NaN | 57971.0 | 6421.0 | 5889.0 | 5775.0 | 7228.0 | NaN | 6433.0

(c)

Figure 32: An example of plausible alternative prediction that is labeled as
incorrect due to limited coverage of the evaluation metric.

##  Appendix M Data Card for the Training Data of PaChiNCo

We provide a data card for the training data of PaChiNCo as outlined in § 4,
following the structure presented in Chowdhery et al. (2022). We also report
training data composition in Tab. 7.

Motivation
---
For what purpose was the dataset created? Who created the dataset? Who funded the creation of the dataset? |  The dataset was created for training code and language models by a team of researchers at Google.
Composition
What do the instances that comprise the dataset represent (e.g., documents, photos, people, countries)? |  Dataset comprises of Python source code files and Jupyter notebooks from GitHub, filtered by license so as to exclude code with restrictive licenses.
How many instances are there in total (of each type, if appropriate)? |  The data makeup is given in Table 7.
Does the dataset contain all possible instances or is it a sample (not necessarily random) of instances from a larger set? |  The dataset is a small (random) subset of a larger set.
What data does each instance consist of? |  Each instance is encoded content of a source code file.
Is there a label or target associated with each instance? |  No, there are no labels associated with each instance.
Is any information missing from individual instances? |  No.
Are relationships between individual instances made explicit? |  No.
Are there recommended data splits? |  We use random splits for the training, validation, and test.
Are there any errors, sources of noise, or redundancies in the dataset? |  • Python files were near deduplicated at the file level using a custom implementation of minhash algorithm, so lower level redundancies (lines, code blocks) may still exist. • Some files were misclassified in the license tagging and filtration process given that license classification algorithm can have false positives and negatives.
Is the dataset self-contained, or does it link to or otherwise rely on external resources? |  The dataset is self-contained.
Does the dataset contain data that might be considered confidential? |  No.
Does the dataset contain data that, if viewed directly, might be offensive, insulting, threatening, or might otherwise cause anxiety? |  Given the dataset contains source code, it is not likely there is any offensive text in it, however no explicit measures are in place to eliminate such data if it were present.
Collection Process
How was the data associated with each instance acquired? |  The data was collected from publicly available sources.
What mechanisms or procedures were used to collect the data? |  The data was collected using a variety of software programs to extract and clean source code files.
If the dataset is a sample from a larger set, what was the sampling strategy? |  The dataset is small subset of publicly available code from Github, sampled randomly.
Who was involved in the data collection process? |  A team of researchers at Google.
Over what timeframe was the data collected? |  April - July 2022
Were any ethical review processes conducted? |  No.
Preprocessing, cleaning, and labeling
Was any preprocessing, cleaning, or labeling of the data done (e.g., discretization or bucketing, tokenization, part-of-speech tagging, SIFT feature extraction, removal of instances, processing of missing values)? |  License filtration, quality filtration and deduplication were applied to the source code files.  • License classification was done using Google License Classifier library. Source code files with restricted licenses were filtered out. • Python files were deduplicated at the file level using a custom variant of minhash algorithm. Locality sensitive hashes of file content were used to create partitions of potentially duplicate files based on collisions in the hash buckets. For each pair in the partitions, Jaccard Similarity and Edit Distance scores were calculated to create an "edge" for a pair whenever the scores are higher than the specified threshold. This was followed by application of connected components algorithm to return the sets of duplicates. • Jupyter notebooks were first deduplicated following the same procedure as deduplicating Python files, and then deduplicated at individual cell level against the evaluation dataset (§ 4).
Is the software used to preprocess, clean, or label the instances available? |  No.
Uses
Has the dataset been used for any tasks already? |  Yes, we use the dataset for pre-training other code and language models.
Is there a repository that links to any or all papers or systems that use the dataset? |  No.
What (other) tasks could the dataset be used for? |  The dataset can be used for training of other code and language models.
Is there anything about the composition of the dataset or the way it was collected and pre-processed/cleaned/labeled that might impact future uses? |  The dataset is static in nature and thus will become progressively more “stale”. It will not include any new source code repositories that were created/updated later on Github.
Are there tasks for which the dataset should not be used? |  This should not be used for any unacceptable code or language modeling use cases e.g. generating code or language with toxic/biased connotations.
Distribution
Will the dataset be distributed to third parties outside of the entity (e.g., company, institution, organization) on behalf of which the dataset was created? |  No.
Language | Tokens | Source Files
---|---|---
Python | 63,786,481,126 | 60,397,107
Jupyter Notebooks | 9,613,648,619 | 3,796,713
Table 7: Data Composition of the fine-tuning data for PaChiNCo.

##  Appendix N Detailed Prompting Examples

In this section we provide detailed examples of prompts used in our
experiments. As in § 5.2, there are two categories of experiments in § 5,
namely prompting using notebook context (§ 5.3) and few-shot prompting with
extra exemplars pre-pended to notebook context (§ 5.4). Here, we list the
prompts for $\bm{u}_{2}$ in Fig. 1 in these two types of prompting
experiments.

##### Prompting using Notebook Context

In the basic setup without extra few-shot exemplars, the prompt basically
consist of all the prior notebook context, including NL descriptions of schema
information and previous rounds of problems. Fig. 39 shows the complete prompt
for $\bm{u}_{2}$ in Fig. 1 in this setup.181818This prompt is not exactly the
same as the one in our dataset. It is adapted to align with the illustrative
example in Fig. 1 At inference time, a code LM will complete the last code
cell after the cell delimiter ‘# In[ ]:’. Note that for InCoder we follow
Fried et al. (2022) and use a special template to linearize notebooks
(Appendix D).

##### Prompting using Additional Few-shot Exemplars

We have four prompting styles for few-shot experiments. Here, we show the
prompt prefix (§ 5.2) for $\mathrm{Vanilla~{}Code}$ and
$\mathrm{Step}\textrm{-}\mathrm{by}\textrm{-}\mathrm{Step}$$+\mathrm{Explanations}$
prompting, as the remaining two styles are just simplified version of the
latter by removing inline explanations ($\mathrm{SbS}+\mathrm{Preamble}$) and
preambles ($\mathrm{Step}\textrm{-}\mathrm{by}\textrm{-}\mathrm{Step}$).

A prompt in this setup is the concatenation of a prompt prefix (with few-shot
exemplars) and the notebook context (with prior rounds of problems and NL
schema descriptions). The part of a prompt that corresponds to notebook
context as the same as the previous setting (_e.g._ Fig. 39), except that we
insert the preamble # Solution: Let’s solve this problem step-by-step. as
appropriate after the last cell delimiter. For prompt prefix, Fig. 33 gives an
example prompt prefix for
$\mathrm{Step}\textrm{-}\mathrm{by}\textrm{-}\mathrm{Step}$ prompting, while
Fig. 36 shows the same set of few-shot exemplars for $\mathrm{Vanilla~{}Code}$
prompting.

As mentioned in § 5.2, we created three prompt prefixes for each of the four
different styles, and report results averaged over these three restarts. Figs.
33, 34 and 35 show the three groups of prompt prefixes for
$\mathrm{Step}\textrm{-}\mathrm{by}\textrm{-}\mathrm{Step}$, and Figs. 36, 37
and 38 show those for $\mathrm{Vanilla~{}Code}$ prompting. Each prompt prefix
has four exemplars, and some exemplars are shared across different prefixes.
Note that some prompt prefixes in
$\mathrm{Step}\textrm{-}\mathrm{by}\textrm{-}\mathrm{Step}$ also contain one
simple problem that does not require decomposition and explanation (_e.g._
Exercise 3, Fig. 33). We find this to be useful to not bias a model from
generate overly complex code solutions for simpler problems. We did not put
much effort in prompting engineering. Actually, those prompt prefixes were
created before we collected $70\%$ of the dataset,

$\mathrm{Step}\textrm{-}\mathrm{by}\textrm{-}\mathrm{Step}$ Prompt Prefix
(Group 1)

⬇

\par# In[ ]:

\par\parimport pandas as pd

import matplotlib.pyplot as plt

\par\par# In[ ]:

\par\par# You are a professional data scientist. Answer the following
questions using pandas and matplotlib.

\par\par# In[ ]:

\par\par# # Exercise 1

\par\par# In[ ]:

\par\pardf = pd.read_csv(’employee.csv’)

\par\par# In[ ]:

\par\par# Schema of Dataframes:

# Columns in df with example values:

# name (Peter), gender (m), DOB (1992/01/17)

\par\par# In[ ]:

\par\par# Problem: How many male and female employees are born in 1992?

\par\par# In[ ]:

\par\par# Solution: Let’s solve this problem step-by-step.

# Step 1: convert date of birth in to datetime

df[’DOB’] = pd.to_datetime(df[’DOB’])

# Step 2: get the number of male born in 1992

num_male_students = len(df[(df[’DOB’].dt.year == 1992) & (df[’gender’] ==
’m’)])

# Step 3: get the number of female born in that year

num_female_students = len(df[(df[’DOB’].dt.year == 1992) & (df[’gender’] ==
’f’)])

\par\par# In[ ]:

\par\par# # Exercise 2

\par\par# In[ ]:

\par\pardf = pd.read_csv(’scores.csv’)

\par\par# In[ ]:

\par\par# Schema of Dataframes:

# Columns in df with example values:

# Stu_Name (Mike), Engineering (90), English (89), Math (92)

\par\par# In[ ]:

\par\par# Problem: Get the students with an averaged score above 90 for
science subjects.

\par\par# In[ ]:

\par\par# Solution: Let’s solve this problem step-by-step.

# Step 1: Create a new column with the average score of engineering and math

df[’Science_Avg’] = (df[’Engineering’] + df[’Math’]) / 2

# Step 2: Get the rows whose average score is above 90

df_score_above_90 = df[df[’Science_Avg’] > 90]

# Step 3: Return the student name and average scores

result = df_score_above_90[[’Stu_Name’, ’Science_Avg’]]

\par\par# In[ ]:

\par\par# # Exercise 3

\par\par# In[ ]:

\par\pardf = pd.read_csv(’geo.csv’)

\par\par# In[ ]:

\par\par# Schema of Dataframes:

# Columns in df with example values:

# state (WA), capital (Seattle), population (1.4 millon)

\par\par# In[ ]:

\par\par# Problem: What is the population of California?

\par\par# In[ ]:

\par\par# Solution: Let’s solve this problem step-by-step.

result = df[df[’state’] == ’CA’][’population’]

\par\par# In[ ]:

\par\par# # Exercise 4

\par\par# In[ ]:

\par\pardf = pd.read_csv(’phones.csv’)

\par\par# In[ ]:

\par\par# Schema of Dataframes:

# Columns in df with example values:

# model (Pixel 6), brand (Google), price (387), release (2022)

\par\par# In[ ]:

\par\par# Problem: What is the most expensive phone in each brand.

\par\par# In[ ]:

\par\par# Solution: Let’s solve this problem step-by-step.

# Step 1: Group models by their brands.

model_by_brand_df = df.groupby(’brand’)

# Step 2: Find the index of rows that have the highest price in each group

idx = model_by_brand_df[’price’].idxmax()

# Step 3: Get the rows using the index

expensive_models_df = df.loc[idx]

# Step 4: Return the brand name, model and price.

result = expensive_models_df[[’brand’, ’model’, ’price’]]

\par\par# In[ ]:

\par\par# # Exercise 5

\par

$\mathrm{Step}\textrm{-}\mathrm{by}\textrm{-}\mathrm{Step}$ Prompt Prefix
(Group 2)

⬇

# In[ ]:

\par\parimport pandas as pd

import matplotlib.pyplot as plt

\par\par# In[ ]:

\par\par# You are a professional data scientist. Answer the following
questions using pandas and matplotlib.

\par\par# In[ ]:

\par\par# # Exercise 1

\par\par# In[ ]:

\par\pardf = pd.read_csv(’employee.csv’)

\par\par# In[ ]:

\par\par# Schema of Dataframes:

# Columns in df with example values:

# name (Peter), gender (m), DOB (1992/01/17)

\par\par# In[ ]:

\par\par# Problem: How many male and female employees are born in 1992?

\par\par# In[ ]:

\par\par# Solution: Let’s solve this problem step-by-step.

# Step 1: convert date of birth in to datetime

df[’DOB’] = pd.to_datetime(df[’DOB’])

# Step 2: get the number of male born in 1992

num_male_students = len(df[(df[’DOB’].dt.year == 1992) & (df[’gender’] ==
’m’)])

# Step 3: get the number of female born in that year

num_female_students = len(df[(df[’DOB’].dt.year == 1992) & (df[’gender’] ==
’f’)])

\par\par# In[ ]:

\par\par# # Exercise 2

\par\par# In[ ]:

\par\pardf = pd.read_csv(’scores.csv’)

\par\par# In[ ]:

\par\par# Schema of Dataframes:

# Columns in df with example values:

# Stu_Name (Mike), Engineering (90), English (89), Math (92)

\par\par# In[ ]:

\par\par# Problem: Get the students with an averaged score above 90 for
science subjects.

\par\par# In[ ]:

\par\par# Solution: Let’s solve this problem step-by-step.

# Step 1: Create a new column with the average score of engineering and math

df[’Science_Avg’] = (df[’Engineering’] + df[’Math’]) / 2

# Step 2: Get the rows whose average score is above 90

df_score_above_90 = df[df[’Science_Avg’] > 90]

# Step 3: Return the student name and average scores

result = df_score_above_90[[’Stu_Name’, ’Science_Avg’]]

\par\par# In[ ]:

\par\par# # Exercise 3

\par\par# In[ ]:

\par\pardf = pd.read_csv(’geo.csv’)

\par\par# In[ ]:

\par\par# Schema of Dataframes:

# Columns in df with example values:

# state (WA), capital (Seattle), population (1.4 millon)

\par\par# In[ ]:

\par\par# Problem: What is the population of California?

\par\par# In[ ]:

\par\par# Solution: Let’s solve this problem step-by-step.

result = df[df[’state’] == ’CA’][’population’]

\par\par# In[ ]:

\par\par# # Exercise 4

\par\par# In[ ]:

\par\pardf = pd.read_csv(’phones.csv’)

\par\par# In[ ]:

\par\par# Schema of Dataframes:

# Columns in df with example values:

# model (Pixel 6), brand (Google), price (387), release (2022)

\par\par# In[ ]:

\par\par# Problem: What is the most expensive phone in each brand.

\par\par# In[ ]:

\par\par# Solution: Let’s solve this problem step-by-step.

# Step 1: Group models by their brands.

model_by_brand_df = df.groupby(’brand’)

# Step 2: Find the index of rows that have the highest price in each group

idx = model_by_brand_df[’price’].idxmax()

# Step 3: Get the rows using the index

expensive_models_df = df.loc[idx]

# Step 4: Return the brand name, model and price.

result = expensive_models_df[[’brand’, ’model’, ’price’]]

\par\par# In[ ]:

\par\par# # Exercise 5

$\mathrm{Step}\textrm{-}\mathrm{by}\textrm{-}\mathrm{Step}$ Prompt Prefix
(Group 3)

⬇

# In[ ]:

\par\parimport pandas as pd

import matplotlib.pyplot as plt

\par\par# In[ ]:

\par\par# You are a professional data scientist. Answer the following
questions using pandas and matplotlib.

\par\par# In[ ]:

\par\par# # Exercise 1

\par\par# In[ ]:

\par\pardf = pd.read_csv(’olympics.csv’)

\par\par# In[ ]:

\par\par# Schema of Dataframes:

# Columns in df with example values:

# Year (1896), City (Athens), Country (Greece), Nations (14)

\par\par# In[ ]:

\par\par# Problem: Which countries host at least two olympic games?

\par\par# In[ ]:

\par\par# Solution: Let’s solve this problem step-by-step.

# Step 1: Count the number of times each country hosted olympics

count_df = df[’Country’].value_counts()

# Step 2: Find entries with more than 2 counts

filtered_df = count_df[count_df >= 2]

# Step 3: Get the country names as a list

filtered_df.index.tolist()

\par\par# In[ ]:

\par\par# # Exercise 2

\par\par# In[ ]:

\par\pardf = pd.read_csv(’employee.csv’)

\par\par# In[ ]:

\par\par# Schema of Dataframes:

# Columns in df with example values:

# name (Peter), gender (m), DOB (1992/01/17)

\par\par# In[ ]:

\par\par# Problem: How many male and female employees are born in 1992?

\par\par# In[ ]:

\par\par# Solution: Let’s solve this problem step-by-step.

# Step 1: convert date of birth in to datetime

df[’DOB’] = pd.to_datetime(df[’DOB’])

# Step 2: get the number of male born in 1992

num_male_students = len(df[(df[’DOB’].dt.year == 1992) & (df[’gender’] ==
’m’)])

# Step 3: get the number of female born in that year

num_female_students = len(df[(df[’DOB’].dt.year == 1992) & (df[’gender’] ==
’f’)])

\par\par# In[ ]:

\par\par# # Exercise 3

\par\par# In[ ]:

\par\pardf = pd.read_csv(’score.csv’)

\par\par# In[ ]:

\par\par# Schema of Dataframes:

# Columns in df with example values:

# name (John), score (97)

\par\par# In[ ]:

\par\par# Problem: Make a new column "grade" for letter grades (A: 90+, B:
70-90, C: <70) and plot the number of students in each grade.

\par\par# In[ ]:

\par\par# Solution: Let’s solve this problem step-by-step.

# Step 1: Define a function to convert scores to letter grades.

def get_grade(score):

if score >= 90:

return ’A’

elif 70 <= score < 90:

return ’B’

else:

return ’C’

# Step 2: Convert scores to letter grades.

df[’grade’] = df.score.apply(get_grade)

# Step 3: Count the number of students by grade.

count_df = df[’grade’].value_counts()

# Step 4: Visualize in a bar chart.

count_df.plot(kind=’bar’)

\par\par# In[ ]:

\par\par# # Exercise 4

\par\par# In[ ]:

\par\pardf = pd.read_csv(’phones.csv’)

\par\par# In[ ]:

\par\par# Schema of Dataframes:

# Columns in df with example values:

# model (Pixel 6), brand (Google), price (387), release (2022)

\par\par# In[ ]:

\par\par# Problem: What is the most expensive phone in each brand.

\par\par# In[ ]:

\par\par# Solution: Let’s solve this problem step-by-step.

# Step 1: Group models by their brands.

model_by_brand_df = df.groupby(’brand’)

# Step 2: Find the index of rows that have the highest price in each group

idx = model_by_brand_df[’price’].idxmax()

# Step 3: Get the rows using the index

expensive_models_df = df.loc[idx]

# Step 4: Return the brand name, model and price.

result = expensive_models_df[[’brand’, ’model’, ’price’]]

\par\par# In[ ]:

\par\par# # Exercise 5

$\mathrm{Vanilla~{}Code}$ Prompt Prefix (Setup 1)

⬇

# In[ ]:

\par\parimport pandas as pd

import matplotlib.pyplot as plt

\par\par# In[ ]:

\par\par# You are a professional data scientist. Answer the following
questions using pandas and matplotlib.

\par\par# In[ ]:

\par\par# # Exercise 1

\par\par# In[ ]: