Rewriting the Code: A Simple Method for Large Language Model Augmented Code Search (2024)

Abstract: In code search, the Generation-Augmented Retrieval (GAR) framework, which generates exemplar code snippets to augment queries, has emerged as a promising strategy to address the principal challenge of modality misalignment between code snippets and natural language queries, particularly with the demonstrated code generation capabilities of Large Language Models (LLMs). Nevertheless, our preliminary investigations indicate that the improvements conferred by such an LLM-augmented framework are somewhat constrained. This limitation could potentially be ascribed to the fact that the generated codes, albeit functionally accurate, frequently display a pronounced stylistic deviation from the ground truth code in the codebase. In this paper, we extend the foundational GAR framework and propose a simple yet effective method that additionally Rewrites the Code (ReCo) within the codebase for style normalization. Experimental results demonstrate that ReCo significantly boosts retrieval accuracy across sparse (up to 35.7%), zero-shot dense (up to 27.6%), and fine-tuned dense (up to 23.6%) retrieval settings in diverse search scenarios. To further elucidate the advantages of ReCo and stimulate research in code style normalization, we introduce Code Style Similarity, the first metric tailored to quantify stylistic similarities in code. Notably, our empirical findings reveal the inadequacy of existing metrics in capturing stylistic nuances. The source code and data are available at https: //github.com/Alex-HaochenLi/ReCo.

text: Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 1371–1389 August 11-16, 2024 ©2024 Association for Computational Linguistics
In code search, the Generation-Augmented Retrieval (GAR) framework, which generates exemplar code snippets to augment queries, has emerged as a promising strategy to address the principal challenge of modality misalignment between code snippets and natural language queries, particularly with the demonstrated code generation capabilities of Large Language Models (LLMs). Nevertheless, our preliminary investigations indicate that the improvements conferred by such an LLM-augmented framework are somewhat constrained. This limitation could potentially be ascribed to the fact that the generated codes, albeit functionally accurate, frequently display a pronounced stylistic deviation from the ground truth code in the codebase. In this paper, we extend the foundational GAR framework and propose a simple yet effective method that additionally Rewrites the Code (ReCo) within the codebase for style normalization. Experimental results demonstrate that ReCo significantly boosts retrieval accuracy across sparse (up to 35.7%), zero-shot dense (up to 27.6%), and fine-tuned dense (up to 23.6%) retrieval settings in diverse search scenarios. To further elucidate the advantages of ReCo and stimulate research in code style normalization, we introduce Code Style Similarity, the first metric tailored to quantify stylistic similarities in code. Notably, our empirical findings reveal the inadequacy of existing metrics in capturing stylistic nuances. The source code and data are available at https: //github.com/Alex-HaochenLi/ReCo.

