# Generating Code with the Help of Retrieved Template Functions and Stack
Overflow Answers

Dawn Drain
Microsoft Cloud and AI
<EMAIL_ADDRESS>Hu
Microsoft Cloud and AI
<EMAIL_ADDRESS>Wu
Microsoft Cloud and AI
<EMAIL_ADDRESS>Breslav
Microsoft Cloud and AI
<EMAIL_ADDRESS>Sundaresan
Microsoft Cloud and AI
<EMAIL_ADDRESS>Corresponding author Work done as an intern at Microsoft

###### Abstract

We approach the important challenge of code autocompletion as an open-domain
task, in which a sequence-to-sequence code generator model is enhanced with
the ability to attend to reference code snippets supplied by a semantic code
search engine. In this work, we present a novel framework to precisely
retrieve template functions as well as intent-snippet pairs and effectively
train such a retrieval-guided code generator. To demonstrate the effectiveness
of our model designs, we perform extensive experiments with CodeSearchNet
(Husain et al., 2019) which contains template functions and CoNaLa (Yin et
al., 2018) which contains Stack Overflow intent-snippet pairs. We also
investigate different retrieval models, including Elasticsearch, DPR
(Karpukhin et al., 2020), and our fusion representation search model, which
currently holds the number one spot on the CodeSearchNet leaderboard. We
observe improvements by leveraging multiple database elements and further gain
from retrieving diverse data points by using Maximal Marginal Relevance.
Overall, we see a 4% improvement to cross-entropy loss, a 15% improvement to
edit distance, and a 44% improvement to BLEU score when retrieving template
functions. We see subtler improvements of 2%, 11%, and 6% respectively when
retrieving Stack Overflow intent-snippet pairs. We also create a novel Stack
Overflow-Function Alignment dataset, which consists of 150K tuples of
functions and Stack Overflow intent-snippet pairs that are of help in writing
the associated function, of which 1.7K are manually curated.

##  1 Introduction

Code autocompletion is an important and ubiquitous tool for enhancing
developer productivity. Recent advances in NLP, especially the usage of large,
pretrained sequence-to-sequence neural models (Lewis et al., 2019; Raffel et
al., 2019; Bi et al., 2020; Song et al., 2019) have enabled code completion
systems to write larger and more complicated code snippets, such as unit
tests, bug-fixes, transcompiler and arbitrary methods (Tufano et al., 2020;
Clement et al., 2020; cod, 2020).

These models typically have an impoverished context compared to what
professional developers can access. At the most extreme, the input may only be
a function’s signature. In contrast, a developer may consult pre-existing
implementations of related functions, look up API usages, and peruse solutions
to common questions on sites like Stack Overflow. That is, writing code is in
general an open-domain task, whereas researchers typically treat it as closed-
book.

Example of Retrieving a Template
⬇ # Input def add_toc_to_epub(tocless_epub_path, write_path):  """  reads the
epub file, adds a table of contents,  and saves it """  # Retrieved Template
def create_epub(epub_path):  """  Creates an example epub including metadata,
toc, and css and saves it to epub_path """  \parbook = epub.EpubBook()
book.set_title(’Sample book’)  \par# create chapter c1 =
epub.EpubHtml(title=’Intro’,  file_name=’chap_01.xhtml’, lang=’hr’)
c1.content=u”’
Intro heading

Zaba je skocila u baru.
”’  \par#
add chapter book.add_item(c1)  \par# define Table Of Contents book.toc =
(epub.Link(’chap_01.xhtml’,  ’Introduction’, ’intro’),  (epub.Section(’Simple
book’),  (c1, ))  )  \par# add default NCX and Nav file
book.add_item(epub.EpubNcx())  book.add_item(epub.EpubNav())  \par# basic
spine book.spine = [’nav’, c1]  \par# write to the file
epub.write_epub(’test.epub’, book, {})

⬇

\pardef add_toc_to_epub(tocless_epub_path, write_path):

"""

reads the epub file, adds a table of contents,

and saves it

