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.

## introduction 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 @xcite : comparison of gar between passage retrieval and code search. in passage retrieval, the truth (yellow) is included in the generated content. in code search, despite the generated exemplar code satisfies the description of the query, it exhibits noticeable dissimilarity to the true code. 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 subsequently, the original query and these exemplar codes are synthesized to formulate an augmented query. analogously, the rewritten codes, produced by the llms, are merged with the original code, thereby creating a candidate for retrieval. the example delineated in this figure aligns with the one depicted in fig. 1 . @xcite . 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 @xcite , our research concentrates on the former solution, referred as generation-augmented retrieval (gar) @xcite . 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 @xcite @xcite , 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 @xcite and @xcite 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 @xcite . 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. ## related works 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 @xcite @xcite . as we mentioned before, since programming languages and natural languages share few grammatical rules, these methods often suffer from vocabulary mismatch problems @xcite . then, neural models became popular @xcite @xcite . 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 @xcite @xcite . for instance, codebert @xcite utilizes masked language modeling and replaced token detection. codet5 @xcite focuses on generative tasks through bimodal dual generation. unixcoder @xcite integrates the aforementioned generative and understanding pre-training tasks. codet5+ @xcite 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 @xcite with 12 billion parameters indicates the beginning of the code llm era. meanwhile, there are a number of powerful code llms proposed @xcite , 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 @xcite @xcite . among them, code llama @xcite 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 @xcite . 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 @xcite . 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 generation-augmented retrieval, has been proven effective in both question answering and passage retrieval @xcite @xcite . @xcite is the first work to propose gar in question answering. hyde @xcite evaluates gar in passage retrieval under zero-shot setting. query2doc @xcite 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 @xcite and rouge @xcite , 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. crystal-bleu @xcite 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. ## code style similarity 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 @xcite . 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 @xmath0. 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: to ensure symmetry in this metric, we update code distance in variable naming as: for the measurement of code structure, we simply apply tree edit distance (ted) @xcite 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 c 1 and c 2 , which is: ## experimental setups datasets we evaluate reco across various search scenarios and programming languages: online forum stackoverflow conala @xcite , coding challenge competence apps @xcite , general programming problems mbpp @xcite and mbjp @xcite . 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 @xcite . 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 @xcite 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. ## broader impact 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. ## conclusion 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.