heading: 1 Introduction
text: Code search, aimed at retrieving the most semantically relevant code snippets from a codebase according to a specified natural language query, is a common activity that plays an important role in software development (Nie et al., 2016; Shuai et al.,
∗Corresponding author
2020). Retrieving and reusing analogous code fragments from large-scale codebases like GitHub can enhance productivity significantly.
Despite both being sequences of words, matching code queries and natural language queries is challenging as they share few grammatical rules, causing them to fall into two distinct modalities. This grammatical distinction results in limited word overlap, significantly hampering the application of sparse retrieval systems in code search. On the other hand, in dense retrieval systems, the alignment of query and code representations during the training phase assists in alleviating the challenge
1371
(Li et al., 2023). As a result, these systems are capable of encapsulating potential semantic correlations between terminologies employed in programming languages and those in natural languages. However, this potential association becomes challenging to capture if two terminologies rarely manifest together within a query-code pair.
To bridge this gap, one possible solution is to transform the data from one modality to the other. This could involve either generating exemplar codes based on the query or summarizing the functionality of codes in the codebase. Given that natural language queries in code search are often short and ambiguous (Mao et al., 2023; Rahman and Roy, 2021), our research concentrates on the former solution, referred as GenerationAugmented Retrieval (GAR) (Mao et al., 2020). GAR has demonstrated competitive performance in question answering and passage retrieval. In these NLP tasks, a language model is adopted to generate references based on the query to augment it. Similarly, we could use a language model to generate exemplar code snippets that realize the functionalities described in the query. Then the query and exemplar codes are combined to be fed into the retrieval system. With many LLMs demonstrating great intelligence in precisely writing codes (Touvron et al., 2023a,b; OpenAI, 2023b; Zan et al., 2022), performing GAR with LLMs becomes a promising approach for code search.
However, from our preliminary studies, the improvement in performance brought by GAR using LLMs is limited, especially with the high computational cost of LLMs. We argue that answer format influences the performance of GAR on question
answering and code search. In question answering, the correct answer to the question is often unique and can be expressed in limited forms. The generated contents from LLMs, if correct, are usually in the exact same form as the answer. As highlighted in Fig. 1, the matching word “depressive” appears in the reference. On the other hand, code snippets with the same functionality can have diverse formulations, which lowers the chance of matching the code in the codebase, and thus leads to minor improvement of GAR in code search. As shown in Fig. 1, the true code uses Python built-in function Counter to count the number of elements in a list, while the exemplar code snippet does it manually.
To address the mismatch of the generated and ground truth code snippets, we build upon GAR and propose a simple yet effective framework that additionally Rewrites and the Code (ReCo) in the codebase. As shown in Fig. 2, after rewriting, the style of codes in the codebase are normalized by LLMs to align with the exemplar code, thereby facilitating the retrieval. We evaluate ReCo on several code search models across various search scenarios, including coding challenge competence, online programming community, and general programming problems in Python and Java. Experimental results show that ReCo could significantly boost the performance of sparse retrieval systems (up to 35.7%) and dense retrieval systems in both zero-shot (up to 27.6%) and fine-tuning (up to 23.6%) settings.
Furthermore, we propose a novel evaluation metric, dubbed Code Style Similarity, to quantitatively measure the disparity in code style. Our metric validates ReCo’s capability in aligning the style of code
within the codebase with that of code generated by LLMs. Conventional metrics like BLEU (Papineni et al., 2002) and CodeBLEU (Ren et al., 2020) are deemed less appropriate as they calculate similarity based on exact-matched tokens of the given two code snippets. In contrast, Code Style Similarity evaluates style from three distinct perspectives: variable naming, API invocation, and code structure, based on edit distance (Ristad and Yianilos, 1998). Our experiments show that Code Style Similarity exhibits superior explanatory power than existing metrics in measuring the style deviation of code from the dataset and that generated from LLM.

