CERT: Continual Pre-training on Sketches for Library-oriented Code Generation

1 Introduction

Code generation, aiming to generate a code snippet for a given natural language description, is a longstanding challenge in the artificial intelligence community. Usually, to train a code generation model with good performance, the massive amount of code snippets paired with natural language descriptions are indispensable [Sun et al., 2019; Lu et al., 2021]. However, it is costly and time-consuming to annotate such a dataset. To alleviate this problem, inspired by GPT-3's powerful zero-shot natural language generation ability [Brown et al., 2021], recent years have witnessed a trend to train large language models using large-scale code corpora (e.g., GitHub), and expect these models to work well directly on code generation tasks, without fine-tuning on expensive text-code pairs. For example, Codex shows that a 12B parameters language model can solve 28.8% of standalone Python programming problems 1 .

∗Work done during the internship at Microsoft Research Asia.

In this paper, we focus on investigating whether and how language models pre-trained on code corpora (without finetuned on pairwise labelled data) can generate library-oriented code snippets rather than standalone ones. During software development, it is a common practice for programmers to reuse third-party libraries (e.g., Pandas and NumPy) to implement needed functionalities. It is not easy for programmers to learn how to use these libraries properly. For example, according to our statistics, more than 40% of StackOverflow questions with "Python" tag also have at least one library tag. Moreover, for library-oriented code generation, the necessity of training the model without pairwise labelled data is raised, as programmers usually need to reuse different libraries in different scenarios, and it is extremely costly to label sufficient text-code pairs that cover most of these libraries.

1 It is measured on HumanEval [Chen et al., 2021a] with pass@1.

Compared to standalone code snippets, library-oriented code snippets are more likely to share similar sketches. Sketch is the code structure after anonymizing the userdefined terms in the code, such as variable names, method names, constants, etc., which has also been identified as an API usage pattern in previous research litterateurs on software data mining [Zhong et al., 2009; Wang et al., 2013; Niu et al., 2017]. An example is shown in Figure 1. After anonymizing variables and constants, multiple code snippets using the Pandas APIs may have the same (or similar) sketch. Based on this observation, a natural idea to improve libraryoriented code generation is to decompose this task into two subtasks: generating the sketch and then filling in the details. Many methods based on this idea have been proposed in different code generation tasks (e.g., Coarse-to-Fine [Dong and Lapata, 2018] and PLOTCODER [Chen et al., 2021b]) and have shown that this idea can effectively improve the quality of generated code snippets. However, these methods are proposed for the fine-tuning process, in which high-quality textcode pairs are required to derive supervision signals for the two-step generation. Therefore, in our scenario that no pairwise labelled data is provided, a research question arises: how to leverage the insight of sketching to enhance the language model pre-training on unlabelled code corpora, thus improving the quality of generated library-oriented code snippets?

To meet the challenge, we propose CERT (for sketCher and gEneRaTor), a continual pre-training approach on sketches for library-oriented code generation. In CERT, a sketcher firstly focuses on predicting a sketch, which omits user-defined details; then, a generator uses the sketch as a prompt to generate the complete code. Both the sketcher and the generator are continually pre-trained based on a base language model for code, using unlabelled code corpora rather than pairwise labelled data. In addition, we craft two evaluation benchmarks for Python libraries, called PandasEval and NumpyEval, each including 101 programming problems using Pandas and NumPy, respectively. We perform extensive experiments on CERT. Results indicate that CERT has superior performance on library-oriented code generation. We further draw several insights via thorough analysis.

2 Task Formulation

Before diving into the details of our proposed approach, we start with a formal description of the task. Code generation is

import pandas as pd def normalizeſdfƀ: # Normalization using pandas # We simply subtract the mean and divide by standard deviation on df.iloc[:,0,-ɨ] obj with axis is zero. # Return the normalized dataframe func_ = lambda x: ſx-x.meanſƀƀ / x.stdſƀ df.iloc[:,0:-ɨ] = df.iloc[:,0:-ɨ].applyſfunc_, axis=0ƀ return df import numpy as np x = np.arrayſ[[ɨ], [ɩ], [ɪ]]ƀ # Numpy Vector ſN,ɨƀ dimension -> ſN,ƀ dimension conversion out = x.reshapeſɪ,ƀ

