你好,游客 登录
背景:
阅读新闻

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

[日期:2018-01-25] 来源:机器之心  作者: [字体: ]

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

  1. King + (Woman - Man) = Queen

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

  1. vectors = np.array(embeddingmodel.embeddings)

  2. ranks = np.dot(query,vectors.T)/np.sqrt(np.sum(vectors**2,1))

  3. mostSimilar = []

  4. [mostSimilar.append(idx) for idx in ranks.argsort()[::-1]]

  5.  

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

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

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

制作一个索引

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

  1. '''

  2. Usage: python2 make_annoy_index.py

  3. --embeddings=<embedding path>

  4. --num_trees=<int>

  5. --verbose

  6. Generate an Annoy index and lmdb map given an embedding file

  7. Embedding file can be

  8. 1. A .bin file that is compatible with word2vec binary formats.

  9. There are pre-trained vectors to download at https://code.google.com/p/word2vec/

  10. 2. A .gz file with the GloVe format (item then a list of floats in plaintext)

  11. 3. A plain text file with the same format as above

  12. '''

  13.  

  14. import annoy

  15. import lmdb

  16. import os

  17. import sys

  18. import argparse

  19.  

  20. from vector_utils import get_vectors

  21.  

最后一行里非常重要的是「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 索引。

  1. '''

  2. function create_index(fn, num_trees=30, verbose=False)

  3. -------------------------------

  4. Creates an Annoy index and lmdb map given an embedding file fn

  5. Input:

  6. fn - filename of the embedding file

  7. num_trees - number of trees to build Annoy index with

  8. verbose - log status

  9. Return:

  10. Void

  11. '''

  12. def create_index(fn, num_trees=30, verbose=False):

  13. fn_annoy = fn + '.annoy'

  14. fn_lmdb = fn + '.lmdb' # stores word <-> id mapping

  15.  

  16. word, vec = get_vectors(fn).next()

  17. size = len(vec)

  18. if verbose:

  19. print("Vector size: {}".format(size))

  20.  

  21. env = lmdb.open(fn_lmdb, map_size=int(1e9))

  22. if not os.path.exists(fn_annoy) or not os.path.exists(fn_lmdb):

  23. i = 0

  24. a = annoy.AnnoyIndex(size)

  25. with env.begin(write=True) as txn:

  26. for word, vec in get_vectors(fn):

  27. a.add_item(i, vec)

  28. id = 'i%d' % i

  29. word = 'w' + word

  30. txn.put(id, word)

  31. txn.put(word, id)

  32. i += 1

  33. if verbose:

  34. if i % 1000 == 0:

  35. print(i, '...')

  36. if verbose:

  37. print("Starting to build")

  38. a.build(num_trees)

  39. if verbose:

  40. print("Finished building")

  41. a.save(fn_annoy)

  42. if verbose:

  43. print("Annoy index saved to: {}".format(fn_annoy))

  44. print("lmdb map saved to: {}".format(fn_lmdb))

  45. else:

  46. print("Annoy index and lmdb map already in path")

  47.  

我已经推断出 argparse,因此,我们可以利用命令行启用我们的脚本:

  1. '''

  2. private function _create_args()

  3. -------------------------------

  4. Creates an argeparse object for CLI for create_index() function

  5. Input:

  6. Void

  7. Return:

  8. args object with required arguments for threshold_image() function

  9. '''

  10. def _create_args():

  11. parser = argparse.ArgumentParser()

  12. parser.add_argument("--embeddings", help="filename of the embeddings", type=str)

  13. parser.add_argument("--num_trees", help="number of trees to build index with", type=int)

  14. parser.add_argument("--verbose", help="print logging", action="store_true")

  15.  

  16. args = parser.parse_args()

  17. return args

  18.  

添加主函数以启用脚本,得到 make_annoy_index.py:

  1. if __name__ == '__main__':

  2. args = _create_args()

  3. create_index(args.embeddings, num_trees=args.num_trees, verbose=args.verbose)

  4.  

现在我们可以仅利用命令行启用新脚本,以生成 Annoy 索引和对应的 lmdb 图!

  1. python2 make_annoy_index.py

  2. --embeddings=<embedding path>

  3. --num_trees=<int>

  4. --verbose

  5.  

写向 量Utils

我们在 make_annoy_index.py 中推导出 Python 脚本 vector_utils。现在要写该脚本,Vector_utils 用于帮助读取.txt, .bin 和 .pkl 文件中的向量。

写该脚本与我们现在在做的不那么相关,因此我已经推导出整个脚本,如下:

  1. '''

  2. Vector Utils

  3. Utils to read in vectors from txt, .bin, or .pkl.

  4. Taken from Erik Bernhardsson

  5. Source: https://github.com/erikbern/ann-presentation/blob/master/util.py

  6. '''

  7. import gzip

  8. import struct

  9. import cPickle

  10.  

  11. def _get_vectors(fn):

  12. if fn.endswith('.gz'):

  13. f = gzip.open(fn)

  14. fn = fn[:-3]

  15.  

  16. else:

  17. f = open(fn)

  18.  

  19. if fn.endswith('.bin'): # word2vec format

  20. words, size = (int(x) for x in f.readline().strip().split())

  21.  

  22. t = 'f' * size

  23.  

  24. while True:

  25. pos = f.tell()

  26. buf = f.read(1024)

  27. if buf == '' or buf == 'n': return

  28. i = buf.index(' ')

  29. word = buf[:i]

  30. f.seek(pos + i + 1)

  31.  

  32. vec = struct.unpack(t, f.read(4 * size))

  33.  

  34. yield word.lower(), vec

  35.  

  36. elif fn.endswith('.txt'): # Assume simple text format

  37. for line in f:

  38. items = line.strip().split()

  39. yield items[0], [float(x) for x in items[1:]]

  40.  

  41. elif fn.endswith('.pkl'): # Assume pickle (MNIST)

  42. i = 0

  43. for pics, labels in cPickle.load(f):

  44. for pic in pics:

  45. yield i, pic

  46. i += 1

  47.  

  48.  

  49. def get_vectors(fn, n=float('inf')):

  50. i = 0

  51. for line in _get_vectors(fn):

  52. yield line

  53. i += 1

  54. if i >= n:

  55. break

  56.  