heading: 2 Related Works
text: Code Search Models The development of code search models could be split into three stages. Traditional methods, also denoted as sparse retrieval, employ information retrieval techniques to match words between queries and codes (Hill et al., 2011; Yang and Huang, 2017; Satter and Sakib, 2016). As we mentioned before, since programming languages and natural languages share few grammatical rules, these methods often suffer from vocabulary mismatch problems (McMillan et al., 2011). Then, neural models became popular (Gu et al., 2021; Cambronero et al., 2019; Gu et al., 2018; Husain et al., 2019). They all employ a framework where queries and codes are encoded by neural encoders separately into a joint representation space.
Recently, transformer-based pre-trained models significantly outperformed previous methods, since they can be trained on large-scale unlabelled corpus with self-supervised pre-training tasks. Many novel pre-training tasks are proposed to let models have a better general understanding of codes (Guo et al., 2021; Li et al., 2022b,c; Shi et al., 2022). For instance, CodeBERT (Feng et al., 2020) utilizes masked language modeling and replaced token detection. CodeT5 (Wang et al., 2021) focuses on generative tasks through bimodal dual generation. UniXcoder (Guo et al., 2022) integrates the aforementioned generative and understanding pre-training tasks. CodeT5+ (Wang et al., 2023b) employs the same architecture as CodeT5 and pretrains it with span denoising, causal language modeling, contrastive learning, and text-code matching from both unimodal code data and bimodal codetext data.
Large Language Models As the model parameters and size of training corpora of those transformer-based pre-trained models scale up to billions, they appear to demonstrate remarkable intelligence in understanding and generating codes. As a milestone, Codex (Chen et al., 2021) with 12 billion parameters indicates the beginning of the Code LLM era. Meanwhile, there are a number of powerful Code LLMs proposed (Zan et al., 2022), though most of them are not publicly accessible. Recently, ignited by OpenAI’s ChatGPT (OpenAI, 2023a), a bunch of excellent open-sourced models also contribute to the thriving of Code LLMs (Roziere et al., 2023; Nijkamp et al., 2022; Luo et al., 2023). Among them, Code LLaMA (Roziere et al., 2023) has attracted significant attention because it is a collection of efficient Code LLMs ranging from 7B to 34B parameters. At the same time, some LLMs that are not specifically trained for code exhibit surprising abilities in code intelligence as well, such as GPT3.5 (OpenAI, 2023a) and LLaMA (Touvron et al., 2023a,b). This can be attributed to the inclusion of code snippets in the unlabeled training corpus.
LLMs for Retrieval While LLMs are designed for token generation, their direct application to retrieval tasks such as code search is not suitable. Indeed, there have been attempts to amalgamate the search query and all the candidates together as input, subsequently requesting the LLMs to rank the candidates within the input (Qin et al., 2023). However, the constraint on input sequence length impedes its applicability to large-scale retrieval tasks.
One indirect way is to ask LLMs to generate some references and expand the search query with them. This framework, denoted as GenerationAugmented Retrieval, has been proven effective in both question answering and passage retrieval (Mao et al., 2020; Gao et al., 2022; Wang et al., 2023a). Mao et al. (2020) is the first work to propose GAR in question answering. HyDE (Gao et al., 2022) evaluates GAR in passage retrieval under zero-shot setting. query2doc (Wang et al., 2023a) extends GAR to fine-tuning. Our research findings suggest that the Generation-Augmented Retrieval (GAR) method does not substantially enhance the efficiency of code search, primarily due to the significant stylistic difference between exemplar code and true code.
Code Generation Evaluation Metrics A suitable automatic evaluation metric is vital to the growth of code generation. It is used to measure the lexical similarity between the generated hypothetical code and the true reference code. Initially, metrics such as BLEU (Papineni et al., 2002) and ROUGE (Lin, 2004), originally designed for machine translation, were utilized in the realm of code generation. However, subsequent scholarly discourse posits that these metrics overlook the syntactic and semantic nuances inherent to code. Hence, to consider those features, CodeBLEU (Ren et al., 2020) adds terms that calculate Abstract Syntax Tree similarity and data-flow similarity. CrystalBLEU (Eghbali and Pradel, 2022) sets weights for tokens according to their frequency. They find that high-frequency tokens are often meaningless hence assigning lower weights. These metrics are widely adopted in code generation evaluation, yet they are not suitable for measuring the style difference between two codes due to shared syntactic verbosity.

heading: 3 Methodology

heading: 3.1 Preliminaries
text: Code search aims to retrieve code snippets that are semantically most pertinent to a specified query. Given a search query q and a code snippet c in the fixed codebase, an encoder G is used to map the query and the code to a shared representation space. We calculate the similarity between query and code by dot product, which could be formulated as:
sim(q, c) = ⟨G(q), G(c)⟩ = ⟨vq,vc⟩, (1)
where vq and vc are representation vectors of q and c, respectively. Finally, codes in the codebase are ranked according to the similarity score. Note that in a code search system, code representations vc can be calculated and stored in advance.

