如果你是一名软件开发人员,你的浏览器历史记录里可能充满了像“how to reverse a list in Python (如何在 Python 中反转列表) ”或“pandas dataframe drop duplicates (pandas dataframe 去重) ”这样的搜索记录。这种根据自然语言描述找到合适代码片段的过程,被称为代码检索 (Code Retrieval)

为了让 AI 模型擅长这项任务,它们需要海量的训练数据。这些数据由成对的内容组成: 用户查询 (自然语言) 和相应的代码片段 (编程语言) 。但这些数据从何而来?

历史上,我们通常有两种选择:

  1. 爬取 GitHub: 我们获取代码函数,并将其“文档字符串 (docstrings) ” (注释) 视为查询。这种方法通过可扩展性好,但不准确,因为正式的文档字符串看起来一点也不像混乱的用户搜索查询。
  2. 爬取 Stack Overflow: 我们获取用户的问题和被采纳的回答。这里的查询很真实,但代码质量参差不齐。

我们需要一个平衡点: 来自代码仓库的高质量代码,配上真实的、类似用户的查询。在中国科学技术大学的研究论文 “Optimizing Code Retrieval: High-Quality and Scalable Dataset Annotation through Large Language Models” 中,研究人员提出了一种新颖的解决方案。他们开发了一种方法,利用大型语言模型 (LLM) 为现有代码生成合成的、高质量的查询,从而产生了一个名为 Query4Code 的新数据集。

在这篇文章中,我们将探讨为什么 LLM 难以在孤立状态下阅读代码,研究人员如何利用图论解决“缺失上下文”问题,以及这个合成数据集如何超越传统数据源。

错位: 文档字符串 vs. 查询

要理解核心问题,我们首先需要看看模型目前训练所用的数据。大多数开源数据集依赖于文档字符串 (docstrings)

文档字符串是对函数的正式解释,通常详细说明参数和返回类型。而查询则是用户的问题或意图的陈述。

图 1: 代码片段及其对应的查询和文档字符串示例。

图 1 所示,这里存在语义鸿沟。文档字符串描述了函数是如何工作的 (“从 Jupyter notebook 导出内容…”) ,而查询则是问用户想要达到什么目的 (“如何导出内容…”) 。为了构建一个理解用户的搜索引擎,我们需要看起来像“Query (查询) ”部分的数据,而不仅仅是“Docstring (文档字符串) ”部分。

研究人员提出了一个问题: 我们能不能直接让 LLM (比如 GPT-3.5 或 CodeLlama) 看代码并为我们写一个用户查询?

答案是肯定的,但有很大的局限性。LLM 虽然强大,但不是全知全能的。如果你给 LLM 提供一个孤立的函数,它往往无法理解该函数实际上是做什么的。

挑战: 上下文缺口

研究人员进行了初步分析,以了解 LLM 为何无法正确标注代码。他们确定了两个主要罪魁祸首: 库内函数调用 (Intra-repository function calls)第三方 API 调用 (Third-party API calls)

1. 库内函数调用

现实世界的代码是模块化的。函数 A 调用函数 B,函数 B 又调用函数 C。如果你让 LLM 解释函数 A,但不给它看函数 B 的代码,LLM 只能靠猜。

表 1: 库内和第三方库 API 调用数量及比例统计。

表 1 显示了这个问题的规模。在分析的代码库中,近一半 (46.5%) 的函数涉及调用同一代码库中的其他函数。

这种上下文缺失的影响是巨大的。研究人员在有和没有被调用函数上下文的情况下测试了 LLM (GPT-3.5 和 CodeLlama) 。

图 2: 不同数量的库内调用对查询标注质量的影响。

图 2 揭示了结果。蓝色和橙色柱状图代表生成查询的质量得分。请注意,当包含上下文 (被调用函数的代码) 时 (w/ Context) ,质量得分显著跃升。调用链越复杂 (X 轴) ,模型就越依赖额外的上下文。

2. 第三方 API 调用

现代编程涉及将各种库 (如 NumPy、PyTorch 或不起眼的工具库) 粘合在一起。虽然 LLM 在预训练期间见过流行的库,但在面对小众或不流行的 API 时却很吃力。

图 4: 不同流行程度的第三方 API 对 LLM 理解能力的影响。

图 4 展示了“流行度偏差”。X 轴代表库的流行程度。Y 轴是 LLM 的理解得分。正如你所见,对于低流行度的 API (图表左侧) ,模型表现不佳。它们根本不知道 random_niche_lib.do_thing() 到底是做什么的。

解决方案: 上下文感知标注流程

基于这些发现,研究人员提出了一套复杂的流程来生成 Query4Code 数据集。他们不仅仅是把代码扔给 ChatGPT;他们设计了一个系统来提供所需的精确上下文。

图 3: 我们的标注方法概览。(a) 库中的文件。(b) 解析得到的函数调用图。(c) 解析得到的 API 调用及其对应流行度。(d) 基于调用关系和当前 API 调用构建标注上下文。(e) 标注方法流程。

图 3 全面展示了他们的方法。让我们分解一下关键创新。

步骤 1: 解析与调用图

首先,他们将代码库视为一个网络,而不是文件列表。他们解析代码以构建函数调用图 (Function Call Graph) (图 3b) 。这个图谱清晰地映射出了谁调用了谁。

步骤 2: 用于上下文的拓扑排序

这是巧妙之处。如果函数 Train 调用函数 LoadData,在不理解 LoadData 之前你就无法理解 Train

