Revisiting Code Search in a Two-Stage Paradigm

TL;DR

This paper introduces TOSS, a two-stage code search framework that combines IR, bi-encoder, and cross-encoder models to improve accuracy and efficiency in code retrieval tasks across multiple languages.

Contribution

TOSS effectively fuses different code search methods in a two-stage process, achieving state-of-the-art accuracy with improved efficiency over existing approaches.

Findings

TOSS outperforms baseline models with an MRR of 0.763.

TOSS is efficient and effective across multiple programming languages.

Compared to six data fusion methods, TOSS achieves superior results.

Abstract

With a good code search engine, developers can reuse existing code snippets and accelerate software development process. Current code search methods can be divided into two categories: traditional information retrieval (IR) based and deep learning (DL) based approaches. DL-based approaches include the cross-encoder paradigm and the bi-encoder paradigm. However, both approaches have certain limitations. The inference of IR-based and bi-encoder models are fast, however, they are not accurate enough; while cross-encoder models can achieve higher search accuracy but consume more time. In this work, we propose TOSS, a two-stage fusion code search framework that can combine the advantages of different code search methods. TOSS first uses IR-based and bi-encoder models to efficiently recall a small number of top-k code candidates, and then uses fine-grained cross-encoders for finer ranking. Furthermore, we conduct extensive experiments on different code candidate volumes and multiple programming languages to verify the effectiveness of TOSS. We also compare TOSS with six data fusion methods. Experimental results show that TOSS is not only efficient, but also achieves state-of-the-art accuracy with an overall mean reciprocal ranking (MRR) score of 0.763, compared to the best baseline result on the CodeSearchNet benchmark of 0.713. Our source code and experimental data are available at: https://github.com/fly-dragon211/TOSS.