heading: 3.2 ReCo
text: Building on GAR, ReCo not only generates exemplar codes based on the query but also rewrites the codes in the codebase.
Generating and Rewriting Code First, we elucidate the process of generating exemplar codes. Given a query q, we employ few-shot prompting (a.k.a in-context learning) (Brown et al., 2020) to generate an exemplar code snippet. The prompt consists of an instruction “Please generate a
Java/Python code snippet according to the given description.” and K randomly sampled query-code pairs from the training set. In this paper, we set K = 4. The instruction and in-context samples are denoted as GEN, enabling us to derive the exemplary code cq as follows:
cq = LLM(q, GEN). (2)
In the procedure of rewriting the code c, we initially summarize the code into a natural language description, represented as qsum. This can be achieved by changing the instruction in GEN to “What is the main purpose of the Java/Python code snippet?”. We denote it as SUM. Subsequently, similar to generating exemplar codes, we consider qsum as the query q1. The entire process leading to the acquisition of the rewritten code, denoted as cc, is as follows:
qsum = LLM(c, SUM), (3)
cc = LLM(qsum, GEN). (4)
Detailed examples are provided in Appendix A to further elucidate the prompt.
Sparse Retrieval Query q and code c are appended with exemplar code cq and rewritten code cc in a sparse retrieval system, respectively. Since we could generate multiple code snippets as augmentation, to retain the original semantics of q and c, we simply repeat them for N times which is equal to the number of augmented codes. Take the query as an example, the augmented search query q+ could be expressed as:
q+ = concat({q}×N, {cq1, cq2, . . . , cqN}). (5)
Similarly, we could get the augmented code c+. In application, q+ is fed to the sparse retrieval system as the search query and c+ are candidates in the codebase.
Dense Retrieval InfoNCE loss (Van den Oord et al., 2018) is widely adopted in fine-tuning because it can pull together the representations between the query and its corresponding code while pushing away the representation of negative codes (Li et al., 2023, 2022a). During training, we take
1The process of rewriting is implemented via a summarizethen-generate approach, as we have observed that merely instructing LLMs to rewrite the original codes does not result in significant alterations.
other in-batch codes as negative samples for a query (Huang et al., 2021). With augmented query q+ and augmented code c+, InfoNCE loss L can be described as:
L = −E [ log
exp(v+qi · v+ci) exp(v+qi · v+ci) + ∑n j ̸=i exp(v + qi · v+cj)
] ,
(6)
where n is the batch size, v+q and v + c are augmented representations of q+ and c+, respectively. For augmented representations, we calculate the expectation of all the generated content according to the chain rule. Take the exemplar code as an example, we have:
E[vcq] = E[G(cq)] = E[G(LLM(q, GEN))]. (7)
Here we assume the distribution of vcq is unimodal since the preferred style of LLM is consistent when generating codes. Then, we employ average pooling between the representation of vcq and vq to get the augmented representation v+q . The total process can be described as:
v+q = 1
2N
 N ·G(q) +
∑
cq∼LLM(q,GEN) G(cq)
  ,
(8) where N is the number of exemplar codes. Similarly, we can get the augmented representation v+c of each code. During evaluation, vq and vc in Eq.(1) are replaced by v+q and v + c .
Theoretical Insights We offer theoretical insights to differentiate GAR and our proposed ReCo. Each code in the codebase is an implementation of a specific query, which we denote as c ∼ P (q). Here P denotes a real-world distribution between queries and codes. LLM also defines a probability distribution over queries. Thus, exemplar codes can be considered to follow cq ∼ LLM(q). We could find that c and cq are sampled from two different distributions given q. This accounts for the occasional divergence between the true code and the exemplar code for the same query, as illustrated in Fig. 1.
The rewritten code follows cc ∼ LLM(qsum). If the query q and code c are identical in semantics and qsum correctly reflect the functionality of code c, we could approximate the distribution of LLM(qsum) as LLM(q). Once the exemplar
code and the rewritten code are both sampled from LLM(q), the expectation of LLM-generated content becomes more similar, which is reflected in the style of generated codes.

