Facebook发布基于NLP的代码搜索工具

来源 | Facebook AI Blog

作者 | Sonia Kim, Hongyu Li, and Satish Chandra

译者 | 核子可乐

编辑 | 陈思

导读：事实证明，如果工程师们能够轻松找到代码示例来指导自己完成特定的编码任务时，他们的工作效率与成果都会更令人满意。对于某些问题，例如“如何以编程方式关闭或隐藏 Android 系统的软键盘？”等，则完全可以从 Satck Overflow 等流行资源库中获取信息。但是，某些专有代码或者 API（或者是使用那些不太常见的编程语言编写的代码）中往往存在一些极其小众的问题，而工程师们也很难从主流论坛中找到答案。

为了满足这一需求，Facebook 开发出一种代码搜索工具，能够将自然语言处理（NLP）与信息检索（IR）技术直接应用于源代码文本。他们将这款工具命名为神经代码搜索（NCS），其能够接收自然语言查询并直接返回从码语料库中检索到的相关代码片段。

我们设置了一个前提，即随着众多大型代码库的上线，与开发人员所查询之代码片段相关的内容有可能已经存在于现有大型代码库之内。在本篇文章中，我们将介绍完成这项任务所用到的两种模型：  

• NCS 是一套无监督模型，结合了 NLP 与 IR 技术。
• UNIF 是 NCS 的扩展，当拥有良好的监督数据可供训练使用时，它会利用监督神经网络模型以提高性能。

利用开源 Facebook AI 工具（包括 fastText、FAISS 以及 PyTorch），NCS 与 UNIF 得以将自然语言查询与代码片段表示为向量，而后训练网络，使得语义相似的代码片段与查询向量在向量空间之内紧密相连。通过这些模型，我们可以直接从代码语料库当中找到代码片段，以便有效回答工程师提出的问题。为了评估 NCS 与 UNIF，我们在 Stack Overflow 上使用了我们新创建的公共查询数据集。我们的模型能够正确回答该数据集中的问题，例如：  

• 如何关闭 / 隐藏 Android 软键盘？
• 如何在 Android 当中将位图转换为可绘制图片？
• 如何删除整个文件夹及其内容？
• 如何处理活动中的后退按钮？

NCS 的性能表明，这种相对简单的方法能够在源代码领域带来很好的收效。UNIF 的表现则证明，当拥有可用标记数据时，简单的监督学习方法能够带来显著的额外助益。在配合其它 Facebook 构建的系统（例如 Aroma 与 Getafix）时，该项目能够为我们的工程师提供广泛且不断增长的 ML 驱动型工具包，以帮助他们更加高效地编写并管理代码。

NCS 如何使用嵌入

NCS 模型利用嵌入来捕获程序语义（在本示例中，即代码片段的意图），并以连续向量进行表示。通过适当的计算，其即可在向量空间当中使语义相似的实体拥有彼此接近的期望属性。在以下示例当中，我们面对的是两个不同的方法体，其都与关闭或隐藏 Android 软键盘有关（前方提出的第一个问题）。由于二者具有相似的语义含义，因此即使其代码行内容并不完全相同，同样能够在向量空间中以彼此接近的点表示。

上图展示了向量空间中类似的代码片段是如何聚集的

我们利用这一概念来构建 NCS 模型。从高级角度来看，模型生成期间的每个代码片段都会以方法级粒度嵌入至向量空间之内。在模型构建完成之后，特定查询将映射至同一向量空间之内，其中向量距离将用于估算代码片段与查询内容的相关性。本节将更具体地描述模型生成以及搜索检索流水线，具体如下图所示。

此图展示了 NCS 的整体模型生成与搜索检索过程

模型生成

为了生成模型，NCS 必须提取单词、构建单词嵌入，而后构建文档嵌入。（这里的「文档」指方法体。）

提取单词

NCS 从源代码当中提取单词，并对其标记化以生成单词的线性序列。此处显示的示例数据来自 GitHub 上基于 Apache 2.0 许可协议的公开代码。

为了生成用于表示方法体的向量，我们需要将源代码视为文本，并从以下语法类别当中进行提取：方法名称、方法调用、枚举、字符串文字以及注释。而后，我们根据标准英语惯例（例如空格、标点符号等）以及与代码相关的标点符号（曲线与峰谷）对其进行标记。例如，对于上图中的方法体“pxToDp”，我们可以将源代码视为单词的集合，即“将 dp px 中的像素转换为 dp，获取资源，获取显示指标。”

对于我们语料库中的每个方法体，我们都可以使用这种方式进行源代码标记并学习每一个单词的嵌入。在这一步骤之后，我们将从每个方法体中提取到的单词列表整理为类似自然语言文档的形式。

