Python版近傍検索NGT（ngtpy）の使い方

高次元ベクトルデータの近傍検索ソフトウェアNGT（2019年度情報処理学会業績賞受賞）の研究開発を行っているYahoo! JAPAN研究所の岩崎です。

Python版NGTとして、もう1年以上前にpipでインストールできるngtpyをリリースしましたが、このTech Blogでは紹介していなかったので紹介します。以前からctypesを使ったものがありますが、ctypesの部分が遅くてNGTの高速性が損なわれていたのでpybind11を使ったngtpyをリリースしました。現在ctypes版はサポートしていないので、まだctypes版を使っている方は、このタイミングでぜひngtpyに移行してください。

インストールは簡単です。

pip3 install ngt

なお、現在、Linux版とmacOS版のみです。

それでは、早速 fastText のデータを使った実例を紹介します。まず、以下のようにデータをダウンロードします。

curl -O https://dl.fbaipublicfiles.com/fasttext/vectors-english/wiki-news-300d-1M-subword.vec.zip
unzip wiki-news-300d-1M-subword.vec.zip

扱い易いように次のスクリプトでオブジェクト（objects.tsv）と単語（words.tsv）に分離します。

with open('wiki-news-300d-1M-subword.vec', 'r') as fi,\
     open('objects.tsv', 'w') as fov, open('words.tsv', 'w') as fow:
    n, dim = map(int, fi.readline().split())
    fov.write('{0}\t{1}\n'.format(n, dim))
    for line in fi:
        tokens = line.rstrip().split(' ')
        fow.write(tokens[0] + '\n')
        fov.write('{0}\n'.format('\t'.join(tokens[1:])))

次のスクリプトでインデックスを生成します。

import ngtpy

index_path = 'fasttext.anng'
with open('objects.tsv', 'r') as fin:
    n, dim = map(int, fin.readline().split())             # 次元数の抽出
    ngtpy.create(index_path, dim, distance_type='Cosine') # 空のインデックスの生成
    index = ngtpy.Index(index_path)                       # インデックスのオープン
    for line in fin:
        object = list(map(float, line.rstrip().split('\t')))
        index.insert(object)                              # オブジェクトの登録
index.build_index()                                       # インデックスのビルド
index.save()                                              # インデックスの保存

次元数と距離（ここではコサイン類似度）を指定してngtpy.createで空のインデックスを生成し、ngt.Indexで生成したインデックスをオープンします。index.insertでオブジェクトを登録します。この段階では登録されたオブジェクトのインデックスは生成されません。index.build_indexで登録されたオブジェクトのインデックスを一括で生成します。生成されたインデックスをindex.saveで保存します。

ビルドされたインデックスは次のスクリプトで検索できます。

import ngtpy

with open('words.tsv', 'r') as fin:
    words = list(map(lambda x: x.rstrip('\n'), fin.readlines()))

index = ngtpy.Index('fasttext.anng')                # インデックスのオープン
query_id = 10000
query_object = index.get_object(query_id)           # オブジェクトの取得

result = index.search(query_object, epsilon = 0.10) # 検索
print('Query={}'.format(words[query_id]))
for rank, object in enumerate(result):
    print('{}\t{}\t{:.6f}\t{}'.format(rank + 1, object[0], object[1], words[object[0]]))

まず、単語リストをロードしています。ここでは登録されているオブジェクトの中から10,000番目のオブジェクトをクエリとして利用します。そのために10,000番目のオブジェクトをindex.get_objectで取得して、これをクエリとしてindex.searchでインデックスを検索します。このスクリプトを実行すると以下のように出力されるはずです。なお、index.searchのepsilonを大きくすると検索時間は増えますが精度が上がります。

Query=Doctors
1       10000   0.000000        Doctors
2       4631    0.244096        Doctor
3       79542   0.258944        Medics
4       2044    0.263412        doctors
5       339397  0.274972        Doctoring
6       20667   0.280508        Physicians
7       80646   0.292580        Dentists
8       24255   0.292752        Nurses
9       9480    0.322195        Scientists
10      623160  0.330500        Practioners

以上で使い方はだいたい理解していただけると思います。詳しくは、ngtpy レファレンスやOptimization Examples Using Pythonをご覧ください。また、NGTの概要に関してはNGT -グラフ型インデックスによる近似近傍検索の基本-をご覧ください。