蚂蚁CodeFuse代码大模型技术解析：基于全仓库上下文的代码补全

背景

2023年CodeFuse完成了百亿级别的代码大模型从0到1的预训练，配合指令微调、量化部署等一系列配套技术，成功将AI大模型能力应用到多个下游研发场景，助力生产提效。在众多下游产品中，CodeFuse代码补全插件直接触及研发过程中最核心的编码场景，因此对开发效率的影响最显著。目前，CodeFuse代码补全插件是CodeFuse系列产品中用户数量最多、留存率最大，调用AI能力最多的产品。

目前，大部分代码语言模型在预训练阶段以文件为基本单位，随机选择代码文件拼接固定长度后组成训练样本。常见的代码评测数据集（HumanevalX、MBPP）也以单文件为主：基于待补全位置（Task Hole）的前缀（Prefix）使用Left-to-Right方式推理，或者同时使用前缀（Prefix）和后缀（Suffix）采用FIM（Fill In the Middle）方式推理。然而，实际的开发场景通常以代码仓库（Repository）为基本单位。在目前常见的业务设计模式下，大量与当前编辑文件存在依赖关系的内容散落在仓库内的的其他文件中，模型仅使用当前编辑文件内容预测会存在上下文不足的问题，从而导致补全结果存在幻觉、不准确等问题。目前，比较常见的解决思路是使用RAG的方法抽取一定代码片段作为上下文指导模型推理，但此方案同时会带来"上下文-延迟困境"的挑战。即丰富的上下文带来的效果提升和更长的提示内容增加的推理时间之间的权衡，这种权衡在IDE场景中可能会影响用户的实际体验。

为了解决上述问题，本文提出一种仓库级别代码补全框架RepoFuse：通过对实际编程的总结抽象，我们的方法从仓库中抽取了两种关键的跨文件上下文：基于代码相似性分析的相似上下文（Similar Context），用于识别功能相近的代码段；以及语义上下文（Semantic Context），提供类别分类和API交互的语义理解。然而，如此大量的信息可能导致模型的输入过于冗长，影响推理生成的效率。为此，RepoFuse采用了一种基于相关性引导的上下文选择策略（Relevance-Guided Context Selection）指导模型Prompt的构建。这种技术有选择地筛选出与当前任务最相关的上下文，将上下文精炼为简洁的Prompt，既能适应有限的上下文长度，又能确保高完成度的准确性。本文使用常见的开源代码模型在CrossCodeEval的Java和Python数据集上进行了实验，结果表明在完全匹配指标（Exact Match）上RepoFuse有3.01到3.97的提升（与当前主流开源工具的SOAT对比）。

相关工作

目前业界的相关工作主要沿用RAG的思路，代表工作如下表所示。具体来说，每一个方法需要回答以下三个问题：搜什么、怎么搜以及怎么用：

名称	搜什么	怎么搜	怎么用
RLPG	启发式搜索规则	判别模型选择搜索规则	Prompt Engineering
RepoFusion	同RLPG	同RLPG	Fusion-In-Decoder
ReACC	外部知识库相似片段	相似度搜索	Prompt Engineering
RepoCoder	仓库内相似片段	生成 + 相似度迭代搜索	Prompt Engineering
CrossCodeEval	仓库内相似片段	相似度搜索	Prompt Engineering
RepoBench	语义依赖信息	AST