构建单词嵌入

我们使用 fastText 为词汇语料库中的所有单词构建单词嵌入。FastText 利用两层密集神经网络计算这些向量表示，而该神经网络则可在大型语料库上进行无监督训练。具体来讲，我们使用到 skip-gram 模型，其中目标标记的嵌入被用于预测固定窗口大小之内的上下文标记嵌入。在以上示例当中，由于嵌入对应“dp”标记与窗口大小为 2，因此 skip-gram 模型将学习预测标记“pixel”、“in”、“px”以及“to”。其目标是学习嵌入矩阵 Tq ∈ R|Vc|×d，其中的|Vc |代表语料库的大小，d 为单词嵌入维度，T 中的第 k 行即代表 Vc 中第 k 个单词的嵌入。

在该矩阵当中，如果相应的单词经常出现在相似的上下文当中，则代表两个向量彼此靠近。我们使用该语句的逆向形式帮助定义语义关系：具有更接近的向量单词则应该具有较为相关的含义。这在 NLP 文献当中被称为分布假设，我们相信源文本中也存在相同的概念。

构建文档嵌入

下一步是使用方法体中存在的单词表达方法体的整体意图。为此，我们需要对方法体中单词集的单词嵌入向量进行加权平均。我们将此称为文档嵌入。

这里，d 代表的是方法体中的一组单词，vw 是单词 w 的 fastText 单词嵌入，C 是包含所有文档的语料库，u 则为标准化函数。

我们使用的术语“频率”是指逆文档频率函数（TF-IDF），负责为给定文档中的特定单词分配权重。其目标在于突出文档中最具代表性的单词——如果某个单词在文档当中经常出现，则其具有更高的权重；但如果它在语料库的过多文档中出现，则会受到权重惩罚。

在完成此步骤之后，我们已经将语料库中的每个方法体索引至文档向量表达中，模型生成也就此完成。

搜索查询

搜索查询往往表达为自然语言句子，例如“关闭 / 隐藏软键盘”或者“如何创建没有标题的对话框”。我们以与源代码相同的方式对查询内容进行标记，并使用相同的 fastText 嵌入矩阵 T。我们直接对各单词的向量表示进行平均化以建立起查询语句的文档嵌入；所有与核心内容无关的词汇都会被删除。接下来，我们使用标准相似性扫完算法 FAISS 来查找与查询内容具有最近余弦距离的文档向量，并返回前 n 个结果（再加上一些处理后的排名，后文将进一步做出解释）。

两个方法体与查询被映射至同一向量空间中靠近在一起的点上。这意味着查询与这两个方法体具有语义相似性，且后两者与查询内容相关。

结果

我们利用 Stack Overflow 上的问题评估 NCS 的性能，将标题作为查询内容，并将答案中的代码片段作为所需要的代码答案。面向特定查询，我们测量自己的模型能否从 GitHub 存储库集合当中的第五条、前五条以及前十条结果中找到正确答案（分别在下表中标记为 Answered @ 1,5,10）。我们还报告了平均倒数等级（MRR），用以衡量 NCS 能够在第几个结果当中找到正确答案。在我们创建的 Stack Overflow 评估数据集的总计 287 个问题当中，NCS 能够在前十条结果中正确回答的问题为 175 个；这相当于整个数据集的 60% 以上。我们还将 NCS 性能与传统 IR 技术 BM25 进行了比较。从表中可以看出，NCS 的表现要优于 BM25。

NCS 拥有良好问题回答能力的案例之一，是“从应用程序中启动 Android Market”，其中从 NCS 处返回的最优先结果为：  

private void showMarketAppIn() {try {startActivity(new Intent(Intent.ACTION_VIEW, Uri.parse("market://detailsid=" + BuildConfig.APPLICATION_ID)));} catch (ActivityNotFoundException e) {startActivity(new Intent(Intent.ACTION_VIEW, Uri.parse("http://play.google.com/store/apps/detailsid="+ BuildConfig.APPLICATION_ID)));}}

(此片段来自 GitHub 上的公开代码，依据 MIT 许可共享)

UNIF: 探索监督方法

NCS 的关键在于其使用了单词嵌入。由于 NCS 是一种无监督模型，因此具有以下几项显著优点：它可以通过搜索语料库进行学习，并能够快速轻松完成训练。NCS 假定查询中的单词来自与源代码中提取之单词相同的域，因为查询与代码片段都映射至相同的向量空间。然而，实际情况并非总是如此。例如，在查询“获取内部内存空间”时，其中的所有单词都不会出现在以下代码片段中。因此，我们需要将查询词“free space”映射至代码中的“available”一词处。  

