Enhancing Domain-Specific Code Completion via Collaborative Inference with Large and Small Language Models

Large language model-based code completion has demonstrated excellent performance, but still encounters challenges in capturing domain-specific knowledge for more precise completion within specific domains, i.e., domain-specific code completion. Prior work has studied fine-tuning techniques or retrieval-augmented techniques for this task. Nevertheless, it requires a lot of computational resources to fine-tune large language models (LLMs), and the cost can increase quadratically with the model size. Retrieval-augmented techniques face difficulties in accurately and adaptively retrieving relevant information. Moreover, considering that code completion tools work in real time, how to utilize large language models more efficiently poses challenges. To tackle these challenges, in this paper, we first conduct preliminary experiments and observe that the code completion results of a small model fine-tuned within a specific domain complement those of a large model. Building on this insight, we propose a collaborative framework to effectively combine large and small models for better domain-specific code completion. Specifically, we fine-tune a small code model instead of a large model with the PEFT method, reducing the overhead of fine-tuning. We utilize a well-designed classifier to facilitate the adaptive combination of distinct completion results. The classifier relies on features in various dimensions, such as the similarity between the completed code and the context, and is used to adaptively determine how to combine the tokens predicted by the large and small models for better code completion. Evaluation results show that our approach achieves an average improvement in the exact match of 7.42% and 4.67% over the state-of-the-art baselines in the intra-project and intra-domain code completion scenarios, respectively. Furthermore, compared to the state-of-the-art domain-specific code completion approach FT2Ra, the inference speed of our approach is 1.40 times faster, and the average space requirement drops from 25.98G to 13.69G. These advantages make our approach much more accessible and efficient.