算法集锦（13）|自然语言处理| Python代码的语义搜索引擎创建

现代搜索引擎的力量非常强大，可以让你瞬间从互联网中获取想要的知识。但是，现有技术也存在着无法忽视的局限性，比如搜索非文字内容或者内容难以用“关键词”描述时，都难以达到预期的搜索效果。更进一步，现有搜索技术难以让用户实现“语义”搜索，即通过文字内容的意义来检索相关内容。

今天，我们分享一个简单易行的算法，可以实现对任意对象的语义搜索。具体来说，该算法创建了一个系统，可以对python代码进行语义搜索，但该方法也可以推广到其他内如(例如图片或视频等)。

搜索内容为“Ping REST api and return results”，引擎返回了合理的检索内容，尽管检索到的代码或注释中并没有包含Ping，REST或api等关键词。

该例子展示了语义搜索的强大：我们可以结合关键词以及关键词代表的意义来最大限度的找到想要的内容。语义搜索的深刻意义在于：即使我们不熟悉代码或者难以找到合适的关键词，依然可以确保用户检索到需要的内容。

创建一个共享向量空间

在深入技术细节之前，从直观上了解语义搜索是如何实现的，是非常有意义的。其中心思想是：将想要搜索的内容（如代码）变换到共享向量空间（shared vector space）中。

算法的目标是将代码映射到自然语言的向量空间中，然后利用余弦相似性（Cosine Similarity）将代表相似意义的代码聚类的一起，而不相关的内容则会分布在较远的坐标上。我们提供的方法可以利用预训练模型提取代码特征，然后再调试（fine-tuning）该模型从而实现将潜在的代码映射到自然语言向量空间。

本文将分为5个具体步骤介绍算法。下面将演示这些步骤，当您在本教程中继续学习时，这些步骤将是一个有用的参考。在完成本教程之后，有必要重新检查这个图，以确认所有步骤是如何结合在一起的。

步骤1 获得和清洗数据

BigQuery是谷歌收集和存储的开源数据集（在GitHub上），可以用于各种有趣的数据科学项目。本项目就是采用的该数据集，当您注册一个谷歌云帐户时，他们会给您300美元，这足以查询此练习的数据。获取这些数据非常方便，因为您可以使用SQL查询来选择要查找的文件类型，以及关于repos的其他元数据。

收集这些数据之后，我们需要将这些文件解析为code-docstring(代码、文档字符串)对。对于本教程，一个代码单元将是顶级函数或方法。我们将匹配后的code-docstring对作为模型的训练数据，以便对代码进行处理(稍后将详细介绍)。我们还去掉了所有注释，只保留代码。这项工作任务量很大，但在Python的标准库中有一个名为ast的程序库，它可以用来提取函数、方法和文档字符串。利用ast库，我们可以先将代码转换成抽象语法树（Abstract syntax tree，AST），然后再使用Astor库将AST反转换成代码，从而达到从代码中删除注释的目的。

def tokenize_docstring(text):
    """Apply tokenization using spacy to docstrings."""
    tokens = EN.tokenizer(text)
    return [token.text.lower() for token in tokens if not token.is_space]
def tokenize_code(text):
    """A very basic procedure for tokenizing code strings."""
    return RegexpTokenizer(r'\w+').tokenize(text)
def get_function_docstring_pairs(blob):
    """Extract (function/method, docstring) pairs from a given code blob."""
    pairs = []
    try:
        module = ast.parse(blob)
        classes = [node for node in module.body if isinstance(node, ast.ClassDef)]
        functions = [node for node in module.body if isinstance(node, ast.FunctionDef)]
        for _class in classes:
            functions.extend([node for node in _class.body if isinstance(node, ast.FunctionDef)])
        for f in functions:
            source = astor.to_source(f)
            docstring = ast.get_docstring(f) if ast.get_docstring(f) else ''
            function = source.replace(ast.get_docstring(f, clean=False), '') if docstring else source
            pairs.append((f.name,
                          f.lineno,
                          source,
                          ' '.join(tokenize_code(function)),
                          ' '.join(tokenize_docstring(docstring.split('\n\n')[0]))
                         ))
    except (AssertionError, MemoryError, SyntaxError, UnicodeEncodeError):
        pass
    return pairs

我们将数据分为训练集、验证集和测试集，以便开展模型训练。为了追踪每个（代码、文档）对，算法中特意设置了lineage文件。

步骤2： 利用Seq2Seq模型创建代码摘要

可以采用GitHub issue summarizer 来创建sequence-to-sequence模型来总结代码。不同的是这里用python代码替代issues数据，用文档字符串代替issue标题。

但是，与GitHub的issue文本不同，代码不是自然语言。为了完全的显现代码中蕴含的信息，我们的采用了领域指定优化（domain-specific optimizations）方法，比如tree-based LSTM和语法标记策略（syntax-aware tokenization）。采用上述方法，我们可以很便捷的将代码像自然语言一样处理，并获得合理的结果。

