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

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]]

'''
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

'''
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:
            for word, vec in get_vectors(fn):
                a.add_item(i, vec)
                id = 'i%d' % i
                word = 'w' + word
                txn.put(id, word)
                txn.put(word, id)
                i += 1
                if verbose:
                    if i % 1000 == 0:
                        print(i, '...')
        if verbose:
            print("Starting to build")
        a.build(num_trees)
        if verbose:
            print("Finished building")
        a.save(fn_annoy)
        if verbose:
            print("Annoy index saved to: {}".format(fn_annoy))
            print("lmdb map saved to: {}".format(fn_lmdb))
    else:
        print("Annoy index and lmdb map already in path")

'''
private function _create_args()
-------------------------------
Creates an argeparse object for CLI for create_index() function
Input:
    Void
Return:
    args object with required arguments for threshold_image() function
'''
def _create_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--embeddings", help="filename of the embeddings", type=str)
    parser.add_argument("--num_trees", help="number of trees to build index with", type=int)
    parser.add_argument("--verbose", help="print logging", action="store_true")

    args = parser.parse_args()
    return args

if __name__ == '__main__':
    args = _create_args()
    create_index(args.embeddings, num_trees=args.num_trees, verbose=args.verbose)

python2 make_annoy_index.py \
    --embeddings=<embedding path> \
    --num_trees=<int> \
    --verbose

'''
Vector Utils
Utils to read in vectors from txt, .bin, or .pkl.
Taken from Erik Bernhardsson
Source: https://github.com/erikbern/ann-presentation/blob/master/util.py
'''
import gzip
import struct
import cPickle

def _get_vectors(fn):
    if fn.endswith('.gz'):
        f = gzip.open(fn)
        fn = fn[:-3]

    else:
        f = open(fn)

    if fn.endswith('.bin'): # word2vec format
        words, size = (int(x) for x in f.readline().strip().split())

        t = 'f' * size

        while True:
            pos = f.tell()
            buf = f.read(1024)
            if buf == '' or buf == '\n': return
            i = buf.index(' ')
            word = buf[:i]
            f.seek(pos + i + 1)

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

            yield word.lower(), vec

    elif fn.endswith('.txt'): # Assume simple text format
        for line in f:
            items = line.strip().split()
            yield items[0], [float(x) for x in items[1:]]

    elif fn.endswith('.pkl'): # Assume pickle (MNIST)
        i = 0
        for pics, labels in cPickle.load(f):
            for pic in pics:
                yield i, pic
                i += 1


def get_vectors(fn, n=float('inf')):
    i = 0
    for line in _get_vectors(fn):
        yield line
        i += 1
        if i >= n:
            break

'''
Usage: python2 annoy_inference.py \
    --token='hello' \
    --num_results=<int> \
    --verbose
Query an Annoy index to find approximate nearest neighbors
'''
import annoy
import lmdb
import argparse

VEC_LENGTH = 50
FN_ANNOY = 'glove.6B.50d.txt.annoy'
FN_LMDB = 'glove.6B.50d.txt.lmdb'

a = annoy.AnnoyIndex(VEC_LENGTH)
a.load(FN_ANNOY)
env = lmdb.open(FN_LMDB, map_size=int(1e9))

'''
private function calculate(query, num_results)
-------------------------------
Queries a given Annoy index and lmdb map for num_results nearest neighbors
Input:
    query           - query to be searched
    num_results     - the number of results
Return:
    ret_keys        - list of num_results nearest neighbors keys
'''
def calculate(query, num_results, verbose=False):
    ret_keys = []
    with env.begin() as txn:
        id = int(txn.get('w' + query)[1:])
        if verbose:
            print("Query: {}, with id: {}".format(query, id))
        v = a.get_item_vector(id)
        for id in a.get_nns_by_vector(v, num_results):
            key = txn.get('i%d' % id)[1:]
            ret_keys.append(key)
    if verbose:
        print("Found: {} results".format(len(ret_keys)))
    return ret_keys

'''
private function _create_args()
-------------------------------
Creates an argeparse object for CLI for calculate() function
Input:
    Void
Return:
    args object with required arguments for threshold_image() function
'''
def _create_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--token", help="query word", type=str)
    parser.add_argument("--num_results", help="number of results to return", type=int)
    parser.add_argument("--verbose", help="print logging", action="store_true")

    args = parser.parse_args()
    return args

if __name__ == '__main__':
    args = _create_args()
    print(calculate(args.token, args.num_results, args.verbose))

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

['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']