测试 Annoy 索引和 lmdb 图

我们已经生成了 Annoy 索引和 lmdb 图,现在我们来写一个脚本使用它们进行推断。

将我们的文件命名为 annoy_inference.py,得到下列依赖项:

  1. '''

  2. Usage: python2 annoy_inference.py

  3. --token='hello'

  4. --num_results=<int>

  5. --verbose

  6. Query an Annoy index to find approximate nearest neighbors

  7. '''

  8. import annoy

  9. import lmdb

  10. import argparse

  11.  

现在我们需要在 Annoy 索引和 lmdb 图中加载依赖项,我们将进行全局加载,以方便访问。注意,这里设置的 VEC_LENGTH 为 50。确保你的 VEC_LENGTH 与嵌入长度匹配,否则 Annoy 会不开心的哦~

  1. VEC_LENGTH = 50

  2. FN_ANNOY = 'glove.6B.50d.txt.annoy'

  3. FN_LMDB = 'glove.6B.50d.txt.lmdb'

  4.  

  5. a = annoy.AnnoyIndex(VEC_LENGTH)

  6. a.load(FN_ANNOY)

  7. env = lmdb.open(FN_LMDB, map_size=int(1e9))

  8.  

有趣的部分在于「calculate」函数。

1. 从 lmdb 图中获取查询索引;

2. 用 get_item_vector(id) 获取 Annoy 对应的向量;

3. 用 a.get_nns_by_vector(v, num_results) 获取 Annoy 的最近邻。

  1. '''

  2. private function calculate(query, num_results)

  3. -------------------------------

  4. Queries a given Annoy index and lmdb map for num_results nearest neighbors

  5. Input:

  6. query - query to be searched

  7. num_results - the number of results

  8. Return:

  9. ret_keys - list of num_results nearest neighbors keys

  10. '''

  11. def calculate(query, num_results, verbose=False):

  12. ret_keys = []

  13. with env.begin() as txn:

  14. id = int(txn.get('w' + query)[1:])

  15. if verbose:

  16. print("Query: {}, with id: {}".format(query, id))

  17. v = a.get_item_vector(id)

  18. for id in a.get_nns_by_vector(v, num_results):

  19. key = txn.get('i%d' % id)[1:]

  20. ret_keys.append(key)

  21. if verbose:

  22. print("Found: {} results".format(len(ret_keys)))

  23. return ret_keys

  24.  

再次,这里使用 argparse 来使读取命令行参数更加简单。

  1. '''

  2. private function _create_args()

  3. -------------------------------

  4. Creates an argeparse object for CLI for calculate() function

  5. Input:

  6. Void

  7. Return:

  8. args object with required arguments for threshold_image() function

  9. '''

  10. def _create_args():

  11. parser = argparse.ArgumentParser()

  12. parser.add_argument("--token", help="query word", type=str)

  13. parser.add_argument("--num_results", help="number of results to return", type=int)

  14. parser.add_argument("--verbose", help="print logging", action="store_true")

  15.  

  16. args = parser.parse_args()

  17. return args

  18.  

主函数从命令行中启用 annoy_inference.py。

  1. if __name__ == '__main__':

  2. args = _create_args()

  3. print(calculate(args.token, args.num_results, args.verbose))

  4.  

现在我们可以使用 Annoy 索引和 lmdb 图,获取查询的最近邻!

  1. python2 annoy_inference.py --token="test" --num_results=30

  2.  

  3. ['test', 'tests', 'determine', 'for', 'crucial', 'only', 'preparation', 'needed', 'positive', 'guided', 'time', 'performance', 'one', 'fitness', 'replacement', 'stages', 'made', 'both', 'accuracy', 'deliver', 'put', 'standardized', 'best', 'discovery', '.', 'a', 'diagnostic', 'delayed', 'while', 'side']

  4.  

代码

本教程所有代码的 GitHub 地址:https://github.com/kyang6/annoy_tutorial

原文地址:https://medium.com/@kevin_yang/simple-approximate-nearest-neighbors-in-python-with-annoy-and-lmdb-e8a701baf905

收藏 推荐 打印 | 录入:Cstor | 阅读:
相关新闻      
本文评论   查看全部评论 (0)
表情: 表情 姓名: 字数
点评:
       
评论声明
  • 尊重网上道德,遵守中华人民共和国的各项有关法律法规
  • 承担一切因您的行为而直接或间接导致的民事或刑事法律责任
  • 本站管理人员有权保留或删除其管辖留言中的任意内容
  • 本站有权在网站内转载或引用您的评论
  • 参与本评论即表明您已经阅读并接受上述条款