我们训练issue summarizer模型的目的不是对代码进行汇总，而是从中提取代码的特征。从技术上讲，该步骤是可选的，我们可以直接跳过该步骤，直接进行模型权重初始化或以下流程。

在后面的步骤中，我们将从这个模型中提取编码器并对它进行微调以完成另一个任务。下面是这个模型的一些输出示例：

可以看到，虽然结果并不完美，但却有力的证明了模型已经学会从代码中提取一些语义意义，这是我们完进行这项任务的主要目标。我们可以使用BLEU度量对这些模型进行定量评估。

需要指出的是，训练Seq2Seq模型以建立代码摘要，并不是构建代码特征提取器的惟一技术。例如，您还可以训练一个GAN，并使用鉴别器作为特征提取器。

步骤3： 训练语言模型来编码自然语言语句

我们已经构建了一种将代码表示为向量的机制，那么就需要一种类似的方法来编码自然语言语句（Nature Language Phrase）。

有许多通用的的模型可以产生高质量的语句嵌入(也称为句子嵌入)。例如，谷歌的通用语句编码器（可以在Tensorflow Hub上获取），实现证明该编码器在许多现实的应用中都工作得很好。

这些预训练的模型不仅很方便，而且可以通过微调获取指定区域的词汇表和文档字符串的语义信息。可以用来实现语句嵌入的方法很多，简单的方法如平均词向量（averaging word vector），而那些用于构建通用语句编码器的方法则相对复杂些。

本算法使用了AWD LSTM生成语句嵌入的神经网络模型。该过程通过fast.ai库实现，该程序库提供了非常便捷和快速的方式来创建我们需要的模型。

构建语言模型时，需要仔细考虑用于训练的语料库。通常，使用与待解决问题相关的语料库是最理想的选择，以便能够充分捕获相关的语义和词汇表。对于本算法，stack overflow数据集是一个很适用的语料库，因为它包含了大量的代码讨论的内容。然而，为了保持本算法的简单性，我们采用文档字符串（docstrings）用作我们的语料库。这是次优的，因为关于堆栈溢出的讨论通常包含比一行docstring中更丰富的语义信息。如果读者感兴趣，可以将本算法使用其他语料库进行训练，并检验对最终结果的影响。

在训练语言模型之后，下一个任务是使用这个模型为每个句子生成一个嵌入（embedding）。一种常见的策略是总结语言模型的隐藏状态，例如采用concat pooling方法。然而，为了简单起见，我们将对所有隐藏状态进行平均处理。

下面的代码，可以实现从fast.ai语言模型中提取隐含状态的平均值。

评估语句嵌入的一个好方法是测量这些嵌入对诸如情感分析、文本相似性等下游任务的有效性。通常，可以使用通用基准来度量嵌入的质量。但是，该策略可能不适合本算法，因为我们的数据是来源于特定领域的。现阶段，我们还没有为代码语义查询设计出可以开源的下游任务。在现有情况下，我们只能通过预先的判定来检查语句之间的相似性，来判断这些嵌入是否包含语义信息。

下图展示了一些示例，我们在向量化的docstring中搜索用户提供的短语的相似性。

需要注意的是，这只是一个合理性检查—更严格的方法是度量这些嵌入对各种下游任务的影响，并使用它对嵌入质量形成更客观的意见。

步骤4： 将代码向量和自然语言映射到相同的向量空间

步骤4的流程图如下所示。

本步骤中，我们在步骤2中的seq2seq模型中加入Dense Layers层，通过微调使模型可以进行docstring嵌入预测。

在训练这个模型的冻结版本之后，我们解冻所有的层并且训练这个模型几个周期，这有助于微调模型对这个任务的表现。

最后，我们希望对代码进行矢量化，以便构建搜索索引。出于评估目的，我们还将对不包含docstring的代码进行矢量化，以便查看此过程如何很好地推广到我们尚未看到的数据。

步骤5： 创建语义搜索工具

本步骤中，我们结合前面提到的方法来创建一个搜索索引。

在步骤4中，我们向量化了所有不包含任何docstring的代码。下一步是将这些向量放到一个搜索索引中，以便快速检索最近的匹配。实现该功能的一个可行方法是采用python库中的nmslib函数。

构建代码向量搜索索引后，需要一种方法将字符串(查询)转换为向量。为此，可以使用步骤3中的语言模型。为了简化这个过程，我们在lang_model_utils.py 中提供了一个helper类（Query2Emb）。

最后，一旦我们能够将字符串转换为查询向量，我们就可以为这个向量获取最近的匹配对，比如：

idxs, dists = self.search_index.knnQuery(query_vector, k=k)

搜索索引将返回两个条目：

(1)一个索引列表，这些索引是数据集中最近匹配的整数位置

(2)这些邻匹配与查询向量的距离(这里定义索引使用余弦距离)。有了这些信息之后，就可以直接构建语义搜索了。详见代码中的Build Search Index.ipynb。

最后，向您展示下利用本算法实现的代码语义搜索效果。