heading: 4 Code Style Similarity
text: To quantitatively measure the style difference among codes, we propose a novel evaluation metric, dubbed Code Style Similarity (CSSim). To the best of our knowledge, this is the first metric addressing the similarity between two codes from a stylistic perspective. Indeed, there are several evaluation metrics widely adopted in code generation or translation to measure semantic similarity like BLEU and CodeBLEU. Yet they are not suitable for measuring the style similarity.
The basic idea of these metrics is to compare the predicted code snippet against the ground truth by calculating the intersection of contiguous sequences of code tokens (i.e., n-grams). It is recognized that due to the syntactic verbosity and coding conventions inherent to programming languages, two code snippets frequently share numerous n-grams that are incapable of reflecting their stylistic nuances. Besides, the score is calculated based on the exact match of n-grams, which can be deemed excessively rigid. For example, compared with token_count, word_count is expected to be more stylistically similar to words_count. However, both of them will be assigned a score of 0 under 2-gram match.
CSSim addresses style from three perspectives: variable naming, API invocation, and code structure. Variable naming is generally believed as a reflection of the programmer’s preference. For API invocation, similar APIs often exist in various libraries or packages, the choice of APIs also indicates the preference. As for code structure, sometimes the swap of two lines does not influence the operation hence the order should also be considered. Besides, CSSim is calculated based on a softer measurement, edit distance (Ristad and Yianilos, 1998).
API invocation and variable name follow the same process. Here we take the variable name as an example. We first extract all the variables from the code snippet to get V = {vi}Ni=1. For each variable in the set, we find the most similar variable from the other code and take the edit distance between the two as the similarity of this variable. Then, we take the weighted average value of all the variables
as the style distance in variable naming. The whole process can be described as:
DisV1 = 1 ||λ||1 ∑
vi∈V1 λi min vj∈V2 ED(vi, vj), (9)
where V1 and V2 are extracted variables from two codes and ED denotes Edit Distance. λi is normalized inverse document frequency (IDF) because we intend to decrease the impact of common words. To ensure symmetry in this metric, we update code distance in variable naming as:
DisVar = DisV1 +DisV2
2 . (10)
For the measurement of code structure, we simply apply Tree Edit Distance (TED) (Paaßen, 2018) to the Abstract Syntax Tree transformed from the code. Similar to edit distance, TED quantifies the least amount of basic operations (Insertion, Deletion, and Substitution) required to transform one tree into the other. To calculate CSSim, we first calculate the Code Style Distance CSDis between two codes c1 and c2, which is:
CSDis(c1, c2) = DisVar +DisAPI +TED
3 ,
(11) where DisVar, DisAPI, TED ∈ [0, 1] hence CSDis ∈ [0, 1]. We define CSSim = 1− CSDis.

heading: 5 Experimental Setups
text: Datasets We evaluate ReCo across various search scenarios and programming languages: online forum StackOverflow CoNaLa (Yin et al., 2018), coding challenge competence APPS (Hendrycks et al., 2021), general programming problems MBPP (Austin et al., 2021) and MBJP (Athiwaratkun et al., 2022). The first three datasets are written in Python while the last one is written in Java. The statistics of the datasets are shown in Appendix B.1. We take the widely adopted Mean Reciprocal Rank (MRR) as the evaluation metric (Li et al., 2023). MRR is the average of reciprocal ranks of the true code snippets for the given query.
Baselines We apply ReCo on several models: BM25, an enhanced version of TF-IDF, is a statistical measure that matches certain keywords in codes with the given query. CodeBERT is a bimodal model pre-trained on Masked Language
Modeling and Replaced Token Detection. UniXcoder unifies both generating and understanding pre-training tasks to further enhance code representation learning by leveraging cross-modal contents. Contriever (Izacard et al., 2021) is an unsupervised dense information retrieval model that leverages contrastive learning for its training. CodeT5+ is pre-trained on both unimodal code data and bimodal code-text data with a diverse set of pretraining tasks including span denoising, causal language modeling, contrastive learning, and textcode matching.
Compared Metrics We compare Code Style Similarity with several metrics used for measuring semantic similarity. BLEU measures how many words are shared between the generated and the reference sentence based on the modified n-gram precision. ROUGE-L computes the longest common subsequence of words. CodeBLEU is tailored for code snippets by setting higher weights for programming language keywords and considering data-flow and AST match as well.
Implementation Details When prompting LLMs, we randomly sample 4 in-context examples from the training sets and set a temperature of 1. And when we prompt LLMs multiple times for the
same input, each time we resample the in-context example. The maximum length of output for code summarization and generation is 128 and 256, respectively. For sparse retrieval, we use the default implementation from Pyserini (Lin et al., 2021). For dense retrieval, during training, we adopt the default hyperparameters described in the original paper. They are trained for 10 epochs with a batch size of 32. Experiments are conducted on a Nvidia Tesla A100 GPU. Please refer to Appendix B.3 for more details.

