R e ACC : A Retrieval-Augmented Code Completion Framework (2022)

Abstract: Code completion, which aims to predict the following code token(s) according to the code context, can improve the productivity of software development. Recent work has proved that statistical language modeling with transformers can greatly improve the performance in the code completion task via learning from large-scale source code datasets. However, current approaches focus only on code context within the file or project, i.e. internal context. Our distinction is utilizing ”external” context, inspired by human behaviors of copying from the related code snippets when writing code. Specifically, we propose a retrieval-augmented code completion framework, leveraging both lexical copying and referring to code with similar semantics by retrieval. We adopt a stage-wise training approach that combines a source code retriever and an auto-regressive language model for programming language. We evaluate our approach in the code completion task in Python and Java programming languages, achieving a state-of-the-art performance on CodeXGLUE benchmark.

## introduction with the growth of software engineering field, large-scale source code corpus gives a chance to train language models in code domain @xcite . and benefiting from the large transformer models @xcite and pre-training techniques @xcite , a rapid progress has been made in many code-related tasks like code search @xcite , code summarization @xcite , bug fixing @xcite and code completion @xcite @xcite . code completion is considered as an essential feature towards efficient software development in modern integrated development environments (ides). the task is formulated by predicting the following code token(s) based on the code context. traditionally, code completion requires real-time program analysis and recommends type-correct code tokens @xcite . recently, statistical language models trained on large-scale source code data have shown high accuracy in the code completion task. primitive approaches take the given context only @xcite , some methods use richer information, e.g., adding code token types @xcite , abstract syntax tree (ast) structures @xcite , or extended hierarchical context @xcite . however, one key limitation of existing methods is the scope of information they utilize; all the information is bounded in the given input file. this is unnatural from human perspective, as studies demonstrate that programmers tend to reuse an existing code snippet by copying part of code with or without minor modifications to accelerate software development @xcite , leading a software repository usually containing 7-23% cloned codes @xcite . motivated by this phenomenon, in this paper, we argue the utility of extending the information scope beyond the input file, i.e., into a large codebase. we conjecture that using codes with similar semantics as auxiliary information are beneficial to predict the following code tokens. therefore, we propose reacc -a retrieval-augmented code completion framework (see figure 1 ). the code completion task under our framework can be re-formulated by, given a source code corpus for search and an unfinished code snippet to complete, using the unfinished code as a query to retrieve similar code snippets from search corpus, and predicting the following code tokens by reusing the the simplest technique for retrieving code is to build a sparse vector retriever like tf-idf or bm25 @xcite which are both based on keyword matching algorithms. the sparse retriever can capture lexical information and is sensitive to the names of code identifiers. the dense retriever, on the contrary, can capture syntactic and semantic information by mapping a code snippet to a dense vector. in the code completion task, the code retriever is expected to comprehend the source code's intent in order to retrieve the semantically similar codes. on the other hand, considering programmers are prone to copy-and-paste existing code, the retriever should evaluate lexical similarity as well. to that end, we adopt the hybrid retriever @xcite , which combines results of dense and sparse retriever. we employ a dual-encoder model architecture as the dense retriever since the cross-encoder model has a high computational complexity. to achieve a better understanding ability, we initialize our dense retriever with @xcite , which is a pre-trained bert-based programming language understanding model. then we continue pre-training the retriever by contrastive learning to enhance sentence embedding. as the labeled data containing similar code pairs is rare, we utilize various transformations to generate programs with similar functionality for data augmentation. we implement the generator with a decoder-only transformer model. to incorporate the external information from retrieved similar code, we concatenate the obtained code and code context as input. the generator is initialized by codegpt-adapted @xcite which is a domain-adaptation model from gpt-2 @xcite pretrained on code corpus. we evaluate our reacc framework on two benchmark datasets -codexglue @xcite and codenet @xcite , in python and java programming languages. reacc achieves a state-of-the-art performance on both datasets. the experimental results demonstrate that external source code retrieved by our retriever is useful for auto-completing the partial code. to summarize, our main contributions are: ## approach we first introduce the formulation of retrievalaugmented code completion task. then we give detailed descriptions on the retriever and generator in reacc. we show how we continue pretraining @xcite with contrastive learning on code and how we address the problem that there is no labeled data for positive instances of similar programs in section 3.2. in section 3.3 we talk about the way to aggregate retrieved code and code context in the generator. ## experiments: code clone detection in order to evaluate the effectiveness of the codeto-code retrieval module in reacc, we perform code clone detection task which aims to retrieve semantic equivalent programs. in this section, we describe how we create the test dataset for this task and how we evaluate the performance of reacc's retriever. ## experiments: code completion in this section, we evaluate reacc on end-to-end code completion. ## conclusion we propose reacc, a retrieval-augmented code completion framework that utilizes "external" context for the code completion task by retrieving semantically and lexically similar codes from existing codebase. we pre-train a dual-encoder as a retriever for partial code search, which retrieves code fragments given a partial code. our method can adopt any architecture that can perform code completion as the generator. on the codexglue benchmark, reacc achieves a state-of-the-art performance in the code completion task.