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

基于代码预训练的大型语言模型在完成任意程序函数代码方面表现出色。然而，在大型软件开发项目中，正确的代码补全可能依赖于代码仓库中其他位置定义的API调用或函数。

检索增强生成通过从仓库中获取相关上下文来解决这个问题，增强了模型的理解并改进了输出。但执行检索需要时间并会减慢生成速度：这总是最佳选择吗？

在今年国际机器学习会议上发表的一篇论文中，我们研究了这个问题，发现实际上80%的情况下检索并不会提高代码生成的质量。

方法

为了解决这种效率低下的问题，我们微调了一个大型语言模型来判断检索是否可能有所帮助，并根据答案输出两个特殊标记之一。

为了进行微调，我们使用了一个通过从开源许可证仓库中采样代码、随机掩码代码行并从仓库其他地方检索相关代码构建的数据集。然后我们比较了大型语言模型在有和没有额外上下文情况下对掩码代码的重建效果。这些示例根据检索是否改进了生成效果进行标记。

在实验中，我们发现在代码补全任务上，基于我们数据集微调的代码大型语言模型甚至比总是执行检索的模型表现更好——但由于选择性检索，推理速度提高了70%。

技术实现

创建数据集的所有步骤——采样和掩码代码、检索相关代码、有和没有检索上下文的代码生成——都可以自动化，这使得我们的方法具有自监督性：不需要人工标注，并且可以扩展到任意大的数据集规模。

我们尝试了多种从仓库检索上下文信息的方法，包括使用基于Transformer的语义嵌入来匹配代码序列的UniXCoder，以及使用n-gram数据、语法树和代码流语义的CodeBLEU。然而，这两种方法都没有比更高效的Jaccard相似度表现得更好，Jaccard相似度是两个符号序列交集与并集的比率。因此在大多数实验中，我们使用Jaccard相似度进行检索。

对于模型微调，我们使用了"中间填充"机制，其中掩码代码从代码序列中切除，前面和后面的部分用特殊标记标识。训练目标包括在字符串末尾附加掩码代码的输入字符串，同样用特殊标记标识。这使得模型能够利用掩码代码前后的上下文信息；已被证明比训练模型在前面和后面部分之间插入生成代码产生更好的结果。

在微调过程中，我们有两个训练目标：正确重建缺失代码和准确评估检索信息何时有助于重建。

准确性评估

与现有模型（如StarCoder）相比，我们的方法（称为Repoformer）在各种基准测试（包括RepoEval和针对长格式代码补全的新基准CrossCodeEval）中提高了准确性并减少了推理延迟。

延迟评估

我们在现实的"在线服务"设置中展示了Repoformer减少延迟的能力。我们假设工作仓库已被索引。给定包含当前文件的代码补全请求，系统同时启动三个过程：

• 使用Repoformer做出检索决策；
• 使用代码大型语言模型生成没有跨文件上下文的代码补全；
• 检索跨文件上下文并使用它生成代码补全。

在一系列固定选择阈值范围内，Repoformer的选择性检索能够提高准确性和推理速度。性能在各种阈值设置下也保持稳定。

更有趣的是，Repoformer能够作为即插即用的策略模型，减少各种强大代码大型语言模型作为检索增强生成中生成模型的推理延迟。

Repoformer在检索决策中超过85%的准确率确保上下文检索仅在增加价值时使用。

进一步的分析表明，所提出的策略提高了Repoformer对检索的鲁棒性，减少了有害检索，增加了通过检索改进的实例。

致谢

我们非常感谢Wasi Uddin Ahmad和Dejiao Zhang作为该项目导师的贡献。他们的指导，从项目制定到定期会议中的所有伟大建议，产生了重大影响。我们还要感谢其他合著者和匿名ICML评审提供的宝贵反馈，这些反馈真正帮助改进和完善了这项工作。