When Language Model Meets Private Library
Daoguang Zan, Bei Chen, Zeqi Lin, Bei Guan, Wang Yongji, Jian–Guang Lou · 2022
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 Torch-DataEval, MonkeyEval, and BeatNumEval.Experimental results demonstrate the impressive performance of our framework.1 .