to solve a programming problem: generate target code based on context. Context contains natural language problem description in the form of code comments, and a code snippet that includes statements such as import, function header and variable definition; target code is a code snippet that solves the programming problem described in the context. Formally, let x = (x1, x2, · · ·, xN) denote the context, where each xn can be either a code token or a natural language token. Given x, the code generation model can be formulated as y = M(x), where y = (y1, y2, · · ·, yM) denotes the target code and each ym is a code token.

For standalone code generation, the programming problem is expected to be solved by a code snippet without using third-party libraries; conversely, for library-oriented code generation, the target code y contains library API calls. Two examples of library-oriented programming problems can be found in Figure 2. Note that carefully labelled context and target code pairs are indispensable for model fine-tuning, while our proposed approach only requires continual pre-training on unlabelled code corpora.

3 Methodology

In this section, we introduce our base models, followed by the details of our proposed approach CERT.

3.1 Base Models

Codex [Chen et al., 2021a] is a milestone pre-trained model that can generate decent code, but it is not publicly available. Several attempts have been made to reproduce Codex's powerful code generation capability, e.g., CodeClippy 2 and CodeParrot 3 , but their performance in Python are not satisfactory. To this end, we present PYCODEGPT, a pre-trained language model, which has the ability to generate pretty good standalone Python code, for example, achieving 8.33% pass@1 on HumanEval [Chen et al., 2021a]. Specially, PYCODEGPT is a 110M parameters model based on GPT-Neo [Black et al., 2021]. We collected 60.6M raw python files with a total size of 330GB. After a series of data pre-processing strategies, such as de-duplicating python files, cleaning and formatting the contents, etc., the final pre-training corpus contains about 13.0M high-quality python files with the size of 96GB. PYCODEGPT is pre-trained for 200K steps and 100B

2 https://github.com/CodedotAl/gpt-code-clippy

3 https://huggingface.co/transformersbook/codeparrot

tokens on a cluster of 16 NVIDIA V100 GPUs with 32GB memory. The pre-training time is about 2 days. We summarize the three key points that make PYCODEGPT powerful: 1) a large amount of carefully cleaned data for pre-training; 2) a newly trained tokenizer, which is specialized in python; and 3) a resampling strategy that prioritizes high-quality data. Besides PYCODEGPT, we also regard CODEGEN (MONO 350M) [Nijkamp et al., 2022] as one of our base models, which is by far the best performing publicly available model on HumanEval 4 .

3.2 CERT

As mentioned in Section 2, code generation is to generate target code y based on context x. Since we observe that library-oriented code snippets are more likely to share similar sketches, we present a novel approach CERT and decompose the code generation model M into two modules: a sketcher MS and a generator MG. Figure 3 shows the overview of CERT with a concrete example in Pandas. Given x as the input, the sketcher predicts s, which is the sketch of the target code y. The sketcher generates multiple candidate sketches (200 in our experiments) and we choose the one that appears the most. Then, the input of the generator is the concatenation of s and x. Formally, the process of CERT can be written as s = MS(x) and y = MG([s; x]). Note that if the sketch s is already a complete code snippet without anonymous symbols, we directly take it as the final prediction instead of using the generator; and if the sketch s is an empty sequence, we directly feed x into the generator.

We build the sketcher and generator on the top of the base model (PYCODEGPT or CODEGEN) by continual pretraining. At first, we extract the python files that use a specific library (e.g., Pandas) from the whole pre-training corpus (13.0M files mentioned in Section 3.1), and obtain the subcorpus denoted by D. Then, we will detail the continual pretraining process of the sketcher and generator for this library.

4 CODEGEN was released during the review period of this paper.

