Retrieval-augmented code completion for local projects using large language models

📄 arXiv: 2408.05026v2 📥 PDF

作者: Marko Hostnik, Marko Robnik-Šikonja

分类: cs.SE, cs.LG

发布日期: 2024-08-09 (更新: 2025-06-15)

备注: 30 pages, 15 figures; Accepted manuscript for Expert Systems with Applications

DOI: 10.1016/j.eswa.2025.128596

💡 一句话要点

提出检索增强的代码补全方法，利用小型LLM在本地项目中实现高效代码补全。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 代码补全 大型语言模型 检索增强生成 本地项目 Jaccard相似度

📋 核心要点

  1. 商业LLM在代码补全中面临隐私和计算资源限制，小型LLM在本地部署和应用成为研究热点。
  2. 论文提出基于Jaccard相似度的上下文RAG方法，以及检索增强的RETRO模型，提升小型LLM的代码补全能力。
  3. 实验表明，上下文RAG方法在代码补全任务中优于RETRO架构，并显著提升了代码补全的性能。

📝 摘要（中文）

大型语言模型（LLM）在软件开发者中日益普及。然而，商业解决方案和LLM的使用存在隐私和计算需求问题。本文着重于使用相对较小且高效的LLM（1.6亿参数），适用于本地执行并能从本地项目进行检索增强。我们训练了两个基于Transformer的开源模型，即生成式的GPT-2和检索自适应的RETRO，使用开源Python文件进行训练，并通过实验比较它们，证实了基于嵌入的检索的优势。此外，我们通过上下文检索增强生成（RAG）来提高模型的性能，该方法使用token的Jaccard相似性来检索代码片段。我们在更大的模型上评估了上下文RAG，并确定尽管其简单性，该方法比使用RETRO架构更合适。实验结果表明，上下文RAG将代码补全基线提高了超过26%，而RETRO将类似大小的GPT-2基线提高了12%。我们强调了适当的tokenization在充分发挥LLM在代码补全中的潜力方面的关键作用。

🔬 方法详解

问题定义：论文旨在解决在本地项目中利用小型语言模型进行高效代码补全的问题。现有商业LLM方案存在隐私泄露风险，且计算资源需求高昂，难以在本地部署。因此，需要一种能够在本地运行、资源消耗低，同时又能提供高质量代码补全建议的解决方案。

核心思路：论文的核心思路是利用检索增强生成（RAG）技术，通过检索本地项目中的相关代码片段，为小型LLM提供上下文信息，从而提高代码补全的准确性和相关性。同时，探索RETRO架构在代码补全任务中的应用，并与RAG方法进行比较。

技术框架：论文主要包含以下几个阶段：1) 使用开源Python文件训练GPT-2和RETRO模型；2) 实现基于Jaccard相似度的上下文RAG方法，用于检索相关代码片段；3) 将检索到的代码片段作为上下文输入到LLM中，生成代码补全建议；4) 评估不同方法在代码补全任务中的性能。

关键创新：论文的关键创新在于：1) 探索了基于Jaccard相似度的简单高效的上下文RAG方法，并证明其在代码补全任务中优于更复杂的RETRO架构；2) 强调了tokenization在代码补全任务中的重要性，并提出了相应的优化策略。

关键设计：上下文RAG方法的关键设计在于使用Jaccard相似度来衡量代码片段之间的相似性，并选择相似度最高的代码片段作为上下文。RETRO模型的关键设计在于引入了检索模块，用于从外部知识库中检索相关信息，并将其融入到生成过程中。论文还关注了tokenization策略的选择，并进行了实验验证。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，基于Jaccard相似度的上下文RAG方法在代码补全任务中取得了显著的性能提升，相比基线模型提升超过26%，优于RETRO架构。同时，RETRO模型也比同等规模的GPT-2模型提升了12%，验证了检索增强的有效性。研究还强调了tokenization在代码补全中的重要作用。

🎯 应用场景

该研究成果可应用于本地代码编辑器、IDE等软件开发工具中，为开发者提供更智能、更高效的代码补全服务。尤其适用于对数据隐私有较高要求的场景，例如企业内部项目开发。未来，可以将该方法扩展到其他编程语言和开发场景，进一步提升代码补全的智能化水平。

📄 摘要（原文）

The use of large language models (LLMs) is becoming increasingly widespread among software developers. However, privacy and computational requirements are problematic with commercial solutions and the use of LLMs. In this work, we focus on using relatively small and efficient LLMs with 160M parameters that are suitable for local execution and augmentation with retrieval from local projects. We train two open transformer-based models, the generative GPT-2 and the retrieval-adapted RETRO, on open-source Python files, and empirically compare them, confirming the benefits of embedding-based retrieval. Furthermore, we improve our models' performance with In-context retrieval-augmented generation (RAG), which retrieves code snippets using the Jaccard similarity of tokens. We evaluate In-context RAG on larger models and determine that, despite its simplicity, the approach is more suitable than using the RETRO architecture. Experimental results indicate that In-context RAG improves the code completion baseline by over 26%, while RETRO improves over the similarly sized GPT-2 baseline by 12%. We highlight the key role of proper tokenization in achieving the full potential of LLMs in code completion.