我正在参加「掘金·启航计划」
导语
本文介绍了一种名为 RepoCoder 的仓库级代码完成框架,它通过利用仓库级信息,结合相似度检索器和预训练的代码语言模型,提高了代码完成的准确性和效率。同时,该文还提出了一个新的基准 RepoEval,用于评估代码完成的性能。实验结果表明,RepoCoder 在各种设置下都显著提高了代码完成的基线,并且比传统的检索增强代码完成方法更有效。
- 会议:Arxiv 2023
- 链接:arxiv.org/abs/2303.12…
摘要
仓库级代码补全任务是基于更广泛的仓库上下文继续编写未完成的代码。而对于自动化代码完成工具来说,利用分散在不同文件中的有用信息是困难的。本文提出了RepoCoder以应对这一挑战。它通过整合基于相似度的检索器和预训练的代码语言模型,简化了仓库级代码完成过程,允许有效利用仓库级信息进行代码完成,并具有在不同粒度级别生成代码的能力。此外,RepoCoder 还利用了一种新颖的迭代检索生成范式,弥合了检索上下文和预期完成目标之间的差距。作者还提出了一个新的基准 RepoEval,它包括最新的高质量实际仓库,涵盖了行、API 调用和函数体补全场景。通过使用不同的代码检索器和生成器的组合来测试 RepoCoder 的性能。实验结果表明,在所有设置中,RepoCoder 在zero-shot 代码补全基线上显著提高了超过 10%,并始终优于普通的检索增强代码完成方法。此外,通过全面分析验证了 RepoCoder 的有效性,并为未来的研究提供了有价值的见解。
1 简介
实际生产中,开发人员在编程过程中需要注意仓库中的其他文件,因此自动化工具也需要考虑仓库中更广泛的上下文以完成未完成的代码,这提出了具有挑战性的仓库级代码补全任务。仓库中的代码文件通常具有相互关联的依赖关系,这种模块化和自定义方式会给自动化代码完成工具带来困难。本文提出了 RepoCoder,这是一个简单、通用且有效的框架,用于仓库级代码完成任务。RepoCoder利用基于相似度的检索器简化了代码补全过程,消除了手动工程和代码分析工具的工作。此外,该框架将仓库级信息直接整合到预训练语言模型的提示中,使其具有推广到不同领域的仓库的能力,无需进行模型调整。同时,使用语言模型还便于生成更高质量和更大粒度的代码,例如完成行、API调用和函数主体。
RepoCoder采用一种新方法弥合检索上下文和预期完成目标之间的差距。如图1所示,该方法涉及迭代检索生成范式,首先只使用未完成的代码执行代码检索并生成中间完成状态。接下来,使用中间代码执行第二次检索迭代。这样可以有效地捕捉输入上下文和预期完成目标之间的关系,从而在仓库中实现更准确的检索。图2是一个例子,说明了设计的合理性,其中模型预测调用 COLMAP() API 的语句,但API名称和参数不正确。通过使用预测语句重新从仓库中进行检索,我们可以获取正确的API签名并成功完成代码。
此外,本文提出了 RepoEval 基准,用于评估仓库级代码完成任务,它使用最新的高质量仓库从Github构建。RepoEval缓解了仓库级场景中缺乏已建立基准的问题,并首次涵盖了三个代码完成细粒度级别:行、API调用和函数主体。我们还利用仓库中的单元测试进行评估,以克服相似度度量的不准确性。实验结果表明,RepoCoder 在所有设置下都显著提高了zero-shot代码补全性能超过了10%。而迭代检索-生成框架仅使用一次检索-生成迭代也始终提高了基于检索的生成的性能。作者还对 RepoCoder 的有效性和局限性进行了全面分析,提供了有价值的未来研究洞见。主要贡献总结如下:
- 提出了RepoCoder,这是一个简单但有效的框架,用于仓库级代码完成任务。
- 提出了一种新颖的迭代检索-生成范式,来弥合检索上下文和预期完成目标之间的差距。
- 创建了RepoEval基准,用于评估仓库级代码补全任务,提供不同粒度级别的评估和基于单元测试的更准确的评估。
- 实验表明,RepoCoder显著优于zero-shot代码完成范式,并提高了基于检索的生成的性能。
2 方法
2.1 整体框架
RepoCoder实际上结合了两次迭代,它首先对输入(待补全的代码段)X到数据库中进行检索(假设检索器为R),得到第一次检索的结果C1=R(X),然后将<C1,X>输入给语言模型M,得到生成的内容Y1=M(C1,X)。之后,将Y1拼接在X后面,然后前向截断超出总长的内容(可以认为是丢弃了X前面的内容,只保留了X尾部的部分内容)再次进行一次检索得到第二次检索的结果C2=R(X,Y1)。最后,将第二次检索的内容C2和X进行拼接,然后输入语言模型M得到最终的输出。在这整个过程中,使用到了两个模块retrieval-augmented generation和generation-augmented retrieval。一个简单的示意图如下图所示:
2.2 Generation-Augmented Retrieval
RepoCoder框架中使用的检索器可以是任何能够在给定查询时搜索相关文档的模型。RepoCoder使用滑动窗口构建检索数据库,该窗口迭代地扫描仓库中的文件,并提取适合窗口大小Sw的连续代码行。滑动窗口在每次迭代中移动固定数量的行,即滑动大小Ss(注意:Sw > Ss)。
在第一次检索迭代中,将使用未完成代码X的最后Sw行来构建查询,并检索出最相似的代码片段C1 = R(X)。然而,检索上下文与预期完成目标之间存在差距,因为检索是基于X进行的,但预期目标是续写X。其中一种可能的解决方案是通过将每个代码片段向下移动Ss行来修改C1。虽然这种方法在先前的研究中被证明是有效的,但是不加区分地移动所有检索到的代码片段并不总是适当的。为了缓解这个问题,本文提出了生成-检索范式,即通过生成的代码Y1来增强检索查询。尽管对于新仓库缺乏定制信息,但预训练的代码语言模型已经展示了出色的通用领域理解和生成能力。生成的代码Y1尽管不确定其正确性,但可以作为检索过程的有价值补充信息。因此,在第二次检索迭代中,查询是通过连接X的最后(Sw-Ss)行和Y1的前Ss行来构建的。最终,可以获得基于Y1的检索结果C2 = R(X,Y1)。
2.3 Retrieval-Augmented Generation
RepoCoder框架中使用的生成器可以是任何能够预测给定特定提示的后续标记的预训练语言模型。将存储库中的全局上下文和目标文件中的局部上下文纳入代码补全至关重要。这将为模型提供基础信息并帮助其推广到未见过的存储库。具体而言,RepoCoder从存储库中检索最相关的代码示例C,并将它们与待补全的代码X连接起来。为了便于阅读和理解,创建了一个提示模板,无缝集成X和C,如图3所示。检索到的代码片段按其与查询的相似性得分升序呈现(因为截断是从前面截断的,故采用升序)。每个代码片段都带有其原始文件路径,并且包含在提示中的代码片段的最大数量K取决于可用的提示长度。最终,提示包含了完成代码所需的本地和全局信息。
3 Benchmark构建
本文提出了一个新颖的RepoEval基准,以评估存储库级别场景下的代码完成工具。通过使用来自Github的最新高质量存储库创建了基准,并涵盖了三个级别的代码完成粒度:行、API调用和函数体。作者还利用存储库中的单元测试来评估已完成函数的正确性,而不仅仅使用基于相似性的度量标准。
RepoEval从Github上精心挑选了Python存储库,满足以下标准:具有开源许可证;创建于2022年1月1日之后;非分支原始存储库;超过100个星标;Python编写的文件占总数的80%以上;具有明确的单元测试。此外,我们随机选择存储库以防止潜在的偏见,并构建了包括行、API调用和函数体完成的三个不同数据集。有关所选存储库的更多信息,请参见表1。
4 实验设置
作者使用了两种类型的检索器:离散的Bag-of-words模型和OpenAI提供的embedding模型Text-Embedding-Ada-002。对于生成器,作者采用了Code-Davinci-002和三种不同size的CodeGen模型(350M、2B、6B)。
对于不同的方法,考虑下面的四种设定:
- Zero-shot;不检索,直接生成;
- RG-1:即只检索1次就进行生成;
- RepoCoder:经过完整的Generation-Augmented Retrieval和Retrieval-Augmented Generation;
- Oracle:将Ground-truth放到Retrieval-Augmented Generation模块进行检索得到最终的生成,这是一种理论上限;
作者采用了多种类型的评估指标,分别基于相似度的Exact Match (EM)和Edit Similarity (ES),和基于执行准确率的Pass Rate (PR)(其中PR只用于评估函数体级别的数据集)。
5 实验结果
5.1 行和API补全数据集
表2a,b分别展示了使用不同检索器得到的实验性能表现,与零-shot基线、神谕方法和基准检索增强方法相比,RepoCoder在行和API完成数据集上始终显著提高了性能。具有350M参数的CodeGen模型,当与RepoCoder集成时,甚至优于zero-shot Code-Davinci-002模型。而且稀疏检索器与密集检索器具有相同的性能。
5.2 函数体补全数据集
RepoCoder在函数体数据集上的表现与行和API上类似。在大多数存储库中,RepoCoder相对于zero-shot基线有显著的提高,并且比仅使用一次检索-生成迭代的方法有更好的表现。此外,RepoCoder的表现接近于Oracle方法,再次展示了该方法的有效性。
6 分析
6.1 更多轮次的迭代
作者还实验了增加更多轮次的检索-生成步骤,得到的实验表现如表4所示(RepoCoder实际上就是2次迭代),可以看到,在使用四个语言模型和稀疏检索器的情况下,扩展的RG迭代的RepoCoder始终优于单迭代检索增强方法。然而,随着RG迭代次数的增加,性能并不会继续提高。
6.2 检索得到的代码片段的统计值
作者对检索增强中的关键因素进行了研究。在API补全数据集上的实验中发现,较高的行重叠与更好的性能呈正相关,而较高的重叠率通常表明存储库中存在更相关和常见的代码片段。此外,大多数检索到的代码片段与完整答案之间的token重叠率超过了70%,RepoCoder的表现与完整答案的重叠率之间存在正相关关系,这些结果强调了检索过程的有效性。
6.3 上下文位置对检索的影响
本文对检索增强中的上下文位置进行了研究。通过定位检索到的代码片段的原始源,研究结果表明(表5),大多数代码片段位于定义的五个类别中,其中绝大多数来自于在方案中归类为“相似导入”、“相似名称”或“当前目录”的文件。此外,与行完成相比,检索到的用于API完成的代码片段更有可能来自于我们方案中分类的文件,表明在API完成情景中需要更多地获取其他文件中的信息。
6.4 仓库中的代码重复的影响
作者对代码重复率和RepoCoder的性能之间的关系进行了研究。通过计算存储库中的重复代码行数与总代码行数的比率,结果表明RepoCoder的表现与代码重复率之间存在相关性,但其相关性并不绝对。即使存储库具有相似的重复率,存在不同的性能收益。这可能是由于作者基于计数相同行的粗略估计重复率,或者是由于语言模型能够从上下文中学习并即使没有相似的代码示例也能表现良好的能力所致。
6.5 误差分析
作者进一步探讨了RepoCoder的有效性和局限性。实验结果表明,在使用zero-shot、RG-1、RG-2和Oracle方法时,Exact Match(EM)和Pass Rate(PR)指标均得到了改进,但是更高级别方法仍然有失败的情况。通过手动案例研究,发现大部分失败是由于误导性的代码检索导致的,同时模型的预测并不总是适合检索。此外,语言模型对给定提示敏感,同时EM分数评估可能存在误差。因此,需要利用单元测试来考虑代码的实际功能而不仅仅依赖于完全匹配的需求。
7 相关工作
Global Context in Code Completion
在软件存储库中整合全局上下文是代码补全工具长期以来的挑战。传统的代码完成技术通常通过分析代码来识别候选建议,然后对它们进行重新排序,但这种方法的灵活性不足以生成任意粒度的代码。另一种研究方法将代码完成视为语言建模任务,其中下一个标记是基于给定上下文生成的。已经提出了各种方法来将全局上下文纳入语言模型的训练中,包括n-gram、LSTMs和Transformers。虽然收集标注数据和调整模型以适用于各种应用程序的过程仍然很昂贵。近年来,大型预训练语言模型在代码完成任务中引起了很大关注。Shrivastava等人(2022)的一项研究也探索了存储库级别的代码完成情景。尽管该研究具有创新性,但仍依赖于不灵活的启发式算法和分类器训练来构建提示。这凸显了利用预训练模型进行代码完成的持续挑战,需要继续进行研究。
Joint Modeling Retrieval and Generation
尽管预训练语言模型具有非凡的能力,但其离线训练范式可能导致缺乏定制和最新信息。最近的研究开始探索在知识密集型任务中联合检索和生成的建模,例如问答和对话生成。这一范式也被扩展到代码生成中,将检索到的文档或代码示例纳入生成过程中。随着语言模型变得越来越复杂,越来越多的研究趋向于在上下文中联合检索和生成,将预训练语言模型视为冻结的黑匣子。此外,一些研究调查了利用模型的预测作为补充上下文来指导检索过程的方法。在这项工作中,我们展示了将代码检索和生成联合在迭代范式中,可以作为存储库级别代码完成任务的有效方法。
8 总结与未来工作
本文提出了RepoCoder,这是一个简单而有效的存储库级别代码补全框架。利用基于相似度的检索器和预训练语言模型,RepoCoder充分利用了存储库级别的信息。通过迭代的检索和生成,RepoCoder可以弥合检索上下文和目标之间的差距,从而提高代码完成性能。通过在RepoEval基准测试上进行了严格的实验表明,RepoCoder在zero-shot代码补全性能方面始终显著改善,并优于普通的检索增强生成方法。通过全面的分析,作者还提供了有关RepoCoder合理性和局限性的有价值的见解。由于其简单性、通用性和有效性,RepoCoder有潜力成为实际软件开发中的重要工具。在未来的工作中,一个目标是解决RepoCoder目前的局限性,继续提高其可用性和鲁棒性。