用抽象语法树分块提升代码生成的准确性和连贯性
cAST: Enhancing Code Retrieval-Augmented Generation with Structural Chunking via Abstract Syntax Tree
- 基于抽象语法树递归分割代码，保持语义完整
- 在RepoEval上召回率提升4.3点，SWE-bench通过率提升2.67点
- 适用于多语言、多任务场景，适合需要高质量代码生成的研究者
检索增强生成（RAG）已成为大规模代码生成的关键技术，通过外部代码库提供上下文以提升生成结果的真实性。然而，现有RAG流程中对文档进行分块的策略——即把代码拆分为可检索单元——仍存在严重问题。传统按行分块的方法常破坏函数结构，导致函数被拆散或无关代码被合并，从而降低生成质量。本文提出一种基于抽象语法树（AST）的结构感知分块方法（cAST），通过递归分解大节点并合并同级节点，在满足大小限制的前提下生成自包含且语义连贯的代码块。该方法在多种编程语言和任务上均表现优异，显著提升代码生成性能：在RepoEval检索任务中，Recall@5提升4.3点；在SWE-bench生成任务中，Pass@1提升2.67点。研究证明了结构感知分块对扩展检索增强型代码智能的重要性。
原文摘要 · Abstract (English)
Retrieval-Augmented Generation (RAG) has become essential for large-scale code generation, grounding predictions in external code corpora to improve actuality. However, a critical yet underexplored aspect of RAG pipelines is chunking -- the process of dividing documents into retrievable units. Existing line-based chunking heuristics often break semantic structures, splitting functions or merging unrelated code, which can degrade generation quality. We propose chunking via Abstract Syntax Trees (\ourwork), a structure-aware method that recursively breaks large AST nodes into smaller chunks and merges sibling nodes while respecting size limits. This approach generates self-contained, semantically coherent units across programming languages and tasks, improving performance on diverse code generation tasks, e.g., boosting Recall@5 by 4.3 points on RepoEval retrieval and Pass@1 by 2.67 points on SWE-bench generation. Our work highlights the importance of structure-aware chunking for scaling retrieval-enhanced code intelligence.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。