File path = Environment.getDataDirectory();StatFs stat = new StatFs(path.getPath());long blockSize = stat.getBlockSize();long availableBlocks = stat.getAvailableBlocks();return Formatter.formatFileSize(this, availableBlocks * blockSize);

(以上代码片段来自 Stack Overflow 的公开代码，依据 CC-By-SA 3.0 许可共享)

利用收集自 14005 篇 Stack Overflow 帖子整理而成的数据集，我们分析了查询内容中的单词与源代码中单词之间的重叠。我们发现，查询当中存在 13972 个唯一单词，而源代码域中的单词量则不到一半（6072 个单词）。这表明如果查询包含源代码当中不存在的单词，那么我们的模型将无法正确发挥检索作用。因此，我们决定删除这些范围外单词。这一观察结果，促使我们将目光转向监督学习，尝试将查询中的单词映射至源代码中的单词。

我们决定尝试使用 UNIF，即 NCS 技术的有监督最小扩展，用以弥合自然语言单词与源代码单词之间的空白。在该模型当中，我们使用监督学习技术修改单词嵌入矩阵 T，分别为代码与查询标记生成两个嵌入矩阵 Tc 与 Tq。我们还利用基于 attention 的权重方案替代代码标记嵌入中的 TF-IDF 加权方法。

UNIF 模型的运作方式

我们利用与 NCS 相同的（c,q）数据点集合对 UNIF 进行训练，其中 c 与 q 分别代码代码与查询标记。（请参阅以下章节以了解该数据集的更多细节信息。）该模型架构可描述为以下形式。设 Tq ∈ R|Vq|×d 与 Tc ∈ R|Vc |×d 是两个嵌入矩阵，分别将自然语言描述与代码标记中的各个单词映射至长度为 d 的向量（Vq 为查询词汇语料库，Vc 为代码词汇语料库。）这两个矩阵使用相同的初始权重 T 进行初始化，并在训练期间各自独立修改（与 fastText 相对应）。为了将每个代码标记向量组成成文档向量，我们使用 attention 机制计算其加权平均值。其中 attention 权重 ac ∈ Rd 是在训练过程中学习到的一个 d 维向量，并且充当 TF-IDF 的对应实体。给定一组代码单词嵌入向量{e1, ….., en}，用于各个 ei 的 attention 权重 ai 将计算如下：

而后将文档向量计算为由 attention 权重加权的单词嵌入向量总和：

为了创建查询文档向量 eq，我们计算查询单词嵌入的简单平均值，方法与 NCS 类似。我们的训练过程通过经典的反向传播过程学习参数 Tq、Tc 与 ac。

此图展示了 UNIF 网络

检索的工作方式与 NCS 相同。首先给定查询，我们利用以上解释过的方法将其表示为一个文档向量，并使用 FAISS 查找与查询内容具有最近余弦距离的文档向量。（在原则上，UNIF 也将受益于 NCS 所给出的处理后排名。）

利用 NCS 进行结果比较

我们将 NCS 与 UNIF 在处理 Stack Overflow 评估数据集时得出的结果进行比较，以查看该模型是否正确能够在第一、前五以及前十项结果中找到正确答案，并观察其 MRR 得分。通过下表可以看出，UNIF 大大改善了 NCS 找到问题答案所需要的结果数量。

这再次证明监督学习技术在具有理想的训练语料库时，能够带来令人印象深刻的搜索性能。例如，使用搜索查询“如何退出应用程序并显示主屏幕？”时，NCS 返回的结果是：  

public void showHomeScreenDialog(View view) {Intent nextScreen = new Intent(getApplicationContext(), HomeScreenActivity.class);startActivity(nextScreen);}

UNIF 则给出了相关度更高的代码段：  

public void clickExit(MenuItem item) {Intent intent=new Intent(Intent.ACTION_MAIN);intent.addCategory(Intent.CATEGORY_HOME);intent.setFlags(Intent.FLAG_ACTIVITY_NEW_TASK);metr.stop();startActivity(intent);finish();}

（第一段来自基于 Apache 2.0 许可的 GitHub 公开可共享代码，第二段来自基于 MIT 许可的 GitHub 公开可共享代码)

再来看另一个例子，“如何获取 ActionBar 高度？”NCS 返回的答案为：  

public int getActionBarHeight() {return mActionBarHeight;}

UNIF 同样给出了相关度更高的代码片段：  

public static int getActionBarHeightPixel(Context context) {TypedValue tv = new TypedValue();if (context.getTheme().resolveAttribute(android.R.attr.actionBarSize, tv, true)) {return TypedValue.complexToDimensionPixelSize(tv.data,context.getResources().getDisplayMetrics());} else if (context.getTheme().resolveAttribute(R.attr.actionBarSize, tv, true)) {return TypedValue.complexToDimensionPixelSize(tv.data,context.getResources().getDisplayMetrics());} else {return 0;}}