heading: 6 Results
text: Overall Results The results are shown in Table 1. It is worth noting that our experiments encompass the use of various LLMs and multiple instances of both exemplar codes and rewritten codes for conducting ablation studies. Here we report the best performance when equipped with ReCo. Comprehensive results can be found in Appendix D. We can observe that ReCo significantly outperforms GAR on both supervised and unsupervised models across diverse search scenarios. With ReCo, the non-neural model BM25 can have competitive performance compared to neural models under zero-shot setting. And ReCo could boost the performance of zero-shot neural models similar to supervised models. We also evaluate ReCo on Contriever, a passage retrieval model that is not specifically trained for code-related tasks. We argue that ReCo can also benefit general-purpose retrieval models. Note that compared with GAR, ReCo does not bring any additional computation in real-time search because the rewritten code could be pre-processed and stored in the codebase.
Comparison among Evaluation Metrics To demonstrate the superiority of Code Style Similarity, we prove that existing metrics are not effective
Dataset Random w/ the best exemplar code
in measuring the style similarity between two codes by contradiction. If existing metrics are effective, they should satisfy two necessary conditions: 1) the variation of metric scores ∆MetricScore between Metric(cq, cc) and Metric(cq, c) is in the same direction with code search performance gap between ReCo and GAR (∆MRR = MRRReCo− MRRGAR). This is because once the rewritten code is closer to the exemplar code in style, the code search performance should improve accordingly. 2) If we only choose the best one with the highest Metric(cq, c) among multiple exemplar codes, the performance should be significantly better than randomly selecting one exemplar code.
For the first condition, we analyze the results from BM25 on the four datasets with different LLMs including GPT3.5, Code Llama-7B, 13B, and 34B. Here we do not take the results from dense retrieval systems because neural models can capture the potential relationship among similar tokens. The numerical scores under different evaluation metrics are shown in Appendix D. Regression plots are shown in Fig. 3. According to our first condition, ∆MetricScore and ∆MRR should be consistent, which means that points are expected to be scattered on Quadrant I and III. We can see that most of the points in CodeBLEU, ROUGE-L, and BLEU are scattered on Quadrant IV. In other
words, when the rewritten codes are considered to be more similar to the exemplar code by these metrics, the performance of ReCo, on the contrary, drops compared with GAR. Points from Code Style Similarity mostly fall on Quadrant I and III and the regression line nearly passes through the origin.
For the second condition, we analyze the results from BM25 equipped with GAR. For each query, we calculate the metric score between its exemplar codes and the true code in the codebase and then select the one with the highest metric score. The performance on four datasets after selecting the best exemplar code generated by Code Llama-7B is shown in Table 2. We can observe that compared with random selection, the improvement brought by CodeBLEU, ROUGE-L, and BLEU is not significant generally. On the contrary, Code Style Similarity outperforms other settings. To better understand the preference of different evaluation metrics, we also conduct case studies in Appendix C.
In conclusion, our findings indicate that existing metrics for measuring code style similarity fall short when subjected to two contradictory conditions. Conversely, Code Style Similarity (CSSim) demonstrably satisfies these criteria, highlighting its superior effectiveness in quantifying stylistic similarities in code. Furthermore, we observe a clear positive correlation between code style similarity as measured by CSSim and improvement in MRR for code search, thereby validating ReCo’s motivation that style normalization is advantageous.
Using only LLM-generated codes To further demonstrate that the exemplar code and rewritten code are similar in style, we conduct experiments to only use these LLM-generated codes in the retrieval system. In other words, we use exemplar code to retrieve rewritten codes. The results of UniXcoder under fine-tuning and zero-shot settings on four
datasets are shown in Table 3. We can see from the results that only using LLM-generated codes can reach competitive performance compared with additionally using original queries and codes, and even outperform the setting only using original queries and codes, which indicates that the consistent style in exemplar code and rewritten code have made retrieval easier.
Impact of Different LLMs We explore the effect of using different LLMs in ReCo. The results on BM25 are shown in Table 4. The full results including other retrieval models are in Appendix D. GPT3.5’s number of parameters is not released publicly but is estimated at around 175 billion. Generally, we observe that larger LLMs yield greater improvements. However, a decrement in performance is noted with the application of Code Llama-34B. This decrement is attributed to the model’s propensity to generate code not only for the prompted fifth example but also for the initial four in-context examples. Consequently, the generated code is often truncated due to output length limitations.
Impact of Number of Generated Codes We also explore the effect of the numbers of generated exemplar codes and rewritten codes in ReCo. The outcomes of BM25 on CoNaLa and MBPP are depicted in Fig. 4, while a comprehensive compilation of results, inclusive of other retrieval models and datasets, can be found in Appendix D. Our observations indicate a marginal enhancement when LLMs are tasked with generating more codes. We discern that the multiple codes generated exhibit similarities, with minor variations, attributable to the self-consistent style of each LLM. To address this, our future work will investigate controlled prompts that steer LLMs towards generating code with controlled stylistic variations, thereby enhancing the diversity of code generation. Fig. 4 also illustrates that the improvement tends to diminish as the quantity of generated codes increases. In practical applications, it is essential to weigh the
trade-off between performance enhancement and the incremental costs associated with generation.