Sketcher. Given the library-oriented sub-corpus D, we perform the sketching operation on each file d ∈ D. An example is shown in the upper part of Figure 5. The sketching operation is used to anonymize the user-defined terms in the code file with our pre-defined symbols. The file after sketching is denoted as ¯d. We design three different types of sketching operations: 1) only anonymizing the user-defined constants (Default CERT); 2) only anonymizing the user-defined names, including function names, class names, and variable names (CERT-N); and 3) anonymizing both the user-defined constants and names (CERT-NC). For example, in Figure 5, the constant 'user 1' is anonymized with the pre-defined symbol 'string'. The details of pre-defined symbols are shown in Figure 4. Then, we continually pre-train the base model on the library-oriented corpus after sketching, and we obtain the sketcher model. The pre-training objective is the same as that of the base model. We pre-train the model for 100K steps on a cluster of 8 NVIDIA V100 GPUs with 32GB memory.

Generator. In order to prepare the pre-trained corpus for the generator, we firstly split the original file d and the sketching file ¯d into K blocks 5 , and obtain d = (d1, d2, · · ·, dK) and ¯d = ( ¯ d1, ¯ d2, · · ·, ¯ dK). Each block is a relatively complete code snippet, such as a function or a class. Note that before splitting, we remove the natural language code comments from the sketching file ¯d. Then, the two files are cross-merged to give a merged file ˆd = ( ¯ d1, d1, ¯ d2, d2, · · ·, ¯ dK, dK). This is to mimic the process of having a sketch as a prompt for each block. An example is shown in the lower part of Figure 5. Then, the base model is continually pre-trained on all the merged files ˆd and we obtain the generator model. As with the sketcher model, we continually pre-train for 100K steps.

4 Benchmark Construction

Third-party libraries are widely used in reality, while little work has been done to evaluate library-oriented code generation. To meet this challenge, we craft PandasEval and NumpyEval, two benchmarks for library-oriented code generation in Python. Each sample in the benchmarks is a programming problem consisting of context and target code. The programming problems are solved using libraries, where Pandas is for PandasEval, and NumPy is for NumpyEval. The benchmarks are expected to be diverse, authentic, high quality, moderately difficult, and unseen during pre-training.

In order to craft programming problems using libraries, we refer to StackOverflow 6 , a Q&A website for programmers. There are plenty of real-world programming problems posted by real users, which helps us to improve the authenticity of our data. Specifically, we search for posts using the library tag on StackOverflow, and select those with high votes. To ensure quality, we only refer to posts with accepted answers. We go through a post's question and its accepted answer, then manually organize them into the form needed for our benchmarks, containing both context and target code. We also polish all programming problems so that the problem descriptions are clear and the codes are correct. Note that we keep the intentions and the descriptions of the programming problems consistent with the posts to the maximum extent. Finally, two programmers with more than three years of coding experience in the library are invited to act as code generation models and check the quality of the data.

5 We use pip-tools: autopep8, docformatter and redbaron.

As a result, we craft 101 programming problems for PandasEval and NumpyEval, respectively. Each programming problem is equipped with test cases for evaluation. For the programming problems in the form of a function, such as the bottom one in Figure 2, we create 20 test cases for each of them. For the others that contain no functions, such as the top one in the Figure 2, we provide 1 test case to check the correctness of predicted variable (e.g., out in Figure 2). In total, 64% programming problems in PandasEval and 30% in NumpyEval are equipped with 20 test cases. In addition, we craft programming problems that refer to StackOverflow rather than GitHub, and also carefully organize and polish the problems, so that we can ensure they are unseen by the pretrained models.

5 Experiments

In this section, we evaluate CERT on PandasEval and NumpyEval to verify its effectiveness.

EvaluationMetrics. We use pass@k as the metrics. When k code samples are generated per problem, pass@k indicates the fraction of correct ones. But computing pass@k in this way may have high variance. Hence, we follow Chen et al. [2021a] to generate n ≥ k code samples per problem

6 https://stackoverflow.com

(n = 200 in our experiments) and count the number of correct samples c. If n - c < k, then pass@k = 1; otherwise, pass@k = 1 - ∏ni=n-c+1 (1 - k/i). Note that a predicted code is correct if it can pass all the test cases.

