选择性检索增强代码库级补全技术

通过选择性检索增强代码库级代码补全

基于大型语言模型的代码补全技术虽然在通用函数补全方面表现优异，但在大型软件开发项目中面临挑战——正确的代码补全可能需要依赖代码库中其他位置定义的API调用或函数。检索增强生成（RAG）通过从代码库中获取相关上下文来解决这一问题，但检索过程会拖慢生成速度。

在国际机器学习会议（ICML 2024）上发表的研究表明，80%的情况下检索并不能提高代码生成质量。为此，研究团队开发了一种自监督方法，通过微调语言模型来智能决定何时需要检索操作。

方法创新

自监督数据构建

通过从开源代码库采样代码、随机掩码代码行，并从代码库其他位置检索相关代码，构建训练数据集。通过对比使用和不使用额外上下文时的代码重建效果，标注样本是否因检索而改善。

检索策略优化

实验发现，基于Jaccard相似度的检索方法在效率上优于基于Transformer语义嵌入的UniXCoder和基于多维度匹配的CodeBLEU。该方法通过计算符号序列交集与并集的比例来实现高效检索。

模型训练机制

采用"中间填充"训练机制：将掩码代码从序列中切除，使用特殊标记标识前后段落，训练目标是在字符串末尾重建被掩码的代码。这种机制允许模型同时利用前后文信息，相比直接在中间插入代码的方式效果更佳。

性能表现

准确性提升

在RepoEval和CrossCodeEval等基准测试中，该方法（命名为Repoformer）在精确匹配（EM）、编辑相似度（ES）和单元测试通过率（UT）等指标上均优于StarCoder等现有模型。

延迟优化

在实际在线服务场景中，系统并行执行三个过程：使用Repoformer做出检索决策、不使用跨文件上下文生成代码、检索跨文件上下文并生成代码。选择性检索策略使推理速度提升70%，同时保持准确性。

泛化能力

Repoformer可作为即插即用的策略模型，有效降低各种强代码LLM在RAG中的推理延迟。检索决策准确率超过85%，确保仅在需要时进行上下文检索。

技术特点

• 自监督学习：无需人工标注，可扩展到任意规模的数据集
• 选择性检索：智能判断检索必要性，避免不必要的计算开销
• 结构感知：未来计划采用结构感知分块技术进一步提升语义检索性能

该方法显著提高了代码补全的效率和准确性，为大型代码库的智能开发工具提供了新的技术方向。
更多精彩内容 请关注我的个人公众号 公众号（办公AI智能小助手）或者 我的个人博客 https://blog.qife122.com/
公众号二维码