Optimizing Code Retrieval: High-Quality and Scalable Dataset Annotation through Large Language Models (2024)

Abstract: Code retrieval aims to identify code from extensive codebases that semantically aligns with a given query code snippet. Collecting a broad and high-quality set of query and code pairs is crucial to the success of this task. However, existing data collection methods struggle to effectively balance scalability and annotation quality. In this paper, we first analyze the factors influencing the quality of function annotations generated by Large Language Models (LLMs). We find that the invocation of intra-repository functions and third-party APIs plays a significant role. Building on this insight, we propose a novel annotation method that enhances the annotation context by incorporating the content of functions called within the repository and information on third-party API functionalities. Additionally, we integrate LLMs with a novel sorting method to address the multi-level function call relationships within repositories. Furthermore, by applying our proposed method across a range of repositories, we have developed the Query4Code dataset. The quality of this synthesized dataset is validated through both model training and human evaluation, demonstrating high-quality annotations. Moreover, cost analysis confirms the scalability of our annotation method.

## introduction code retrieval aims to find the most relevant code snippet in a database given a user query, facilitating the reuse of programs in the software development process @xcite @xcite and driving recent research on retrieval-augmented code generation @xcite . to achieve good performance in practical applications, the key lies in collecting a wide range of high-quality, dual-modal pairing data between natural language queries and code snippets. an efficient approach to collect code retrieval datasets involves directly gathering code data from online repositories (e.g., github foot_0 ) and processing it to extract code snippets along with their corresponding docstrings. as depicted in figure 1 , since the docstring serves as a description of the function code, it can be utilized as a query. however, a significant difference exists between the docstring and the user's query, resulting in a deviation from queries encountered in real-world scenarios. to bridge this gap and obtain queries that closely resemble those of actual users, some researchers @xcite tend to collect user questions and the corresponding code snippets from programming communities such as stack overflow foot_1 . another approach explored by researchers @xcite involves gathering user search queries from browser logs and subsequently enlisting experts to annotate corresponding code snippets based on these queries. regrettably, the for-mer approach often produces code snippets of inferior quality because of the presence of block and statement-level code within the community. on the other hand, the latter approach allows for the acquisition of a high-quality dataset but proves to be cost-prohibitive and challenging to scale. therefore, we pose a question: can a more efficient, low-cost method be developed to obtain a highquality code retrieval dataset? the formidable capabilities of large language models (llms) present a remarkable opportunity. firstly, previous research @xcite has demonstrated the profound code comprehension ability of llms in various code understanding tasks, such as code summarization @xcite . secondly, existing llms, employing preference alignment techniques @xcite , can generate content that aligns with human preferences. in the domain of search, some studies @xcite have proposed generating the query from the documents, yielding highly promising outcomes. hence, a straightforward approach is to employ llms to generate user-like queries from the code snippets. however, there are some differences between code snippets and traditional documents. for instance, intra-repository function calls refer to the calls between different functions within a repository project, as depicted in figure 1 . function export_nb calls function savefile, which makes it challenging for llms to comprehend function export_nb if only provided as input, without considering the function savefile it calls. additionally, third-party api calls involve invoking functions from external apis, as shown in figure 1 . function export_nb calls the third-party api pythonexporter.from_filename, and llm needs to understand the functionality of this api for a better understanding of the function. in this paper, we first analyze the main factors affecting the quality of annotations for functions in repositories. through preliminary experiments on a development set from 100 selected repositories, we observe that the presence of intra-repository function calls exerts a substantial influence on the quality of annotations, with a greater number of call relationships resulting in a heightened degree of impact. additionally, we uncover that infrequent third-party calls have the greatest impact on annotation quality. this observation may be attributed to the limited pretraining knowledge of llms regarding these external libraries. based on these findings, we propose an annotation algorithm aimed at using llms for high-quality code retrieval query annotations. we start by parsing the relationships of intra-repository function calls and use a topological sorting approach to guide the llm annotation sequence. for third-party function calls, we select third-party functions based on popularity and use web scraping to annotate features of unpopular third-party functions, adding this information to the annotation context. to substantiate the efficacy of our annotation approach, we initially employed our method to obtain a large-scale code retrieval dataset query4code, which includes 237.2k queries and code pairs from 12.3k repositories. we use query4code a pretraining corpus for various code retrieval models. subsequently, comprehensive evaluations on multiple real-world benchmarks confirmed that our method significantly enhances the performance of code retrieval models in real scenarios. ## preliminary analysis the direct use of llms for annotating functions often results in a lack of contextual information about the annotated functions. ## conclusion in this paper, we addressed the trade-off between quality and scalability inherent in the construction methods of previous code retrieval datasets by attempting to generate queries based on large language models (llms). initially, we analyzed the key factors affecting the annotation of queries by llms and identified that both intra-repository function calls and third-party api calls significantly impacted annotation quality. based on this understanding, we had designed an annotation algorithm that constructed appropriate contexts by parsing call relationships to generate function queries. moreover, we had utilized existing code snippets to create the query4code dataset. through model validation and manual assessment, the high quality of the query4code dataset was confirmed, and cost analysis had demonstrated the scalability of our annotation approach.