Implementation Details. We implement our approach using PyTorch [Paszke et al., 2019], HuggingFace's transformers library [Wolf et al., 2019], and DeepSpeed 7 . In the training phase of PYCODEGPT, we set the batch size to 10, the window size to 1024, the learning rate to 5e-4, the gradient accumulation steps to 4 and the weight decay to 0.1. The settings of sketcher and generator are the same as PYCODEGPT. We use the mixed-precision of FP16 to accelerate the pre-training. In inference phase, we set the temperature to one of [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]. The best performance is reported across the above hyper-parameters.

5.1 Main Results

Before evaluating CERT, we would like to evaluate our base model PYCODEGPT on HumanEval [Chen et al., 2021a] compared to several advanced pre-trained models. As shown in Table 1, PYCODEGPT (110M) achieves competitive 8.33% pass@1. It largely exceeds other models with comparable parameters, e.g., AlphaCode (89M) [Li et al., 2022], CodeClippy (125M), and CodeParrot (110M), and also is better than the larger model GPT-Neo (1.3B).

Then, our proposed CERT is evaluated on PandasEval and NumpyEval. We train CERT on two base models, including PYCODEGPT and CODEGEN, named PYCODEGPTCERT and CODEGEN-CERT, respectively. For each benchmark, we extract corresponding library-oriented files to train CERT. The file numbers are about 0.61M for Pandas and 2.62M for NumPy. Baselines include our base models PYCODEGPT and CODEGEN; PYCODEGPT-XL and CODEGEN-XL, which are continual pre-trained PYCODEGPT and CODEGEN on the extracted library-oriented files; and advanced pre-trained models for code, like CodeT5 [Wang et al., 2021], CodeGPT [Lu et al., 2021], CodeClippy and CodeParrot. Table 2 summarizes the performance. CERT consistently outperforms all the baselines by a large margin. The absolute improvements over PYCODEGPT and CODEGEN are shown in red, which are significant, for example, 12.69% pass@1 for CODEGEN-CERT and 13.43% pass@1 for PYCODEGPT-CERT on NumpyEval. The results demonstrate the effectiveness of CERT with the idea of leveraging sketches for library-oriented code generation.

7 https://github.com/microsoft/DeepSpeed

Additionally, we would like to investigate the performance of CERT with respect to the number of API calls involved in the target code. We divided the programming problems in each benchmark into four parts based on the number of APIs. As shown in Figure 6, compared to PYCODEGPT, PYCODEGPT-CERT has a steady improvement on each part. It indicates that CERT can improve the performance of libraryoriented code generation of varying difficulties.

5.2 Closer Analysis

We conduct some closer analyses to provide more insights.

Different Types of Sketching. As mentioned in Section 3.2, we propose three types of sketching operations. By default, CERT only anonymizes user-defined constants. The other two types include CERT-N, which anonymizes only user-defined names, and CERT-NC, which anonymizes both user-defined constants and names. As shown in Table 2, CERT with default setting achieves the best performance. This observation may be related to the inherent characteristics of Pandas and NumPy. They are commonly used in data statistics and analysis, often involving manipulation of the data constants. Thus, it is necessary to anonymize userdefined constants. Anonymizing both user-defined constants and names would probably make the sketches too abstract.

Quality of Generated Sketches. Intuitively, it is easier to generate a sketch than a complete code. Thus, we would like to evaluate the quality of sketches generated by the sketcher of CERT. We use exact match accuracy as the metric and include PYCODEGPT and PYCODEGPT-XL for comparison. For PYCODEGPT and PYCODEGPT-XL, we anonymize the user-defined constants in the predicted code to obtain the sketch. As shown in Figure 7, our sketcher surpasses baselines by 15.20% and 14.21% on PandasEval and NumpyEval, respectively. It indicates that the sketcher can generate highquality sketches, and such sketches further benefit the generator. Additionally, the generator does not necessarily require an exactly correct sketch, as the sketch is just a prompt (A case will be discussed in Section 5.3).