(两段结果皆来自 GitHub 上基于 Apache 2.0 许可的可共享公开代码。感兴趣的朋友可以 点击此处 查看更多与 UNIF 性能相关的其它数据）

构建有效的 ML 支持型工具

创建成功机器学习工具的关键之一，在于获取高质量的训练数据集。对我们的模型而言，我们利用到 GitHub 中的大型开源代码库资源。此外，拥有高质量的评估数据集对于评估模型的实际质量同样非常重要。在探索相对较新的研究领域时（例如代码搜索），可用评估数据集的匮乏往往会限制我们评估各种代码搜索工具性能的能力。因此，为了帮助改善该领域的基准性能，我们从 Stack Overflow 当中整理出一套包含 287 个公开数据点的数据集，其中每个数据点由自然语言查询与“黄金”代码片段答案共同组成。

创建一套训练数据集

我们在 GitHub 上挑选了 26109 个最受欢迎的 Android 项目，直接在搜索语料库上训练我们的无监督模型 NCS。随着训练的推进，NCS 返回的代码片段又构成了新的搜索语料库。为了整合 UNIF 模型的监督学习技术，我们需要一对经过校准的数据点以学习映射关系。我们利用一组（c,q）数据点集合对 UNIF 进行训练，其中 q 代表自然语言描述或者查询，c 则代表对应的代码片段。我们从 Stack Overflow 问题标题以及 Stack Exchange 公开发布数据（基于 CC-BY-SA 3.0 许可）的代码片段中整理出这一数据集。在对问题进行各种启发式过滤之后——例如代码片段必须具有 Android 标记，或者必须存在方法调用，或者绝对不可包含 XML 标记等——我最终得到了 45 万 1 千个训练数据点。该数据集与评估查询完全不相交。（这反映出训练数据集的最佳案例可用性；正如我们在前文中所指出，基于文档字符串的训练方法并没能给出理想的结果。）

评估数据集

我们利用 Stack Overflow 评估了 NCS 的有效性。Stack Overflow 是一种非常实用的评估资源，其中包含大量自然语言查询素材，以及可以接受的高支持率答案。我们将特定 Stack Overflow 问题标题作为查询内容，NCS 则从 GitHub 中检索方法列表。在我们创建并改进 NCS 的过程当中，如果 NCS 中至少有一个前 n 项结果与 Stack Overflow 中回答代码片段采用的方法相匹配，我们即将其视为搜索成功。（在我们的评估当中，我们使用了第一、前五以及前十条答案分别进行计算。）

我们使用具有以下标准的脚本对 Stack Overflow 问题进行了筛选：1）问题包含“Android”与“Java”标记；2）代码答案至少得到一位用户的支持；3）实际代码片段至少与我们的 GitHub Android repos 语料库内容具有一项匹配。通过一定程度的手动处理，我们获得了这个包含 287 个问题的数据集。

利用 Aroma 进行自动评估

我们发现，手动评估搜索结果的正确性可能很难以可重复的方式进行，因为答案往往因作者与审查者的思路而有所区别。我们决定使用 Aroma 建立自动评估流水线。Aroma 在搜索结果与实际代码片段间给出相似性评分，超过评估阈值即证明答案能够正确回应查询内容。通过这条流水线，我们得以通过可重复方式评估模型。我们使用 Stack Overflow 上的代码答案作为评估的基本事实依据。

我们不仅利用上述评估比较池 UNIF 与 NCS 之间的差异，同时也将 UNIF 与文献中提到的其它一些代码搜索解决方案进行了比较。（关于更多细节信息，请 点击此处。）

这套工具集正在不断扩展

随着当今生产中大量代码存储库方案的广泛应用，机器学习技术从中提取出足以提升工程师生产力的重要模式与见解。在 Facebook 公司，我们使用的机器学习工具包括 Aroma 代码到代码推荐，以及基于 Getafix 的自动 bug 修复工具。NCS 与 UNIF 正是代码搜索模型当中的典型案例，其能够弥合自然语言模型与相关代码片段之间的表达差异。在这些工具的帮助下，工程师们将能够轻松查找并使用具有高相关性的代码片段，且适用范围涵盖各类专有源代码或者使用频率并不高的编程语言代码。在未来我们希望能够立足更为广泛的领域探索其它深度学习模型，从而进一步提高工程师们的生产效率。

https://ai.facebook.com/blog/neural-code-search-ml-based-code-search-using-natural-language-queries