fastTextのベクトルの類似検索でYAHOOのNGTを使う。

自然言語処理 Python

fasttextにWikipedia全文を食べさせて遊んでいた(だいたい以下の記事の内容)。

qiita.com

「パリ - フランス + 日本」とかそういうやつ(だいたい以下の記事の内容)。

antibayesian.hateblo.jp

1 0.044094 -0.020766 -0.043709 -0.28123 0.13417 0.1673 0.05862 0.22077 -0.12497 0.077398 -0.022736  ...
2 0.1059 -0.050905 -0.10932 -0.22583 0.15278 0.15993 0.034985 0.12569 -0.10522 0.09798 -0.026379  ...
3 0.010506 -0.23203 -0.083458 -0.26819 0.14723 0.2554 -0.032926 0.029112 -0.24196 0.063793 ...
...

↑ fasttextで学習させるとこういうファイルが吐かれる。
最初の整数が単語に割り振ったトークンでそれからfloatが次元数分、ここでは300個並んでいる。
それが約92万行。学習はCPUで2時間くらいで終わる。
92万もあるとファイルサイズが2GBくらいになるのでパソコンでは全力で並列化とかしても読み込みに1分くらいかかってしまう。
さらに、対象のベクトルと全件の類似度を計算するみたいなことをやるのでパソコンでは全力で並列化しても1分くらいかかってしまう。

これだと気軽に使えないし、例によってslack botにするなどしたいのでさくらのVPSの安いやつでも数秒で計算が終わって欲しい。
たぶんアプローチとしては、読み込みも類似計算もこのベクトル空間内で近しいものだけ行うようにインデックス化すればいいんだと思うけど
それは大変だなーとか思ってググっているとすごいツールがあった。

NGT(Neighborhood Graph and Tree)

NGT

簡単にいうとfastextが吐いたような1行に1ベクトルみたいなファイルをインデックス化して、それを使って高速に空間内の距離の検索をしてくれるコマンドラインツール。

インストールは書いてある通りにやれば簡単に出来ると思う。 注意点としては

  • ubuntu, centosなどのLinuxでは動いたけどMacではインストールはできてもコマンド引数が認識されないみたいなエラーが出て動かなかった。
  • githubから最新版をクローンしてくると動かなかったのでおとなしくタグが打ってあるやつを使うといい。

前処理

最初にNGTにベクトル放り込んでインデックス化してもらう。
そのためにfasttextが吐いたやつの下処理をする。

import numpy as np


def sorted_vectors(vec_path):
    vec_file = open(vec_path)
    # 最初の行はヘッダなので捨てる
    vec_file.readline()

    rows = []
    for token, vec in (row.strip().split(maxsplit=1) for row in vec_file):
        # なぜかトークンではないものが混じることがある。
        try:
            rows.append((int(token), vec))
        except ValueError:
            continue

    # 行数とトークンが揃っててほしいので、トークンでソートする。
    rows.sort(key=lambda x: x[0])
    # 1からの連番になることを確認
    for idx, row in enumerate(rows, 1):
        assert idx == row[0]
    yield from (row[1] for row in rows)


def prepare_for_ngt(vec_path, dist_vec):
    dist = open(dist_vec, 'w')
    for row in sorted_vectors(vec_path):
        vec = np.array(row.split(), dtype='float32')
        # ノルムで割ってベクトルを正規化。
        vec /= np.linalg.norm(vec)
        row = list(vec)
        row = ' '.join(map(str, row))
        dist.write('{}\n'.format(row))

正規化をしているのは結果をコサイン類似度と揃えたかったから。 techblog.yahoo.co.jp

NGTは距離しか扱えないため、コサイン類似度と同じ検索結果を得るために事前にベクトルのノルムが1となるように正規化してユークリッド距離を使う方法か、

距離関数はcreateの引数で選択できる。
デフォルトだとユークリッド距離。

これにfasttextが吐いたファイルを与えると1行に1つのベクトルで、トークンでソートされているファイルが出来上がる。

prepare_for_ngt('model.vec', 'ngt_model.vec')

そして以下のコマンドを叩く。
-dは次元数,indexは出力されるディレクトリ名。
$ ngt create -d 300 index ngt_model.vec
詳しいオプションは→https://research-lab.yahoo.co.jp/software/ngt/command.html#create
この所用時間はコア数に応じて数十分~1時間くらい。
もちろん出力されたものは使いまわせる。

検索

$ ngt search index target.vec
taget.vecは検索したいベクトルのファイルで、形式はcreateに使ったファイルと同じ。

Query No.1
Rank    ID  Distance
  1    野球    0
  2    ソフトボール    0.682129
  3    プロ野球    0.708267
  4    サッカー    0.739252
  5    軟式野球    0.754248
  6    少年野球    0.755179
  7    バスケットボール    0.765375
  8    高校野球    0.788223
  9    リトルリーグ    0.788254
  10   アマチュア野球    0.794894
  11   アメリカンフットボール    0.799319
  12   ホッケー    0.805228
  13   ベースボール    0.822246
  14   ラグビー    0.822764
  15   アメフト    0.823869
  16   硬式野球    0.832693
  17   バスケ    0.834052
  18   大学野球    0.835599
  19   卓球    0.835896
  20   投手    0.836089

結果がこういう形で表示される(本当に正しいかわかりづらかったのでトークンを単語に変換している)。
indexに放り込んだベクトルから近い順に返ってくる。
さくらのVPSの下から2番目のやつでも3秒くらいで返ってくる。

まとめ

このツールは、界隈では自分が紹介するまでもなく有名で、かつ使われているのかもしれないけれどリリースが結構前な割には(商用利用出来るようになったのは去年だけど)そんなに情報が多くなかったので、ただ使ってみただけだけれども書いてみた。
多次元ベクトルのインデックス化は明らかに僕のできる範囲を超えているので感動してしまった。ごちそうさまです。

あと一応試せるように使用したモデルデータを置いておきます。
これは冒頭のqiitaの記事の過程で得られるものだけど。
model.vec.gz - Google ドライブ
wiki__all.vocab.gz - Google ドライブ

おわり