heading: 7 Broader Impact
text: As we stated, the key motivation behind ReCo is to normalize the code style between exemplar code and original code. Indeed, there exist code normalization methods, but they only focus on superficial formatting such as the usage of indentation and naming convention (e.g., from camelCase to snake_case). In this paper, we discuss code style normalization from a deeper perspective, implementation logic, and preference for variable naming or API invocation. We believe that this task has great potential as it could not only benefit code search but also many other code-related tasks like code review and code translation.
In this paper, we adopt LLMs to achieve the goal of style normalization by first summarizing the code snippet and then generating code based on the summary. This is because we find directly asking LLMs to rewrite the code results in very similar outputs. In the process of summarize-thengenerate, models are expected to have great code intelligence hence there is no loss of information, as described in the theoretical insights of ReCo. Yet we are aware of the huge cost brought by LLMs. To decrease the cost, one promising solution is to train models specifically used for code style normalization. These models are considered to have much fewer parameters since much general knowledge in LLMs is not needed. To push forward the research of code style normalization, we propose a suitable evaluation metric, dubbed Code Style Similarity. In our future work, we plan to train such models to improve the efficiency of ReCo.

heading: 8 Conclusion
text: In this paper, we propose ReCo, an LLMaugmented code search framework built on GAR, that additionally rewrites the code in the code base to normalize the code style between exemplar code and code in the codebase. We evaluate ReCo on several code search models across various search scenarios with different programming languages. Experimental results demonstrate the effectiveness of ReCo by significantly boosting the performance of models. To encourage further research works on code style normalization and explain the effect of ReCo, we propose an evaluation metric Code Style Similarity. In our future work, based on this metric, we may develop new models that can more efficiently normalize the code style.

heading: Acknowledgement
text: We thank the anonymous reviewers for their helpful comments and suggestions.
Limitations
There are mainly two limitations of this work. First, although ReCo does not require any additional computation in real-time search compared with GAR, both GAR and ReCo rely on the real-time generation of exemplar codes. Therefore, ReCo and GAR may have limitations when applied to tasks that demand low latency. The latency of generating exemplar codes depends on the time cost of LLM inference. As stated in research works focusing on GAR, over the years the cost of hardware has decreased a lot and there are many works proposed to improve the inference efficiency of LLMs (Gao et al., 2022; Wang et al., 2023a). We believe the efficiency problem of GAR and ReCo will be addressed in the future. The second limitation is that we do not evaluate ReCo on some extremely large-scale codebases like CodeSearchNet (Husain et al., 2019). This is due to the time burden of generating exemplar codes and rewriting codes. For example, according to our estimation, there are 1,005,474 queries in total in CodeSearchNet hence generating one exemplar code for them costs more than two months. To address this limitation, we evaluate ReCo on several search scenarios covering coding challenge competence, online programming community, and general programming problems to show the effectiveness of ReCo.

