检索增强辅助代码生成

Code Search Is All You Need? Improving Code Suggestions with Code Search（ICSE 2024）

1. 背景

现代 IDE（集成开发环境）提供代码自动补全和代码生成技术来提高开发效率。当前的方法主要分为基于信息检索（IR-based）和 基于深度学习（DL-based）两类。IR-based 方法（如 BM25 和 Lucene）通过构建索引并基于词匹配进行搜索，具有高效、易部署的优点，但缺乏对代码语义的深入理解。DL-based 方法利用深度神经网络（DNNs）将查询和代码转换为向量表示，并通过计算余弦相似度进行匹配，能够捕捉代码的语义和结构信息，如 CodeBERT 和基于图的检索方法。

尽管代码检索在开发中很常见，如何系统性地将代码检索与代码生成结合仍未被充分研究。因此，该论文探索了一种检索增强的代码建议框架（Retrieval-Augmented Code Suggestion Framework），并通过大量实验验证了其在代码补全和代码生成任务中的有效性。

2. 研究方法

论文提出的检索增强代码建议框架包括三个主要组件：

• Retriever（检索器）：用于从代码库中检索类似代码。
• Formulator（格式化器）：将检索到的代码与代码上下文整合，转化为模型输入。
• Generator（生成器）：基于格式化后的输入，利用语言模型生成代码建议。

Retriever 是框架的核心组件，负责从大规模代码库中检索与当前代码上下文（如方法头或自然语言描述）相似的代码片段。其设计目标是利用代码搜索技术，为后续的代码生成或补全提供参考示例。

论文研究了三种不同的代码检索策略：

1. Header2Code：使用方法头（method header，即方法名、参数列表和返回类型）作为查询，如图中策略1，检索相似代码。
2. NL2Code：使用自然语言（NL）描述作为查询，检索相应代码。如图中策略2所示，使用“fetch data from database”来检索相应的代码片段。
3. NL2NL：使用自然语言描述搜索相似的自然语言描述，然后获取对应代码。即图中策略3所示，从数据库中搜索自然语言描述（comment或函数summary）相近的代码片段。

此外，论文还研究了两种检索方法：

1. IR-based（基于信息检索）：Lucene（BM25算法），支持三种代码检索策略。

2. DL-based（基于深度学习）：

• CodeBERT：预训练模型，支持跨模态检索（NL↔Code），用于NL2Code策略。
• Sentence-BERT：生成文本嵌入，计算语义相似度，用于NL2NL策略。
• ReACC-retriever：基于RoBERTa，专用于Header2Code策略（仅支持Python）。

Formulator将检索到的代码段与源上下文的代码相结合，然后将其作为输入传输到Generator。考虑到一般DL模型和LLM中的差异，论文使用两种不同的方法来制定Generator的输入信息：

1. 对于一般DL模型，使用拼接的方法合并代码和上下文：

2. 对于LLM，论文使用固定模板生成LLM的提示词，

代码生成器（Generator） 负责利用 Formulator（格式化器） 返回的输入，并执行代码建议任务。论文研究了两种不同类型的语言模型（LMs） 作为生成器：

• 通用深度学习模型（General DL models）用于代码补全（Code Completion）：包括Token-Level补全，即生成器接收格式化后的输入（包括检索到的代码片段和代码上下文），然后预测下一个 token；以及Line-Level 补全，即生成器在给定输入的基础上，预测完整的代码语句。使用到的模型有LSTM、Transformer Decoder(TFM)、Transformer-XL (Tr-XL)和CodeGPT。

• 大语言模型（LLMs）用于代码生成：使用Formulator生成的提示词，生成完整的方法体。实验用的LLM是ChatGPT。

3. 实验设置及结果

用于代码检索（Retriever）的数据集

论文选取了两个大规模代码库：

1. CodeMatcher（Java）：
• 该数据集包含 1,048,2463个 Java 方法（Method）。
• 这些方法来自 40,000+ 个 GitHub 仓库，筛选条件是至少 5 颗星。
• 主要用于代码检索（Code Search）任务。
2. PyTorrent（Python）：
• 该数据集包含 2,841,300个 Python 方法。
• 从 PyPI 和 Anaconda上的 220,000个 Python 包中收集。
• 适用于 Python 代码的检索任务。

此外，论文还参考了 CodeSearchNet数据集（NL-to-Code 任务），用于微调 CodeBERT，以提升 NL2Code代码检索策略的效果。

用于代码生成（Generator）的数据集

论文使用 CodeXGLUE数据集进行训练和评估代码生成任务，此外，为了降低代码泄露风险（data leakage），论文采取了一些措施：

