如何在Python中快速进行语料库搜索：近似最近邻算法

本文经机器之心（微信公众号：almosthuman2014）授权转载，禁止二次转载。

最近，我一直在研究在 GloVe 词嵌入中做加减法。例如，我们可以把「king」的词嵌入向量减去「man」的词嵌入向量，随后加入「woman」的词嵌入得到一个结果向量。随后，如果我们有这些词嵌入对应的语料库，那么我们可以通过搜索找到最相似的嵌入并检索相应的词。如果我们做了这样的查询，我们会得到：

King + (Woman - Man) = Queen

我们有很多方法来搜索语料库中词嵌入对作为最近邻查询方式。绝对可以确保找到最优向量的方式是遍历你的语料库，比较每个对与查询需求的相似程度——这当然是耗费时间且不推荐的。一个更好的技术是使用向量化余弦距离方式，如下所示：

vectors = np.array(embeddingmodel.embeddings)
ranks = np.dot(query,vectors.T)/np.sqrt(np.sum(vectors**2,1))
mostSimilar = []
[mostSimilar.append(idx) for idx in ranks.argsort()[::-1]]

想要了解余弦距离，可以看看这篇文章：http://masongallo.github.io/machine/learning,/python/2016/07/29/cosine-similarity.html

矢量化的余弦距离比迭代法快得多，但速度可能太慢。是近似最近邻搜索算法该出现时候了：它可以快速返回近似结果。很多时候你并不需要准确的最佳结果，例如：「Queen」这个单词的同义词是什么？在这种情况下，你只需要快速得到足够好的结果，你需要使用近似最近邻搜索算法。

在本文中，我们将会介绍一个简单的 Python 脚本来快速找到近似最近邻。我们会使用的 Python 库是 Annoy 和 Imdb。对于我的语料库，我会使用词嵌入对，但该说明实际上适用于任何类型的嵌入：如音乐推荐引擎需要用到的歌曲嵌入，甚至以图搜图中的图片嵌入。

制作一个索引

让我们创建一个名为：「make_annoy_index」的 Python 脚本。首先我们需要加入用得到的依赖项：

'''
Usage: python2 make_annoy_index.py \
    --embeddings=<embedding path> \
    --num_trees=<int> \
    --verbose
Generate an Annoy index and lmdb map given an embedding file
Embedding file can be
  1. A .bin file that is compatible with word2vec binary formats.
     There are pre-trained vectors to download at https://code.google.com/p/word2vec/
  2. A .gz file with the GloVe format (item then a list of floats in plaintext)
  3. A plain text file with the same format as above
'''
import annoy
import lmdb
import os
import sys
import argparse
from vector_utils import get_vectors

最后一行里非常重要的是「vector_utils」。稍后我们会写「vector_utils」，所以不必担心。

接下来，让我们丰富这个脚本：加入「creat_index」函数。这里我们将生成 lmdb 图和 Annoy 索引。

• 1. 首先需要找到嵌入的长度，它会被用来做实例化 Annoy 的索引。
• 2. 接下来实例化一个 Imdb 图，使用：「env = lmdb.open(fn_lmdb, map_size=int(1e9))」。
• 3. 确保我们在当前路径中没有 Annoy 索引或 lmdb 图。
• 4. 将嵌入文件中的每一个 key 和向量添加至 lmdb 图和 Annoy 索引。
• 5. 构建和保存 Annoy 索引。

'''
function create_index(fn, num_trees=30, verbose=False)
-------------------------------
Creates an Annoy index and lmdb map given an embedding file fn
Input:
    fn              - filename of the embedding file
    num_trees       - number of trees to build Annoy index with
    verbose         - log status
Return:
    Void
'''
def create_index(fn, num_trees=30, verbose=False):
    fn_annoy = fn + '.annoy'
    fn_lmdb = fn + '.lmdb' # stores word <-> id mapping
    word, vec = get_vectors(fn).next()
    size = len(vec)
    if verbose:
        print("Vector size: {}".format(size))
    env = lmdb.open(fn_lmdb, map_size=int(1e9))
    if not os.path.exists(fn_annoy) or not os.path.exists(fn_lmdb):
        i = 0
        a = annoy.AnnoyIndex(size)
        with env.begin(write=True) as txn:
声明：所有来源为“聚合数据”的内容信息，未经本网许可，不得转载！如对内容有异议或投诉，请与我们联系。邮箱：marketing@think-land.com