heading: A Complete Prompt
text: Table 11 and Table 12 are two complete prompt examples for generating exemplar codes and summarizing original codes, respectively. Note that in the second step of rewriting original codes, we also adopt the prompt structure of generating exemplar codes but replace the description at last with the summary.

heading: B Experiment Settings
text: B.1 Dataset Statistics
The dataset statistics are shown in Table 5. The numbers here are pairs of queries and their true code snippet. Code search models are asked to distinguish the correct code from the codes from other pairs. Note that in the original APPS dataset, there are 4,284 and 3,515 pairs in the training and test set, respectively. Due to the huge cost of prompting LLMs, we randomly sample a subset for our evaluation.
B.2 MRR Calculation
MRR is the average of reciprocal ranks of the true code snippets for the given query, which could be calculated as:
MRR = 1
|Q|
|Q|∑
i=1
1
Ranki (12)
where Ranki is the rank of the true code for the i-th given query Q.
B.3 Implementation Details For neural models, they all use the same set of hyperparameters. The maximum input length of codes and queries are both set to be 256. Models are trained by Adam and learning rate is set to 5e-6. We adopt mean pooling to get the representation of the whole input sentence to make sure the pooling mechanism is consistent with that during pre-training. The representations are normalized by the L2 norm.
UniXcoder is initialized using the publicly available checkpoint at https://huggingface. co/microsoft/unixcoder-base, Contriever is initialized using https://huggingface.co/ facebook/contriever-msmarco, and CodeT5+ is initialized using https://huggingface. co/Salesforce/codet5p-110m-embedding. CodeBERT is initialized using https: //huggingface.co/microsoft/codebert-base and then pre-trained on the CodeSearchNet dataset (Husain et al., 2019) for 10 epochs. The pre-training setting is the same as in fine-tuning. All the experiments involving model training are running with 3 random seeds 1234, 12345, and 123456 and they all meet p < 0.01 of significance tests.

heading: C Case Study
text: We also conduct a case study to show the superiority of Code Style Similarity. Fig. 5 shows two exemplar codes generated by Code Llama-7B. The true code aims to find the first duplicate element in a given array of integers. Although both the two exemplar codes satisfy the description, their implementation style is different. The first exemplar code is preferred by CodeBLEU, ROUGE-L, and BLEU while the second one is preferred by Code Style Similarity. The true code uses a set to collect seen elements when traversing the list, which is also the logic in the second exemplar code. The first exemplar code implements the function in a different way by checking whether my_list[i] appears in my_list[i+1:]. We think the first code
is preferred by existing metrics because lines 2-4 in the first exemplar code are very similar to lines 4-6 in the true code, which contributes a lot to the metric score.

heading: D Full Results
text: In this section, we report the full experimental results. The results of MetricScore(cq, cc) and MetricScore(cq, c) under different evaluation metrics are shown in Table 6. The results of BM25 are shown in Table 7. The results of fine-tuned UniXcoder and CodeBERT are shown in Table 8 and Table 9, respectively. The results of UniXcoder, Contriever, and CodeT5+ under zero-shot setting are shown in Table 10.

heading: Future Work
text: We believe the efficiency problem of GAR and ReCo will be addressed in the future. The second limitation is that we do not evaluate ReCo on some extremely large-scale codebases like CodeSearchNet (Husain et al. , 2019). This is due to the time burden of generating exemplar codes and rewriting codes. For example, according to our estimation, there are 1,005,474 queries in total in CodeSearchNet hence generating one exemplar code for them costs more than two months. To address this limitation, we evaluate ReCo on several search scenarios covering coding challenge competence, online programming community, and general programming problems to show the effectiveness of ReCo