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

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. 1 * Corresponding Author. 1 Our Code and Dataset is available at https://github. com/smsquirrel/queryAnnotation def export_nb(nb_path): exporter = PythonExporter() output, res = exporter.from_filename(nb_path) if 'outputs' in res: for filename, content in res['outputs'].items(): savefile(filename, content) return output Code Query How to export the content of a Jupyter Notebook file. Docstring Export content from a Jupyter notebook file. Parameters: -nb_path : The file path of the Jupyter notebook to be exported.