Python
機械学習
TensorFlow
1522480883@sugiyamath

Facebook Researchのfaissで類似検索

Python機械学習TensorFlow

faissとは、類似性検索およびクラスタリングのための高速アルゴリズムを実装したライブラリです。ここでは、文書をベクトル化してfaissのインデクスに登録し、ベクトル化したクエリで検索してみます。

簡単な概念図

Untitled drawing.jpg

jupyterで実行

まず、必要となるライブラリをインポートします。faissは事前にインストールしておいてください。

In[1]:

import faiss
import tensorflow as tf
import tensorflow_hub as hub
from janome.tokenizer import Tokenizer

次に、例文を用意してベクトル化します。

In[2]:

example_sentences = [
    "京都の大学",
    "アメリカの美味しい食べ物",
    "機械学習の本",
    "ビル・ゲイツ",
    "御殿場市民"
]

jtok = Tokenizer()

with tf.Graph().as_default():
    embed = hub.Module("https://tfhub.dev/google/nnlm-ja-dim128/1")
    embeddings = embed(list(map(lambda x: ' '.join([y.surface for y in jtok.tokenize(x)]), example_sentences)))

    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        sess.run(tf.tables_initializer())
        example_vectors = sess.run(embeddings)

faissにベクトルをインデクシングします。

In[3]:

dim = 128
index = faiss.IndexFlatL2(dim)   
index.add(example_vectors)

検索用メソッドを作成します。

In[4]:

def encode_query(text):
    with tf.Graph().as_default():
        embed = hub.Module("https://tfhub.dev/google/nnlm-ja-dim128/1")
        embeddings = embed(list(map(lambda x: ' '.join([y.surface for y in jtok.tokenize(x)]), [text])))

        with tf.Session() as sess:
            sess.run(tf.global_variables_initializer())
            sess.run(tf.tables_initializer())
            return sess.run(embeddings)


def search_by_text(text, index, sents, k=1):
    D, I = index.search(encode_query(text), k)
    return [sents[i] for i in I[0]]

最後に、いくつかクエリで検索して試してみます。

In[5]:

print(
  search_by_text("バカ", index, example_sentences),
  search_by_text("プログラマー", index, example_sentences),
  search_by_text("立命館", index, example_sentences)
)

Out[5]:

['御殿場市民'] ['ビル・ゲイツ'] ['京都の大学']

まあ、大体それっぽく検索できているようです。（この文例の中で最も"バカ"に近い文は"御殿場市民")

ちなみに、search_by_textは、indexの中からクエリのtextに近いk件の文ベクトルを探すメソッドです。

参考

[1] https://github.com/facebookresearch/faiss/wiki/Getting-started
[2] https://www.tensorflow.org/hub/

Edit request
2