CERT for General Code Generation. Technically speaking, CERT can also be used for general code generation tasks, not just the library-oriented ones. Concretely, following the procedure in Figure 5, we can continually pretrain PYCODEGPT using the whole 13.0M python corpus instead of the extracted library-oriented files, and obtain the model we called CERTg. We evaluate CERTg for general code generation on HumanEval compared to the base model PYCODEGPT. As shown in Table 3, they have similar pass@k results. This observation verifies our assumption that library-oriented code snippets are more likely to share similar sketches, so it is beneficial to use sketches as prompts in this situation. But in the general case, it is not useful. Meanwhile the results of CERTg on PandasEval and NumpyEval are in Table 2. CERTg is inferior to CERT, suggesting that extracting library-oriented files is essential for CERT to learn the knowledge of library-oriented sketches.

Evaluation of GPT-3 and Codex. We evaluate GPT-3 and Codex to see how these extremely large models perform on PandasEval and NumpyEval. As shown in Table 4, CERT is competitive with only 110M parameters. Such observation proves CERT's powerful code generation capability in library-oriented programming problems.

5.3 Case Study

For a more comprehensive comparison, we show three cases in Figure 8. We show in turn the context, the golden target code, the predicted code of PYCODEGPT and PYCODEGPT-XL, the sketch generated by PYCODEGPTCERT and the predicted code of PYCODEGPT-CERT. Case 1 is from PandasEval, both PYCODEGPT-CERT's sketcher and generator reach the correct results, while the baselines do not. It reveals that sketcher and generator can work well together. Case 2 is from NumpyEval, the sketcher predicts the correct sketch, which has no anonymous symbols, then this sketch is the final predicted code. It indicates that the sketcher has the ability to predict code without user-defined constants. At last, in Case 3, the sketcher makes a wrong prediction pd.Series([[number*2]*3]), while the correct sketch is pd.Series([[number*2], number, [[number*2], number]*2]). But PYCODEGPT-CERT's generator rectifies it and finally generates the correct code. Since the sketch acts only as a prompt, it is not necessarily to be perfectly correct, which endows the generator with solid robustness.

6 RelatedWork

The most related work is the line of large pre-trained models for code. As for the encoder-style pre-trained models, they cannot be employed directly to generate code, such as CuBERT [Kanade et al., 2020], CodeBERT [Feng et al., 2020], and GraphCodeBERT [Guo et al., 2020]. As for the decoder-style or encoder-decoder-style ones, they are trained on large unlabelled code corpora and can work directly on code generation task, such as CodeT5 [Wang et al., 2021], CodeGPT [Lu et al., 2021], PLBART [Ahmad et al., 2021], PolyCoder [Xu et al., 2022], CODEGEN [Nijkamp et al., 2022], AlphaCode [Li et al., 2022], and Codex [Chen et al., 2021a]. All of them focus on generating standalone code, while we investigate library-oriented code generation. Also, similar to our idea, there are several works leveraging code sketches, for example, Coarse-to-Fine [Dong and Lapata, 2018], BAYOU [Murali et al., 2018], SKETCHADAPT [Nye et al., 2019], and PLOTCODER [Chen et al., 2021b]. However, they require labelled text-code paired data for finetuning, while our models continually pre-train on unlabelled code corpora. For code generation benchmarks, there are few works, including APPS [Hendrycks et al., 2021], HumanEval [Chen et al., 2021a], and PlotCoder's dataset [Chen et al., 2021b]. The former two ones focus on evaluating the capability of generating standalone code, and the last one is primarily devoted to generating plotting APIs and visualization code. PandasEval and NumpyEval are dedicated to evaluating the performance of library-oriented code generation.

7 Conclusion

In this paper, we propose a novel approach CERT for libraryoriented code generation. It leverages the code sketches and consists of a sketcher and a generator. The sketcher and generator are continually pre-trained upon a base model using unlabelled code corpora. Also, we carefully craft two benchmarks to evaluate library-oriented code generation, namely PandasEval and NumpyEval. Experimental results and thorough analysis show the effectiveness of CERT. In future work, we are interested in code generation for private libraries with fewer data.