Hierarchical Context Pruning: Optimizing Real-World Code Completion with Repository-Level Pretrained Code LLMs

📄 arXiv: 2406.18294v2 📥 PDF

作者: Lei Zhang, Yunshui Li, Jiaming Li, Xiaobo Xia, Jiaxi Yang, Run Luo, Minzheng Wang, Longze Chen, Junhao Liu, Min Yang

分类: cs.CL

发布日期: 2024-06-26 (更新: 2024-06-27)

🔗 代码/项目: GITHUB

💡 一句话要点

提出层级上下文剪枝（HCP）策略，优化仓库级预训练代码大模型在真实场景下的代码补全。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 代码补全 代码大语言模型 上下文剪枝 仓库级代码 预训练模型

📋 核心要点

  1. 现有仓库级代码大模型在实际应用中，直接拼接整个代码仓库易超出上下文窗口限制，导致代码补全性能下降。
  2. 论文提出层级上下文剪枝（HCP）策略，在函数级别建模代码仓库，保持拓扑依赖，移除不相关代码，减少输入长度。
  3. 实验表明，HCP策略能显著提升代码补全准确率，同时大幅降低输入长度，适用于多种仓库级代码大模型。

📝 摘要（中文）

最近出现的一些代码大语言模型（Code LLMs）已经在仓库级别的代码数据（Repo-Code LLMs）上进行了预训练，这使得这些模型能够识别仓库结构并利用跨文件信息进行代码补全。然而，在真实的开发场景中，简单地连接整个代码仓库通常会超过这些Repo-Code LLMs的上下文窗口限制，导致性能显著下降。本研究对六个Repo-Code LLMs进行了广泛的初步实验和分析。结果表明，保持文件的拓扑依赖关系并增加补全提示中的代码文件内容可以提高补全准确性；剪除所有依赖文件中函数的具体实现并不会显著降低补全的准确性。基于这些发现，我们提出了一种名为层级上下文剪枝（HCP）的策略，以构建具有高信息量代码内容的补全提示。HCP在函数级别对代码仓库进行建模，保持代码文件之间的拓扑依赖关系，同时删除大量不相关的代码内容，从而显著减少仓库级别代码补全的输入长度。我们将HCP策略应用于六个Repo-Code LLMs的实验中，结果表明，我们提出的方法可以显著提高补全准确性，同时大幅减少输入长度。我们的代码和数据可在https://github.com/Hambaobao/HCP-Coder获取。

🔬 方法详解

问题定义：论文旨在解决仓库级预训练代码大模型在真实代码补全场景中，因上下文窗口限制而导致的性能下降问题。现有方法简单地拼接整个代码仓库作为上下文，但实际仓库的代码量往往超出模型所能处理的范围，导致模型无法有效利用跨文件信息，影响补全准确率。

核心思路：论文的核心思路是，通过对代码仓库进行层级化的上下文剪枝，在保证关键信息（如文件依赖关系和函数签名）的前提下，移除冗余的代码实现细节，从而显著减少输入长度，使模型能够在有限的上下文窗口内更好地利用仓库级信息。

技术框架：HCP策略主要包含以下几个阶段：1) 代码仓库解析：将代码仓库解析为函数级别的抽象表示，提取文件之间的依赖关系。2) 依赖关系维护：构建代码文件之间的拓扑依赖图，确保在剪枝过程中保留关键的依赖关系信息。3) 上下文剪枝：根据预设的策略，移除函数体内的具体实现代码，只保留函数签名和必要的上下文信息。4) 提示构建：将剪枝后的代码上下文与待补全的代码片段拼接，构建最终的补全提示。

关键创新：HCP的关键创新在于其层级化的剪枝策略，它不同于简单的截断或随机采样，而是基于对代码结构的理解，有选择性地保留关键信息，移除冗余信息。这种方法能够在显著减少输入长度的同时，最大限度地保留对代码补全有用的上下文信息。

关键设计：HCP的关键设计包括：1) 函数级别的抽象表示：将代码仓库表示为函数级别的图结构，方便进行依赖关系分析和剪枝操作。2) 拓扑依赖图的构建：确保在剪枝过程中保留关键的文件依赖关系，避免破坏代码的整体结构。3) 剪枝策略的选择：可以根据不同的代码大模型和应用场景，选择不同的剪枝策略，例如，可以根据函数的调用频率或重要性来决定是否保留其实现代码。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，HCP策略在六个Repo-Code LLMs上均能显著提升代码补全准确率，同时大幅减少输入长度。具体而言，HCP能够在保持甚至提升补全准确率的同时，将输入长度减少高达50%以上。例如，在某些模型上，HCP可以将补全准确率提升5-10个百分点，同时将输入长度减少到原来的三分之一。

🎯 应用场景

该研究成果可应用于智能代码补全、代码搜索、代码理解等领域。通过减少输入长度，HCP策略能够提升代码大模型在资源受限环境下的应用能力，例如在移动设备或边缘计算平台上进行代码补全。此外，HCP还可以用于构建更高效的代码索引，加速代码搜索过程，并提高代码理解的准确性。

📄 摘要（原文）

Some recently developed code large language models (Code LLMs) have been pre-trained on repository-level code data (Repo-Code LLMs), enabling these models to recognize repository structures and utilize cross-file information for code completion. However, in real-world development scenarios, simply concatenating the entire code repository often exceeds the context window limits of these Repo-Code LLMs, leading to significant performance degradation. In this study, we conducted extensive preliminary experiments and analyses on six Repo-Code LLMs. The results indicate that maintaining the topological dependencies of files and increasing the code file content in the completion prompts can improve completion accuracy; pruning the specific implementations of functions in all dependent files does not significantly reduce the accuracy of completions. Based on these findings, we proposed a strategy named Hierarchical Context Pruning (HCP) to construct completion prompts with high informational code content. The HCP models the code repository at the function level, maintaining the topological dependencies between code files while removing a large amount of irrelevant code content, significantly reduces the input length for repository-level code completion. We applied the HCP strategy in experiments with six Repo-Code LLMs, and the results demonstrate that our proposed method can significantly enhance completion accuracy while substantially reducing the length of input. Our code and data are available at https://github.com/Hambaobao/HCP-Coder.