Natural Language to Code Generation in Interactive Data Science Notebooks (2023)

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 1078 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. Arcade is publicly available at https://github.com/google-research/arcade-nl2code/ .

## introduction data science is the process of extracting insights from data @xcite , and has become an integral part of decision making and knowledge discovery @xcite . data scientists and machine learning (ml) practitioners often use computational notebooks, which are interactive environments such as jupyter notebooks @xcite and google colab, in their work. data scientists spend a significant amount of time on data wrangling tasks to process raw data into usable forms (illustrated in fig. 1 ), as well as exploratory data analysis (eda) to gain insights * correspondence to pcyin@google.com in which year was the most played game added? @xcite df @xcite .dt.@xmath0.year) & (df['rating']>4)].groupby( df["added"].dt.month)['game'].count() for each month in that year, how many games that has a rating of more than four? @xcite @xmath1.str.contains('fallout')] fallout.groupby(fallout['added'].dt.year).get_group( 2021) @xcite .mean() what is the average maximum completion time for all fallout games added in 2021? what is the amount of games added in each year for each month? (show a table with index as years, columns as months and fill null values with 0) for decision making @xcite . this has motivated research on automating and accelerating the data science workflow in general @xcite , with particular interest in data wrangling and eda tasks @xcite @xcite . meanwhile, large language models (llms) trained on code can assist developers by translating natural language (nl) intents into executable programs @xcite @xcite @xcite , with promising applications in synthesizing code for data wrangling and eda tasks @xcite @xcite . computational notebooks also present unique challenges to llms, as notebooks freely mix nl, code, graphics, and execution results @xcite , and because of their interactivity, notebooks feature multiple interdependent nl-to-code problems @xcite . 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 @xcite , which tend to contain nl text (e.g., exercise questions) that is verbose and elaborate, instead of the concise, ephemeral style that developers write when interacting with code lms @xcite . other datasets assume that the developer provides extra information, such as unit tests or input/output examples @xcite , but such systems pose an extra burden to users who might not normally write such tests or examples during their workflow @xcite . finally, existing datasets usually contain independent tasks with isolated contexts @xcite , or a limited number of contextually dependent problems @xcite , rather than having multiple, related tasks such as in fig. 1 . therefore, there is a need for a benchmark with realistic nl intents, rich notebook context, and a series of interrelated problems, so as to better reflect real-world usage by data scientists. to fill this gap, we present arcade, foot_0 a new benchmark for code generation for data wrangling and eda tasks in computational notebooks ( §3). arcade consists of 1,078 problems spanning across 136 notebooks based on 106 ml datasets. it features a series of nl utterances written by professional data scientists with the intention of interacting with an ai assistant (e.g., green texts in fig. 1 ), with high-quality code solutions using the pandas library. to mitigate the risk of data leakage, 60% of the problems are created from scratch, based on recent ml datasets on kaggle (e.g., the csv file in c 1 , fig. 1 ). foot_1 arcade also challenges llms with grounded language understanding, where a model needs to leverage variable states (e.g., df @xcite in c 2 ) to interpret nl semantics (e.g., "min and max" in u 1 ). finally, problems in arcade are challenging, involving richer data science api usage than existing benchmarks. to demonstrate how arcade can motivate new research on llms for data science, we develop pachinco, a 62b code lm tailored for python computational notebooks, trained on a mixture of nl, source code, and jupyter notebooks data ( §4). pachinco significantly outperforms public code lms on arcade ( §5.2). even so, all models have difficulty on our benchmark, showing that it is a challenging task. further, we explore few-shot prompting strategies to alter the style of model predictions, such as decomposing code into step-bystep structures and adding inline nl explanations. not only is code in this style potentially more understandable to novice data scientists, prompting the model to explain its solutions also improves the diversity of the model's predictions ( §5.3). ## problem statement a computational notebook is an interactive computing environment that allows mixing code, text, and graphics. a notebook consists of a sequence of markdown or source code cells. given a partial notebook context with n cells {c i } n @xmath2.g., u 1 in fig. 1 for @xmath3, we aim to generate code for c n+1 that fulfills the user's intent @xcite . we refer to the pair ({c i }, u) as a problem. this process could proceed sequentially with multiple rounds between the user and a system @xcite , so a single notebook can contain multiple problems. to satisfy subsequent intents (e.g., u 4 ), a system will leverage the updated notebook context (e.g., {c i } 5 as in fig. 1 , problems within a notebook often have interesting dependency structures. they may share execution context (e.g., dataframe df), form semantically coherent turns (e.g., c 4 and c 5 ), or exhibit non-trivial long range data dependencies (e.g., from c 6 to c 2 , or c 7 to c 3 ). these dependency structures are more diverse than existing multi-turn code generation tasks with sequentially dependent problems @xcite . ## pachinco: adapting code lms to computational notebooks we introduce pachinco, an lm for notebooks. base lm pachinco is based on palm, a family of decoder-only lms for nl tasks @xcite . specifically, we use the 62b palm model trained on 1.3t tokens with a mixture of conversational, webpages and code data (section f, chowdhery et al. (2022)). starting with this base lm, we first fine-tune on python source code and then fine-tune further on jupyter notebooks. fine-tuning on python code we first fine-tune the base lm on a corpus of near-deduplicated, permissively-licensed python source code files from github, with 64b tokens in total. we finetune palm for 1 epoch following the hyper parameters setup in chowdhery et al. (2022). this model is already a strong code lm, even outperforming the larger code lm palm-coder 540b on existing program synthesis benchmarks ( §5.2). fine-tuning on notebooks we then perform a second stage of fine-tuning on a large collection of 3.8m jupyter notebooks from github (9.6b tokens). since our evaluation notebooks in the existing tasks split are also from github, we also perform near-deduplication to remove any training notebooks with one cell similiar to any cells in the notebooks in existing tasks to prevent data contamination. we use nbconvert to linearize notebooks into python code. refer to appendix d for details and appendix k for a data card. 5 experiments models we evaluate pachinco and state-ofthe-art public code llms, namely codegen (nijkamp et al., 2022) and incoder (fried et al., 2022). we test both the monolingual (python-only) and the multilingual version of codegen. in-coder may be a more appealing comparison since it is trained on 5gb of jupyter notebooks. inference and metrics we convert each problem into a prompt ( §5.1) and draw samples using nucleus sampling. following chen et al. (2021a), we report pass@k metric, defined as the fraction 2] # schema of dataframes: # columns in df with example values: # stu_name (mike),engineering (90),english (89),math (92) [3] get the students with an average score above 90 for science subjects ## case study: how useful is predicted code with step-wise explanations? finally, we remark that besides improving solution diversity, step-by-step prompting with nl explanations could also potentially help novice data scientists understand model-generated solutions, as shown in the following qualitative case study. first, 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. 17 ), or the underlying schema is less intelligible (e.g., column names with abbreviations, fig. 18 ). additionally, some inline explanations also describe the output of intermediate steps, which is particularly helpful when these steps involve advanced pandas functions whose output structure may not be obvious, such as pd.unstack (fig. 19 ) meanwhile, step-wise nl explanations serve as high-level procedural descriptions of code, which enable users to easily browse through and understand different solution approaches without being distracted by nuances in the actual code implementation (fig. 20 ). moreover, explanations also help users verify the code solutions by identifying potentially incorrect steps (fig. 21 ). 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. ## related work automating data science the amount of expertise required in data science has called for development of systems to automate its lifecycle @xcite . much work has focused on automating feature engineering and tuning of ml models (automl, @xcite , with well-established systems @xcite and benchmarks @xcite . this paper focuses on automating tabular data wrangling and eda tasks, which account for nearly the same amount of code and documentations in notebooks as that for ml-related tasks @xcite . along this line, existing research synthesizes data wrangling programs using i/o examples @xcite or partial table contents @xcite , followed by recent efforts using llms with additional nl specifications @xcite . this paper considers code generation in notebooks with multiple contextually dependent problems (see §3.2 for recent work). in addition, other works have also considered applications such as synthesizing visualization plots @xcite @xcite @xcite . context-driven code generation our work is another application of context-driven code generation, which maps a series of contextually dependent utterances to programs, such as domain-specific logical forms @xcite @xcite , sql queries over databases @xcite @xcite , or generalpurpose pls @xcite . arcade further offers contextually dependent utterances exhibiting non-trivial dependencies ( §2), with target programs defined in a general-purpose pl. ## conclusion in this paper we present arcade, a code generation benchmark for data wrangling and eda tasks in computational notebooks, featuring problems with realistic nl intents and rich contexts. we also develop pachinco, a 62b lm tailored for data science. pachinco outperforms public lms on arcade, while being effective in few-shot learning to improve code style and solution diversity.