Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

When Language Model Meets Private Library

Daoguang Zan, Bei Chen, Zeqi Lin, Bei Guan, Yongji Wang, Jian-Guang Lou

TL;DR

Private libraries present unseen APIs that challenge contemporary code-generation models. The authors introduce a retrieval-then-generation framework (APIRetriever and APICoder) that leverages API documentation to enable code synthesis with private libraries, trained on public-library data to generalize. They also create three benchmarks—TorchDataEval, MonkeyEval, BeatNumEval—to evaluate private-library code generation. Empirical results show that API information and retrieved APIs significantly boost performance, with CodeGenAPI offering notable gains on private-library tasks and human-in-the-loop retrieval providing additional improvements.

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.

When Language Model Meets Private Library

TL;DR

Private libraries present unseen APIs that challenge contemporary code-generation models. The authors introduce a retrieval-then-generation framework (APIRetriever and APICoder) that leverages API documentation to enable code synthesis with private libraries, trained on public-library data to generalize. They also create three benchmarks—TorchDataEval, MonkeyEval, BeatNumEval—to evaluate private-library code generation. Empirical results show that API information and retrieved APIs significantly boost performance, with CodeGenAPI offering notable gains on private-library tasks and human-in-the-loop retrieval providing additional improvements.

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.
Paper Structure (30 sections, 1 equation, 7 figures, 5 tables)

This paper contains 30 sections, 1 equation, 7 figures, 5 tables.

Table of Contents

  1. Introduction
  2. Framework
  3. Methodology
  4. Data Collection
  5. APIRetriever
  6. Training.
  7. Inference.
  8. Human Interaction with APIRetriever.
  9. APICoder
  10. Benchmark Construction
  11. Experiments
  12. Experimental Setup
  13. API Information.
  14. Baselines.
  15. Evaluation Metrics.
  16. ...and 15 more sections

Figures (7)

  • Figure 1: An example of converting PandasEval (public) to MonkeyEval (private). The changed parts are highlighted in yellow. The performance of Codex $12$B and CodeGen-Mono$350$M is shown at the bottom.
  • Figure 2: The overview of our proposed framework.
  • Figure 3: The training process of APIRetriever and CodeGenAPI.
  • Figure 4: Friendly interaction interface for users.
  • Figure 5: The recall rates of retrieved APIs.
  • ...and 2 more figures