Specialized model initialization and architecture optimization for few-shot code search

Context:

Code search aims to find relevant code snippets from a codebase given a natural language query. It not only boosts developer efficiency but also improves the performance of tasks such as code generation and program repair, thus becoming one of the crucial tasks in software engineering.

Objective:

However, recent works are mainly designed for mainstream programming languages with abundant training data. We aim to address the challenges of code search for domain-specific programming languages with limited training data by proposing a novel two-stage, few-shot code search framework named SMIAO.

Method:

SMIAO includes a specialized model initialization and an architecture optimization stage. In the first stage, we first quantitatively identify a mainstream programming language’s dataset that is semantically closest to a target few-shot programming language. Then, we enrich the dataset with hard samples and train an Adapter-GraphCodeBERT model to obtain well-initialized parameters. In the second stage, we first design a search space for the initialized Adapter-GraphCodeBERT model. Then, we employ neural architecture search to optimize the Adapter modules’ positions and quantities in the GraphCodeBERT layers, tailoring for real-world few-shot code search tasks.

Results:

We conduct experiments on a publicly available dataset to demonstrate the effectiveness and rationality of SMIAO. The experimental results show that SMIAO outperforms other state-of-the-art baselines.

Conclusion:

Using mainstream languages’ datasets to initialize Adapter-GraphCodeBERT models, followed by adjusting the quantities and positions of Adapter modules within the GraphCodeBERT layers by neural architecture search, can effectively improve the performance of few-shot code search tasks.