CoRNStack: High-Quality Contrastive Data for Better Code Retrieval and Reranking

TL;DR

CoRNStack is a large, high-quality contrastive dataset for code that improves code retrieval and reranking, leading to better performance in complex real-world software tasks like bug localization.

Contribution

We introduce CoRNStack, a curated contrastive dataset for code, enabling state-of-the-art retrieval and reranking models across multiple programming languages.

Findings

State-of-the-art performance in code retrieval tasks.

Significant improvement in code reranking quality.

Enhanced bug localization in GitHub repositories.

Abstract

Effective code retrieval plays a crucial role in advancing code generation, bug fixing, and software maintenance, particularly as software systems increase in complexity. While current code embedding models have demonstrated promise in retrieving code snippets for small-scale, well-defined tasks, they often underperform in more demanding real-world applications such as bug localization within GitHub repositories. We hypothesize that a key issue is their reliance on noisy and inconsistent datasets for training, which impedes their ability to generalize to more complex retrieval scenarios. To address these limitations, we introduce CoRNStack, a large-scale, high-quality contrastive training dataset for code that spans multiple programming languages. This dataset is curated using consistency filtering to eliminate noisy positives and is further enriched with mined hard negatives, thereby facilitating more effective learning. We demonstrate that contrastive training of embedding models using CoRNStack leads to state-of-the-art performance across a variety of code retrieval tasks. Furthermore, the dataset can be leveraged for training code reranking models, a largely underexplored area compared to text reranking. Our finetuned code reranking model significantly improves the ranking quality over the retrieved results. Finally, by employing our code retriever and reranker together, we demonstrate significant improvements in function localization for GitHub issues, an important component of real-world software development.