"""

# Generated Code

book = epub.read(epub(tocless_epub_path))

book.toc = tuple(book.get_items_of_type(

epub.EpubHtml))

book.spine = [’nav’] + book.spine

epub.write_epub(write_path, book, {})

Figure 1: We consider the task of generating a function’s body given its
signature and docstring. Given that input, we use the docstring as a query to
retrieve a helpful reference function or code snippet, create_epub in this
manually created example, which we concatenate with the input and then feed to
our generator model.
Figure 2: Flow charts of training and inference based on retrieving template
functions. A flexible framework for retrieval-guided code generation which
consists of a neural code search model for template function retrieval and a
template guided code generator. The training of the template functions
searching and the code snippet generator is decoupled, and they work
cooperatively under the help of the retrieval system.

To address the problem, we similarly consider the task of writing functions
given a signature and docstring, and attempt to bridge this gap by enabling
our models to retrieve template functions from GitHub and answers to relevant
questions from Stack Overflow. For both of these databases, we experiment with
several different retrieval models. In all cases we use a pre-trained
sequence-to-sequence transformer as the generator, and condition generation on
the target function’s signature and docstring as well as the retrieved
documents. We attain our best results when using our fusion representation
model to retrieve as many template functions from CodeSearchNet (Husain et
al., 2019) as fit in the context window, with the templates re-ranked
according to their Maximal Marginal Relevance scores (Carbonell and Goldstein,
1998). Figure 2 depicts the training and inference procedures of our
framework.

Rather than retrieve from all of Stack Overflow, our work uses CoNaLa as the
Stack Overflow database (Yin et al., 2018). CoNaLa consists of intent-snippet
pairs, where each intent is the title of a Stack Overflow question concerning
Python, and each code snippet is extracted from the answering question using a
neural model. More specifically, we follow (Xu et al., 2020), use the most
confidently extracted 100K intent-snippet pairs, including 2.3K that were
manually curated.

Experiments show that our method significantly improves the informativeness of
the generated code as well as their relevance to the corresponding query. When
retrieving template functions, we see a 4% improvement to cross-entropy loss,
a 15% improvement to edit distance, and a 44% improvement to BLEU score. We
see subtler improvements of 2%, 11%, and 6% respectively when retrieving Stack
Overflow answers. In addition, We conduct extensive ablation studies to
evaluate the improvement from different retrieval model design. To summarize,
our contributions are as follows:

  * •

We propose a flexible two-stage framework for retrieval-guided code
generation.

  * •

We propose a semantic code search model for template functions retrieval.

  * •

We propose a template-aware code generator that can handle template functions.

  * •

We create a novel Stack Overflow-Function Alignment dataset, which consists of
150K tuples of functions and Stack Overflow intent-snippet pairs that are of
help in writing the associated function, of which 1.7K are manually curated.
We will release the dataset for future research.

##  2 Related Work

Some prior work has considered code generation through the lens of search.
Similarly to our approach, (Hashimoto et al., 2018) experiment with retrieving
a template method given the input function signature and docstring, resulting
in a large boost in BLEU score from 19.2 to 34.7. In order to train the
retriever, they attempt to minimize the edit distance between the retrieved
template and the ground truth function. In the CodeSearchNet challenge,
(Husain et al., 2019) focus more specifically on curating a database of
functions and on performing code-search using natural language queries.
Additionally, they host a leaderboard for various retrieval models, including
several baselines they have open-sourced. They take a function’s docstring as
a query, excluding parameter annotations, and form a database out of functions
denuded of their docstrings. In order to train the retriever, they give a
binary reward, with one point for retrieving the function the docstring
belongs to, and zero points otherwise.

Several works have also explored using Stack Overflow as a database. As SO
questions and answers can be very long, a standard approach is to restrict to
the question’s title and a code block extracted from the corresponding answer.
Among these works, both CODENN (Iyer et al., 2016) and StaQC (Yao et al.,
2018) match the question title to the entire code block, while CoNaLa (Yin et
al., 2018) matches the question title more precisely to the key lines in the
code block. Our work uses CoNaLa as the Stack Overflow database to retrieve,
as it allows for a higher information density, and enables retrieving more
intent-snippet pairs.

In NLP more broadly, much work has focused on open-domain question answering
with Wikipedia as a database. Typically the answer is at most a few words
long, and the retriever is rewarded for returning an article containing that
exact substring. One state of the art retriever is DPR (Karpukhin et al.,
2020), which is based on BERT (Devlin et al., 2019). DPR outperforms a strong
Lucene-BM25 baseline by 9-15%, and attains superior results with as few as one
thousand training examples. The follow-up work RAG pipelines DPR with BART as
a generator, similarly to our approach (Lewis et al., 2020). Open-domain
scores compare quite favorably to those of models constrained to closed-book
question answering. For example, the 11B-parameter T5 achieves a score of 34.5
EM on Natural Questions and the 770M-parameter T5 achieves a score of 28.9 EM
(Raffel et al., 2019). Meanwhile, RAG achieves a score of 44.5 EM with only
626M parameters (Lewis et al., 2020).

##  3 Retrieval-guided Code Generation

###  3.1 Task Definition

The problem is as follows: Given a function’s signature and docstring as
input, produce the rest of the function’s body. We treat this as an open-
domain problem, and experiment both with retrieving template functions and
relevant answers from Stack Overflow. For simplicity we focus on Python
functions, although our approach is language agnostic.

Consider a database $D=\\{d_{1},\dots,d_{m}\\}$, consisting of either template
functions or Stack Overflow answers, and a set of target functions
$T=\\{t_{1},\dots,t_{n}\\}$ that begin with the signatures and docstrings
$\\{s_{1},\dots,s_{n}\\}$. In the closed-book setting we train a code
generator $G_{c}$ so as to minimize the cross-entropy:

$\mathbb{E}_{T}\text{log}P(t_{i}|s_{i},G_{c}).$

In our open domain setting, we seek a retriever model $R$ and generator
$G_{o}$ to minimize:

$\mathbb{E}_{T}\text{log}P(t_{i}|s_{i},G_{o},d_{i,1},\dots,d_{i,k_{i}}),$

where $\\{d_{i,1},\dots,d_{i,k_{i}}\\}$ is the set of documents retrieved by
$R$ given $s_{i}$.

In practice we condition retrieval solely on the docstring, although we
condition on both the signature and docstring during generation along with the
retrieved documents.

###  3.2 Retrieval Models

All of our neural retrieval models are based around nearest neighbor lookup
between a vectorized query and a database of vectorized function templates or
Stack Overflow answers. We take the function’s docstring as the query, as that
is the same input used in CodeSearchNet. In all cases we use separate encoders
for the queries and database elements, although in general one could reuse a
single encoder.

Rather than perform exact nearest neighbor lookup, in practice we use
approximate nearest neighbor search using ANNOY (Bernhardsson, 2020). ANNOY
works by partitioning the entire vector space into subspaces through constant
random projections and then performing nearest neighbor search in the lower-
dimensional subspaces. In our experiments we use angular distance as the
distance metric, set $n\\_trees=1000$, and set $search\\_k=10000$.

In the following section, we will first introduce three retrieval models on
CodeSearchNet in 3.2.1, 3.2.2 and 3.2.3 respectively, and then presents a
retrieval model on Stack Overflow in 3.2.4.

####  3.2.1 ElasticSearch Baseline

We use Elasticsearch as our baseline retrieval model (Gormley and Tong, 2015).
ElasticSearch is a widely used, inverted index-based search engine. Here we
use the NDCG evaluation results reported by CodeSearchNet.

####  3.2.2 NBoW Bi-encoder Baseline

We use the Joint Vector Representations for Code Search proposed by
CodeSearchNet as another baseline. Specifically, we use their open-sourced
NBoW (Husain et al., 2019) model as a bi-encoder.

The model consists of two NBoW encoders, one for transforming the natural
language documentation of a function into a vector representation, and one for
transforming the corresponding code into a vector representation. The natural
language and code are then mapped into a single vector space by training.

Before being put into the model, the data was first processed by splitting the
code by identifiers in the code and splitting the natural language
documentation using BPE. Secondly, the tokenized sequences are entered into
the corresponding encoder separately to obtain a series of vector
representations. After that, the series of vector representations are
condensed into a single vector through one of mean-pooling, max-pooling, and
an attention-like weighted sum mechanism. Finally, the dot product between the
vectors obtained from the natural language documentation and the code is
computed. To train the model, natural language and code snippets from within
the same function are rewarded, while discrepant pairs are penalized.

####  3.2.3 Bi-encoder with Fusion Representation

We present our fusion representation model, which is currently ranked number
one on the CodeSearchNet leaderboard. We build on CSN’s open-sourced NBoW
baseline, and more than double its mean-NDCG score while maintaining the
dimensionality of the embedding space and only adding one linear alignment
layer and two triplets of fusing weights to the parameter count.

Recall that the NBoW baseline trains two encoders, one for the docstring query
and one for the template functions comprising the database. Each encoder
condenses its hidden states $\mathrm{H}$ to a single vector $\mathbf{h}$ using
one of mean-pooling, max-pooling, or a self-attention type of weighted
average. We reuse those three types of pooling, with the modification that we
standardize the third type of pooling by adding a softmax to the attention
scores as below, where $w_{h}$ is a learnable weight vector.

$\gamma={softmax}\left(\mathbf{w}_{\mathrm{h}}^{\top}\cdot\mathrm{H}\right)$
$\mathbf{h}_{attn}=\sum_{\mathrm{j}}\gamma_{\mathrm{j}}\cdot\mathrm{H}_{:\mathrm{j}},\forall\mathrm{j}\in[1,\ldots,\mathrm{m}]$

We fuse all three vectors using a linear combination
$\beta\cdot\sum_{k\in(mean,max,attn)}w_{k}\mathbf{h}_{k}$ with learnable
weights $w_{mean},w_{max},\text{ and }w_{attn}$. $\beta$ allows the language
encoder to scale the entire fused vector. Finally, we add a learnable linear
layer to the code-encoder only, to help align the code-embedding space with
the docstring-embedding space. Empirically, the largest weight is given to
$w_{attn}$.

####  3.2.4 Dense Passage Retrieval

Dense Passage Retrieval (DPR) is a bi-encoder model proposed to solve the
retrieval part of open domain question answering (Karpukhin et al., 2020). As
with our fusion representation model, DPR trains separate encoders to embed
the queries and database documents. Both encoders of DPR are initialized as
pretrained BERT models.

We fine-tune the open-sourced DPR model on our proposed dataset SOFA-auto-
mined (See SOFA for more details), using each function’s docstring as the
query and the corresponding intent-snippet pair as the target database
document. As DPR warmstarts using a BERT model pretrained on English, we form
the intent-snippet embedding using only the intent question-title. We take all
SO-functions pairs in SOFA-auto-mined as positive examples. For each method we
randomly select 15 SO intent-snippet pairs as negative examples. Intuitively,
training directly on retrieving intent-snippets given a function’s docstring
as input will lead to better downstream performance than training with
question titles (the intents) as input.

###  3.3 Generative Model

We use BART-large as the body of our Seq2Seq architecture. BART-large is a
performative transformer with 406M-parameters, including twelve encoder layers
and twelve decoder layers (Lewis et al., 2019). We initialize the weights from
the open-sourced BART-large which was pretrained on English corpus, and
continue pretraining on a mix of masked Python code and Stack Overflow
question-answer pairs. We additionally add a small number of randomly
initialized vocabulary embeddings for common whitespace tokens such as the
four-space, eight-space, and two-tab tokens, which increases the effective
context window length and throughput.

##  4 Data Preparation

###  4.1 Pretraining Corpus

Following best practice in NLP, we conduct extensive pretraining before
finetuning on the task of method generation. We take advantage of the open-
sourced model BART (Lewis et al., 2019), which was pretrained for 40 epochs on
160GB of English text using a denoising objective. In order to adapt the model
to the domain of source code, we augment its tokenizer with whitespace tokens,
such as the four-space token and the two-tab token. We then perform task-
adaptive pretraining on a mix of 4GB of Stack Overflow question-answers pairs,
corresponding to 1.35M questions tagged ‘python’ with answers that have non-
negative scores, and 27GB of raw python masked in the style of T5 (Raffel et
al., 2019). We train for two weeks on sixteen 32GB Tesla V100 GPU’s.

###  4.2 Code Generation Dataset

We introduce the two datasets we use as databases to retrieve from and
reference during generation.

####  4.2.1 Template Functions Database

We first adopt the CodeSearchNet (Husain et al., 2019) dataset of python
functions that have docstrings, and perform the following data cleaning.

First, we remove all class methods. Class methods are generally highly
dependent on their fellow member functions, which are not included as part of
CSN. Furthermore, we do not need to train a retrieval model to select this
class information, as classes are already localized within a single file.

Second, we tokenize functions based on punctuation and spaces, and remove
functions with more than 150 tokens, which are often too complex to generate
in their entirety. Note that we use the more standard BPE tokens during
training and inference.

We split this cleaned version of CSN into training, validation, and test sets,
and use the training set to build the function template database.

Dataset | Train | Validation | Test
---|---|---|---
Non-class CSN | 119480 | 1000 | 1000
Table 1: Number of CodeSearchNet functions after filtering. We exclude class-
methods, methods longer than 150 tokens after splitting on punctuation and
whitespace, and all CSN methods that lack docstrings.

####  4.2.2 Stack Overflow Database

We performed preliminary experiments using all Stack Overflow Python questions
as the database, and attained underwhelming results. Manual inspection
revealed that very few answers were actually of relevance to the target
function body we were trying to generate, even after cursory filtering based
on the number of stars each answer received.

We thus turned towards a more extensively filtered subset of Stack Overflow
answers, as explored previously by (Yin et al., 2018), in their work leading
to the creation of the dataset CoNaLa. Similar to our goal, their aim was to
retrieve code snippets from Stack Overflow that best answered a natural-
language query.

CoNaLa is a superior database for two reasons. First, rather than consisting
of questions and answers in their entirety, CoNaLa keeps only the question
title, which is typically sufficient to gather the question’s intent, as well
as the most relevant code snippet from the question’s answer, which is
determined using machine learning. Restricting to these intent-snippet pairs
increases information density, and allows us to pack many more retrieved
answers into the context window.

Second, CoNaLa is filtered to only contain 600K intent-snippet pairs in python
that address “how-to" questions, such as "how to flatten a list in python", in
contrast to questions addressing a particular bug or error message, or open-
ended questions discussing more philosophical topics.

Following (Xu et al., 2020), we use the 100K highest rated intent-snippet
pairs as our SO-database, along with 2.3K intent-snippet pairs that CoNaLa
manually curated in order to train their relevance scoring model.

Intent | How to convert a list of multiple integers into a single integer?
---|---
Snippet | sum(d * 10 ** i for i, d in enumerate(x[::-1]))
Intent | Averaging the values in a dictionary based on the key
Snippet | [(i, sum(j) / len(j)) for i, j in list(d.items())]
Table 2:  Example intent-snippet pairs from the CoNaLa dataset. Each intent is
the title of a Stack Overflow question, and each snippet is a code snippet
that answers the question, as determined using a neural model.

###  4.3 SOFA, a Novel Code Search Dataset

The input to our retrieval models is a function’s docstring, which represents
a different distribution from Stack Overflow question titles. Thus we are
motivated to create a dataset that more directly caters to our intended
workflow than does CoNaLa. Towards this aim we create a dataset consisting of
(function, (intent, snippet)) pairs, where each SO intent-snippet is selected
so as to be of help in generating the CSN method. We term this the Stack
Overflow-Function Alignment (SOFA) dataset.

In order to pair an intent-snippet with a CSN method, we use our
implementation of the similarity measure introduced in Aroma (Luan et al.,
2019). We manually inspect the suggested method for each of the 2.3K intent-
snippet pairs manually curated in forming CoNaLa and filter out 600 pairs for
which the intent-snippet is unhelpful for generating the corresponding
function or even itself trivial. To scale up we first rank the 600K intent-
snippet pairs in CoNaLa by their confidence scores. We next keep the top 10K
pairs, which we consider to be of high quality. Then for each intent-snippet
pair we retrieve 15 snippets from CSN that are most similar, and discard
snippets which are not in our filtered version of CSN.

Figure 3: Overview of SOFA, which we use to train the retriever DPR-SO. We build on the database CoNaLa, which consists of intent-snippet pairs, where each intent is the title of a Stack Overflow question, and each code snippet is extracted from that question’s answer. We automatically match functions from CodeSearchNet with relevant intent-snippet pairs using the code search tool Aroma. For the subset of CoNaLa intent-snippet pairs that were manually curated, we manually discard unhelpful matches provided by Aroma. SO-intent | How to convert a list of multiple integers into a single integer? |  |  |
---|---|---|---|---
SO-snippet | sum(d * 10 ** i for i, d in enumerate(x[::-1])) |  |  |
CSN-function |  | def lower_endian_to_number(l, base):
---
"""Helper function: convert a list of digits in the given base to a number."""
return sum([d * (base**i) for i, d in enumerate(l)])
|  |
Table 3: Example from our manually curated SOFA benchmark. We include this
CoNaLa intent-snippet pair as it is considered helpful in generating the
CodeSearchNet function.

Stack Overflow Retrieval Example ⬇ # Input def nii_ugzip(imfile, outpath=”):
"""  Uncompress *.gz file """  # Retrieved SOFA intent-snippet pairs # Loading
a file into a numpy array with pythonfile = open(’file.txt’, ’rb’)  \par#
Generate temporary file names # without creating actual file in Python f =
gzip.open(temp_file_name, ’wb’)  \par# gzip a file in Python f =
gzip.open(’/home/joe/file.txt.gz’, ’wb’)  \par# Using openpyxl to read file
from memory file = open(’path/to/file.xlsx’, ’rb’)  wb =
openpyxl.load_workbook(filename=file)  \par# gzip a file in Python f_in =
open(’/home/joe/file.txt’)  ⬇ # Hypothesis without context:  imfile =
os.path.abspath(imfile)  if not os.path.exists(outpath):  os.makedirs(outpath)
with gzip.open(imfile, ’rb’) as f:  nii_data = f.read()  return nii_data ⬇ #
Hypothesis with context:  if not os.path.exists(outpath):  os.mkdir(outpath)
with gzip.open(imfile, ’rb’) as infile:  with gzip.open(outpath, ’wb’) as
outfile:  outfile.write(infile.read())  Figure 4: Example of the benefits of
retrieving intent-snippet pairs from Stack Overflow. The generated hypothesis
with context for this example learns to write into the outpath provided in the
function signature, whereas the hypothesis without context merely reads the
file and returns it.

##  5 Experiments and Results Analysis

We experiment with manipulating four different variables:

  1. 1.

The retrieval model used (Elasticsearch, Bi-encoder with Fusion
representation, or DPR)

  2. 2.

The number of function templates retrieved

  3. 3.

Retrieving the top-k results independently vs. using Maximal Marginal
Relevance (MMR)

  4. 4.

The choice of database (CodeSearchNet or Stack Overflow)

###  5.1 Retrieval Models

We experiment with three different retrieval models on CodeSearchNet:
Elasticsearch, our own fusion representation model, and the DPR model
(Karpukhin et al., 2020).

In Table 4 we compare our fusion representation model against Elasticsearch
and NBoW on CodeSearchNet. Notably, Elasticsearch outperforms the neural
baselines as reported in the CodeSearchNet paper. The fusion representation
model outperforms the ElasticSearch baseline by 4.7% on average NDCG across
the six languages, and 6.6% on python NDCG only. The Fusion Representation
model also outperforms the current state of the art NBoW model by 1.25% on
average NDCG, and 0.9% on python NDCG only, which suggests that our proposd
fusion strategy can better accommodate natural language and programming
language to the same vector space.

| ES | NBoW | Fusion Rep |  |
---|---|---|---|---|---
| NDCG
---
(mean)
33.7% | 37.2% | 38.4% |  |
| NDCG
---
(python)
40.6% | 46.3% | 47.2% |  |
Table 4: Comparison between Elasticsearch, Neural bag-of-words retriever
(NBoW), and our fusion representation model on the Codesearchnet benchmark.
Note that we draw the NBoW results from the CSN leaderboard, rather than from
the weaker NBoW model open-sourced as part of CSN. The mean NDCG score is
taken across six different programming languages.

###  5.2 Number of Examples Retrieved

When experimenting with CodeSearchNet, we consider four different limits on
the number of template functions retrieved.

No-Retrieve: As a baseline, we consider only using the query as the
generator’s input.

Single-Retrieve: Select the method with the highest similarity to the query
and use its signature, docstring, and method body as contextual information.

Threshold-Retrieve: More strictly than in the single-retrieve condition, only
present a template function if its retrieval score is greater than the median
score across all data points. We hypothesize that foregoing mediocre
retrievals will help prevent mode collapse, in which the generator learns to
ignore the retrieved data.

Full-Retrieve: We retrieve as many template functions as fit in the
generator’s context window, presented in decreasing order of retrieval score.
Since the generator is based on BART-large, its context window is 1024 tokens
wide, some of which is consumed by the query function’s signature and
docstring.

In table 5 we experiment with varying the number of template functions
retrieved. We observe monotonically increasing performance with the number
retrieved, except for our control condition in which we retrieve random
templates and observe a small drop in performance. Filling out the entire
context window with retrieved templates is roughly twice as effective as
retrieving a single template. In the Threshold-Retrieve condition we also
experiment with only retrieving a single template if its retrieval score
passes a threshold, which does worse than simply always retrieving a template.

| Ppl | EDist |  | BLEU-4
---
(BPE)
BLEU-4
No retrieve | 1.25 | 1.46 | 17.40 | 5.76
Threshold retrieve | 1.25 | 1.43 | 18.61 | 6.98
Single retrieve | 1.24 | 1.35 | 19.83 | 7.69
Full retrieve | 1.21 | 1.30 | 22.17 | 9.66
| Full random
---
retrieve
1.24 | 1.46 | 16.83 | 5.52

Table 5:  Results for retrieving multiple templates. In all cases we use our
fusion representation model as the retriever and CSN as the database. We
observe monotonically increasing performance with the number of retrieved
templates, and no benefit from only including documents whose retrieval scores
surpass a threshold.

###  5.3 Returning top-k independently vs. MMR

There is a well-known diversity problem when retrieving multiple search
results. That is, there is little benefit in retrieving a set of redundant
templates or explanations, and we would achieve better results by retrieving a
more comprehensive set of data points. One popular technique for retrieving
diverse results is Maximal Marginal Relevance (MMR) (Carbonell and Goldstein,
1998). Given the query $q$, we first retrieve the closest datapoint $d_{1}$ as
usual. Inductively, having already retrieved data points
$\\{d_{1},\dots,d_{i}\\}=:D_{i}\subset D$, we next retrieve

$d_{i+1}:=argmax_{d\in D\setminus D_{i}}\big{(}\lambda(d\cdot
q)+(1-\lambda)\max_{d_{i}\in D_{i}}d\cdot d_{i}\big{)}$

where ‘$\cdot$’ denotes the dot product. Setting $\lambda=1$ recovers the
standard $k$-nearest neighbor approach. For $\lambda<1$ the second term
penalizes documents that are similar to any that have already been retrieved.
In our experiments we set $\lambda=.5$.

As a control, we also experiment with eschewing a retriever and instead
collecting functions at random. As for greedy retrieval and MMR, we provide as
many functions as fit in the 1024-token context window.

In table 6 we achieve our strongest results by not merely greedily adding
templates based on their retrieval scores, but instead by greedily adding
successive templates according to their marginal relevance scores. When
reranking based on MMR, we add a diversity term to each retrieval score based
on the template’s distance from previously retrieved examples. MMR doubles the
benefit from retrieving more than one template, as measured by the edit
distance and BLEU scores.

| Ppl | EDist |  | BLEU-4
---
(BPE)
BLEU-4
Full retrieve | 1.21 | 1.30 | 22.17 | 9.66
MMR retrieve | 1.20 | 1.24 | 25.09 | 11.87

Table 6:  Results for retrieving templates independently vs. using Maximal
Marginal Relevance. In all cases we use our fusion representation model as the
retriever and CSN as the database. We achieve our best results when using MMR
to incrementally add diverse documents, in contrast to the standard greedy
approach.

##  6 Discussion

###  6.1 Choice of Database

We experiment between retrieving relevant template functions from the
CodeSearchNet dataset and retrieving relevant intent-snippet pairs from Stack
Overflow. We achieve poor results when retrieving from all Stack Overflow
questions tagged ‘python’, and thus restrict to the refined subset shared by
CoNaLa. During validation we restrict to using the training set as the
database so as to prevent leakage. During training we additionally exclude any
retrieved template method that is identical to the target method, so as to
prevent degenerate behavior on the part of the retriever.

When retrieving template functions from CodeSearchNet, we achieved our
strongest results when retrieving as many templates as possible and reranking
them using MMR. Thus we use the same approach when experimenting with Stack
Overflow as a database. The intent-snippet pairs we retrieve give a relatively
small boost, as reported in table 7. Indeed, we achieve better results using a
single template from CodeSearchNet compared to filling out the entire context
window with intent-snippet pairs taken from Stack Overflow.

Perhaps the Stack Overflow database gives comparatively worse results since
the retriever is weaker. If so, we might speculate that this is because it is
so much harder to construct ground truth pairings between functions and
intent-snippets. Alternatively the differing results could be due to the
quality of the databases themselves. Search is most useful when recalling
details about the long tail of obscure API usages and when making small
modifications to pre-existing code. We think this information is better
supplied by functions found in the wild than by the kind of supplementary
documentation found on Stack Overflow.

| Ppl | EDist |  | BLEU-4
---
(BPE)
BLEU-4
No retrieve | 1.25 | 1.46 | 17.40 | 5.76
Full ES | 1.24 | 1.44 | 17.20 | 5.82
Full fusion representation | 1.23 | 1.33 | 18.54 | 6.32
Full DPR-SO | 1.24 | 1.36 | 17.56 | 5.52

Table 7:  Evaluation of Stack Overflow as a database. We retrieve as many
answers as fit in the context window and filter those answers using MMR,
inspired by the success of that approach when using template functions from
CSN.

##  7 Conclusion

In this paper, we presented a novel framework pipelining search with
generation for retrieval-guided code generation. Our method uses a neural code
search model for extracting template functions and intent-snippet pairs and
uses a robust programming language generator for code generation. The two
components are trained separately to allow more flexibility and efficiency. In
order to integrate Stack Overflow intent-snippet pairs with the task of
function generation we also created a novel code search dataset, SOFA, which
we will release for future research. Experiments show that our method
significantly outperforms several strong baselines.

## References

  * cod (2020) 2020\.  Codexglue: A benchmark dataset and open challenge for code intelligence.
  * Bernhardsson (2020) Erik Bernhardsson. 2020.  _Annoy: Approximate Nearest Neighbors in C++/Python_.  Python package version 1.16.0.
  * Bi et al. (2020) Bin Bi, Chenliang Li, Chen Wu, Ming Yan, and Wei Wang. 2020.  Palm: Pre-training an autoencoding&autoregressive language model for context-conditioned generation.  _arXiv preprint arXiv:2004.07159_.
  * Carbonell and Goldstein (1998) Jaime Carbonell and Jade Goldstein. 1998.  The use of mmr, diversity-based reranking for reordering documents and producing summaries.  In _Proceedings of the 21st annual international ACM SIGIR conference on Research and development in information retrieval_ , pages 335–336.
  * Clement et al. (2020) Colin B. Clement, Dawn Drain, Jonathan Timcheck, Alexey Svyatkovskiy, and Neel Sundaresan. 2020.  Pymt5: multi-mode translation of natural language and python code with transformers.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing, EMNLP 2020, Online, November 16-20, 2020_ , pages 9052–9065. Association for Computational Linguistics.
  * Devlin et al. (2019) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019.  Bert: Pre-training of deep bidirectional transformers for language understanding.
  * Gormley and Tong (2015) Clinton Gormley and Zachary Tong. 2015.  _Elasticsearch: the definitive guide: a distributed real-time search and analytics engine_.  " O’Reilly Media, Inc.".
  * Hashimoto et al. (2018) Tatsunori B Hashimoto, Kelvin Guu, Yonatan Oren, and Percy S Liang. 2018.  A retrieve-and-edit framework for predicting structured outputs.  In _Advances in Neural Information Processing Systems_ , pages 10052–10062.
  * Husain et al. (2019) Hamel Husain, Ho-Hsiang Wu, Tiferet Gazit, Miltiadis Allamanis, and Marc Brockschmidt. 2019.  Codesearchnet challenge: Evaluating the state of semantic code search.  _arXiv preprint arXiv:1909.09436_.
  * Iyer et al. (2016) Srinivasan Iyer, Ioannis Konstas, Alvin Cheung, and Luke Zettlemoyer. 2016.  Summarizing source code using a neural attention model.  In _Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 2073–2083, Berlin, Germany. Association for Computational Linguistics.
  * Karpukhin et al. (2020) Vladimir Karpukhin, Barlas Oğuz, Sewon Min, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. 2020.  Dense passage retrieval for open-domain question answering.  _arXiv preprint arXiv:2004.04906_.
  * Lewis et al. (2019) Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Ves Stoyanov, and Luke Zettlemoyer. 2019.  Bart: Denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehension.
  * Lewis et al. (2020) Patrick Lewis, Ethan Perez, Aleksandara Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, et al. 2020.  Retrieval-augmented generation for knowledge-intensive nlp tasks.  _arXiv preprint arXiv:2005.11401_.
  * Luan et al. (2019) Sifei Luan, Di Yang, Celeste Barnaby, Koushik Sen, and Satish Chandra. 2019.  Aroma: Code recommendation via structural code search.  _Proceedings of the ACM on Programming Languages_ , 3(OOPSLA):1–28.
  * Raffel et al. (2019) Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. 2019.  Exploring the limits of transfer learning with a unified text-to-text transformer.
  * Song et al. (2019) Kaitao Song, Xu Tan, Tao Qin, Jianfeng Lu, and Tie-Yan Liu. 2019.  Mass: Masked sequence to sequence pre-training for language generation.  _arXiv preprint arXiv:1905.02450_.
  * Tufano et al. (2020) Michele Tufano, Dawn Drain, Alexey Svyatkovskiy, Shao Kun Deng, and Neel Sundaresan. 2020.  Unit test case generation with transformers.  _arXiv preprint arXiv:2009.05617_.
  * Xu et al. (2020) Frank F. Xu, Zhengbao Jiang, Pengcheng Yin, Bogdan Vasilescu, and Graham Neubig. 2020.  Incorporating external knowledge through pre-training for natural language to code generation.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_ , pages 6045–6052, Online. Association for Computational Linguistics.
  * Yao et al. (2018) Ziyu Yao, Daniel S. Weld, Wei-Peng Chen, and Huan Sun. 2018.  Staqc: A systematically mined question-code dataset from stack overflow.  _CoRR_ , abs/1803.09371.
  * Yin et al. (2018) Pengcheng Yin, Bowen Deng, Edgar Chen, Bogdan Vasilescu, and Graham Neubig. 2018\.  Learning to mine aligned code and natural language pairs from stack overflow.  In _2018 IEEE/ACM 15th International Conference on Mining Software Repositories (MSR)_ , pages 476–486. IEEE.