研究人员使用拓扑排序 (Topological Sorting) (论文中的算法 1) 来确定标注顺序。他们从“叶”节点开始——即那些在库内不调用其他任何东西的函数。LLM 首先标注这些函数。

当需要标注更高级别的函数 (如 Train) 时,系统会检索上一步生成的依赖项 (LoadData) 的摘要。这将一个复杂的多层推理问题转化为单层问题。

步骤 3: 处理不流行的 API

对于第三方库,系统会检查 API 调用的流行度 (图 3c) 。

  • 高流行度: 系统假设 LLM 知道它 (例如 torch.mean) 。
  • 低流行度: 系统使用搜索引擎 (DuckDuckGo) 抓取该特定 API 的官方文档。这些文档会被输入到 LLM 的上下文窗口中。

步骤 4: 两阶段生成

研究人员没有直接要求生成查询,而是将任务分解为两个步骤,以减少 LLM 的认知负荷。

  1. 摘要: LLM 首先基于代码 (\(c\)) 生成代码摘要 (\(s\)) 。
  2. 查询生成: 然后 LLM 基于代码和摘要生成用户查询 (\(q\)) 。

Equation 1

这种“思维链 (Chain of Thought) ”方法确保模型在尝试为其编写搜索查询之前,已经完全“理解”了代码。

步骤 5: 反向验证 (过滤器)

LLM 会产生幻觉。为了确保质量,系统包含了一个自我修正机制。生成查询后,系统要求 LLM 反转角色: “给定这个查询,这段代码片段真的能回答它吗?”

Equation 2

模型会给出一个从 0 到 3 的分数 (\(f(q,c)\)) 。

Equation 3

只有得分 1 或 2 (意味着代码满足查询要求) 的配对才会被保留在最终数据集 (\(C_{filtered}\)) 中。这过滤掉了噪声和不相关的代码。

结果: Query4Code

通过在 GitHub 上的高质量 Python 仓库中使用此流程,作者创建了 Query4Code , 包含 23.72 万个查询-代码对。

为了证明该数据集优于之前的数据集,他们进行了广泛的实验。他们在 Query4Code 上预训练了标准模型 (如 CodeBERT 和 UniXcoder) ,并将它们与在 CodeSearchNet (CSN) 上训练的模型进行了比较。

零样本性能

他们在没有任何微调的情况下 (零样本,Zero-Shot) ,在真实基准上测试了模型。这测试了模型的泛化能力。

表 3: 在 Code-SearchNet (CSN) 和 Query4Code (Q4C) 数据集上预训练的代码表示模型的零样本和微调性能比较。

表 3 所示,在 Query4Code (Q4C) 上训练的模型始终优于在 CodeSearchNet (CSN) 上训练的模型。

  • CoSQAWebQueryTest 是源自真实搜索引擎日志的数据集。这里的性能提升非常显著 (例如,CodeBERT 在 CoSQA 上从 56.34 跃升至 59.80) 。
  • 这证实了 LLM 生成的合成查询比标准文档字符串更接近真实用户的搜索行为。

微调性能

研究人员还对模型进行了微调。即使允许模型针对目标任务进行专门学习,从 Query4Code 预训练开始也能提供更好的基础,从而在所有方面获得更高的最终得分 (参见表 3 中的“Fine-Tuning”部分) 。

人工评估

自动化指标固然好,但人类是否认同数据集的质量?研究人员邀请人类专家对生成的配对进行评分。

表 5: 人工评估结果。

表 5 显示了人类评分与模型自我验证评分之间的相关性 (\(r\) 和 \(\tau\)) 。正相关表明自动过滤机制与人类判断高度一致。专家之间的高一致性得分 (0.858) 证实了评估的可靠性。

为什么这很重要: 成本与可扩展性

你可能会想,“用 GPT-4 标注代码听起来很贵。”然而,作者进行了成本分析。

  • 众包: 雇佣人类标注查询-代码对的成本约为 每对 0.20 美元 , 且耗时数分钟。
  • Query4Code 方法: 使用 GPT-3.5-turbo 的成本在 每对 0.001 美元到 0.004 美元之间

这比人工标注便宜近 100 倍 , 同时提供了人类无法比拟的可扩展性。

案例研究: 一个真实的例子

让我们看一个数据集中的具体例子,来看看文档字符串和生成的查询之间的区别。

图 5: 带有文档字符串和标注查询的代码片段示例。

图 5 中,函数 escape_shell_arg 有一个描述参数和类型的文档字符串 ("@param shell_arg", “@type shell_arg: string”) 。

  • 文档字符串: 技术性强,冗长,关注参数。
  • 生成的查询: “Python code for shell argument escaping with single quotes (用于带单引号的 shell 参数转义的 Python 代码) 。”

生成的查询完美地捕捉了意图 。 搜索此函数的用户不会输入“param shell_arg string”;他们会输入与 LLM 生成的内容完全一致的语句。这种语言风格上的细微转变正是现代代码检索模型一直缺失的。

结论

Query4Code 项目表明,我们不一定需要爬取更多数据来构建更好的 AI 工具;有时,我们只需要更好的数据。通过利用 LLM 的推理能力——关键是辅以正确的上下文 (调用图和 API 文档) ——我们可以合成优于从网络上爬取的原生数据的训练数据集。

这种方法弥合了混乱的代码库现实与开发人员自然语言需求之间的鸿沟。虽然这项研究主要集中在 Python 上,但拓扑排序和 API 文档方法可以轻松应用于 Java、C++ 或 TypeScript,为每个人构建更智能、更具上下文感知的编程助手铺平道路。