• 数据去重：去掉重复的代码样本，避免生成器直接“复制”检索到的代码。
• 避免数据泄露：检测数据集中的相同或相似的GitHub 仓库和Python 包，防止检索到的代码与训练数据高度重叠。
• 预处理：
• 使用 tree-sitter 解析代码结构，提取方法头（Header）和方法注释（Comment）。
• 过滤掉无法正确解析的方法代码。
• 使用 Byte Pair Encoding（BPE）对代码进行分词，避免出现过长的 Token。

RQ1: 代码检索是否能提升基于深度学习（DL）的代码补全性能？

代码补全任务通常依赖于通用深度学习模型（General DL Models），如 LSTM、Transformer-XL 和 CodeGPT。该RQ的研究目标是验证代码检索（Retriever）是否能提高代码补全模型的性能，包括Token-Level 和 Line-Level 补全。

上面两张表格分别展示了在一般DL模型在token-level和line-level上Code Completion任务上的表现，结果表明该检索增强框架在两种层级上都能提升任务的表现，且由于整行代码通常更复杂，代码检索增强的效果更明显。其中NL2NL策略最有效，且框架在Python的表现与提升普遍优于Java，表明Python代码可能更容易通过检索找到类似的代码片段。此外，DL-based 检索优于 IR-based 检索，IR-based 检索（Lucene/BM25）的提升有限，因为它仅基于关键字匹配，无法充分理解代码语义，DL-based 检索方法如CodeBERT 和 Sentence-BERT 在跨模态任务（NL2Code 和 NL2NL）上表现更佳，说明深度学习能更好地捕捉代码语义信息，特别适用于NL2Code 和 NL2NL 检索任务。

RQ2: 代码检索是否能提升基于大语言模型（LLMs）的代码生成性能？

LLMs（如 ChatGPT）在代码生成任务上表现卓越，但是否能通过代码检索增强仍不确定。 该RQ的研究目标是验证代码检索增强对 LLMs 的代码生成性能提升效果。

实验结果表明，LLMs在检索增强后，代码生成能力显著提升，且比一般DL模型的提升更加明显。另外，Python 代码生成比 Java 提升更明显，可能是因为Python代码结构更加统一、检索示例的更高，ChatGPT 可能对 Python 代码的理解比Java更深入。在检索策略和检索方法上，NL2NL的提升效果依然最佳，其优势来源于它能找到 “功能相似但代码不同” 的示例，提高 LLM 生成代码的多样性和准确性。且DL-based的检索方法仍优于IR-based检索。

RQ3: 代码检索增强的 ChatGPT 在不同 Shot Numbers 和 Prompt Templates 下的表现如何？

研究发现，当提供 1-2 个检索示例时，ChatGPT 代码生成的 BLEU-4 和 CodeBLEU 得分显著提升，而当提供 4-8 个示例时，性能反而下降或波动。这可能是因为过多的示例可能会让 LLM 失去生成能力，而倾向于直接复制（Overfitting），或是让LLM关注错误的上下文，而非最相关的部分（Information Overload）。因此，结合一组少量但高质量的类似代码片段可产生更大的改进。

论文还研究了不同的 Prompt 结构 对 ChatGPT 代码生成的影响：

• Detailed（详细的）：提供详细的代码上下文和解释。
• Implicit（隐式的）：减少直接指令，而让 ChatGPT 通过上下文推导答案。
• Two-Step（双步骤的）：分两步引导 LLM 先分析问题，再生成代码。

结果表明，“Detailed, Implicit, Two-Step” 提示能最大程度提高 ChatGPT 代码生成的准确性，过于直接的 Prompt 可能会导致 LLM 仅进行复制，而不是生成新的代码。

4. 总结

代码检索增强可以大幅提升代码建议的性能：

• 在代码补全和代码生成任务上均有显著提升。
• 尤其是结合 ChatGPT，生成代码的 BLEU-4 提高 130.8%。

NL2NL 检索策略效果最好：

• 通过自然语言描述匹配类似代码，能最大程度提升代码生成质量。

IR-based 检索效率更高，DL-based 检索效果更好：

• Lucene（IR-based）检索速度最快，适用于低延迟场景。
• CodeBERT（DL-based）在跨模态检索（NL2Code）时表现更优。

ChatGPT 结合检索后，能更好地模仿开发者行为：

• ChatGPT 通过检索获得上下文示例后，不是简单复制，而是能归纳模式并生成更符合语义的代码。

未来研究方向

• 探索检索增强的更多应用场景（如 Bug 修复、代码翻译）。

• 优化代码检索的搜索策略，提升 NL-to-Code 检索效果。

• 研究如何在低延迟场景中高效结合代码检索与代码生成。