When Language Model Meets Private Library (2022)

Abstract: With the rapid development of pre-training techniques, a number of language models have been pre-trained on large-scale code corpora and perform well in code generation. In this paper, we investigate how to equip pre-trained language models with the ability of code generation for private libraries. In practice, it is common for programmers to write code using private libraries. However, this is a challenge for language models since they have never seen private APIs during training. Motivated by the fact that private libraries usually come with elaborate API documentation, we propose a novel framework with two modules: the APIRetriever finds useful APIs, and then the APICoder generates code using these APIs. For APIRetriever, we present a dense retrieval system and also design a friendly interaction to involve uses. For APICoder, we can directly use off-the-shelf language models, or continually pre-train the base model on a code corpus containing API information. Both modules are trained with data from public libraries and can be generalized to private ones. Furthermore, we craft three benchmarks for private libraries, named TorchDataEval, MonkeyEval, and BeatNumEval. Experimental results demonstrate the impressive performance of our framework.

## introduction code generation, automatically generating code snippets based on user descriptions, is one of the long-standing challenges in the software engineering and artificial intelligence communities. with the rapid development of pre-training techniques, a number of language models are pre-trained on large-scale code corpora and able to generate decent code snippets, for example, codex @xcite , alphacode @xcite , codegen @xcite , and incoder @xcite . although making remarkable progress, these models may be biased towards generating code that is similar to the training distribution @xcite . what if one wants to generate code beyond the training distribution? a real-world scenario for programmers is to write code using a private library, which is very common in practice. for example, for security and functionality reasons, companies often build private libraries for internal use only. private libraries provide a number of private apis that have not been seen by the language models and are also not publicly available on any code hosting platform like github. therefore, it is worth exploring whether and how pre-trained language models can generate code using private libraries. it is challenging for existing language models to generate code that uses a private library directly. a practical evidence is shown in figure 1 . we built a pseudo private library named monkey based on a public one named pandas. pandaseval @xcite is a benchmark consisting of 101 pandas programming problems. we convert all pandasrelated keywords in pandaseval into the new version and construct monkeyeval (details in section 4). as seen in figure 1 , codex 12b and codegen-mono 350m show a significant drop in performance on the private monkeyeval compared to their performance on the public pandaseval. for example, codex 12b drops from 18.88% to 1.47% on pass@1, showing the inadequacy of the language models in code generation for private libraries. to meet the challenge, we propose a framework to equip pre-trained language models with the ability to generate code that uses private libraries. as is known, private libraries usually come with elaborate api documentation, which motivates our main idea to mimic the process of a programmer learning to write code using private libraries. this process is also known as api practices in the software engineering field @xcite : first learning the private api documentation and then invoking the apis to implement the needed functionalities. analogically, there are two modules in our framework: an apiretriever first retrieves the useful apis based on the programming problem and the api documentation, and then an apicoder uses these apis to generate code. for apiretriever, we train a dense retriever and also design a friendly interaction to involve users in the loop optionally. for apicoder, we can directly use existing language models of code generation, such as code-gen, to invoke the private apis; furthermore, to better teach a language model how to invoke apis, we also continually pre-train the base model on a code corpus containing api information from public libraries, and obtain our reinforced model called codegenapi. since we only have access to the data of public libraries during training, we expect that apiretriever and apicoder can be generalized to private libraries via learning. to evaluate the code generation for private libraries, we craft three benchmarks, named torch-dataeval, monkeyeval, and beatnumeval. torch-dataeval includes 50 programming problems using the torchdata library. the last two are adapted from pandaseval and numpyeval @xcite , respectively, each consisting of 101 programming problems. extensive experiments on the three benchmarks have revealed that our framework effectively improves the performance of pre-trained language models on code generation for private libraries. we also provide a thorough analysis to facilitate progress in this direction. ## framework first, we would like to define the task of code generation formally. given context, the task aims to generate target code. in figure 1 , context and target code are shown in white and grey backgrounds, respectively. context consists of a comment, which is a natural language description of the programming problem, and a code snippet including import statements, function header, etc. target code solves the programming problem in context. we denote the context by x. code generation model m outputs target code y based on x. for the task of code generation for private library, the context x contains the instruction for using a private library, such as an import statement. the target code y contains the calls of the corresponding private library apis. as mentioned in section 1, private libraries are usually equipped with elaborate api documentation. as a technical reference manual outlining how to use the library, api documentation typically includes a quick start guide, tutorials, and an instruction for each api (e.g., api name, signature, description, parameters, and examples). to take advantage of the api documentation, we propose to mimic the generic process of a programmer coding with private apis, and design a framework to generate code that can invoke private apis. the framework consists of apiretriever and apicoder with the overview shown in figure 2 . given the context, apiretriever m r aims to retrieve possible used apis from the api documentation; and apicoder m c is dedicated to generating code using the retrieved apis. the process can be formalized as @xmath0; x), where a represents the set of information of all proper apis, and each @xmath1. in our implementation, we design the api information to include the api name, signature and description. note that we only use the first sentence of the api ## methodology we have introduced our framework that provides pre-trained models a fantastic way to deal with private libraries. in this section, we present the data collection, followed by the detailed design of our apiretriever and apicoder. ## experiments in this section, we conduct experiments to illustrate the superiority of our proposed framework. ## conclusion in this paper, we propose a novel framework for code generation for private library. there are two modules: for a specific programming problem, apiretriever first finds out the useful private apis from api documentation, and then apicoder leverages these apis to generate the code. we craft three benchmarks, including torchdataeval, monkeyeval, and beatnumeval, for better evaluating private library oriented code generation. the experimental results and thorough analysis demonstrate the reasonableness and effectiveness of our framework. in future work, we would like to explore how to make better use of api documentation for code generation and improve the approach for real use when programming with private libraries.