似た文書をベクトル検索で探し出したい ~SentenceTransformersとFaissで効率的にベクトル検索~
はじめに
この記事では、ベクトル検索で似た文書を検索するコードを解説します。具体的には、Sentence Transformersライブラリを用いてベクトル化、Faissという近似最近傍探索ライブラリを用いて高速な検索を行います。
用語説明
ベクトル検索
... 文書の検索にベクトルを使用する方法。例えば、文書に映画に関する内容が 10 %、音楽が 2 %、俳優が30%含まれていた時、シンプルにそれを表すと [0.1, 0.02, 0.3]というベクトルを作ることができる。Googleの説明が詳しい。
Sentence Transformers
... ベクトル検索に必要なベクトル化を行うためのライブラリ
Faiss
... ベクトル同士の類似度を高速に検索してくれるライブラリ。以下のHakkyさんのページが詳しい。
コード
入力：検索したい文字のリスト、検索される文字のリスト
出力：どの文字が最も類似しているかとそのスコア
という関数を作成してみました。
Github, Google Colab から実行可能です。
# !pip install -Uqq sentence-transformers faiss-gpu
import faiss
import numpy as np
import pandas as pd
from sentence_transformers import SentenceTransformer
def search_nearest_sentence(sentence_refs, sentence_queries, model_name="intfloat/multilingual-e5-base", path_out_faiss_full_index_path=None):
# モデルの選択と読み込み
model = SentenceTransformer(model_name)
# 英語ならば他にもよいモデルがある
# https://huggingface.co/spaces/mteb/leaderboard など参照
# model_name = "thenlper/gte-base"
# model = SentenceTransformer(model_name)
# データのエンコード
embeddings_refrence = model.encode(
sentence_refs,
normalize_embeddings=True,
show_progress_bar=True,
convert_to_tensor=True
)
embeddings_query = model.encode(
sentence_queries,
normalize_embeddings=True,
show_progress_bar=True,
convert_to_tensor=True
)
# Faiss インデックスの構築と保存
# faiss.IndexFlatL2の初期設定として、次元数を設定
faiss_index = faiss.IndexFlatL2(len(embeddings_refrence[0]))
faiss_index.add(embeddings_refrence.detach().cpu().numpy())
if path_out_faiss_full_index_path is not None:
faiss.write_index(faiss_index, path_out_faiss_full_index_path)
# 類似度検索と結果の保存
# 今回は最も類似しているものを返すため1
search_score, idx_list = faiss_index.search(embeddings_query.detach().cpu().numpy().astype(np.float32), 1)
df_out = pd.DataFrame([search_score.flatten(), idx_list.flatten()]).T
df_out.columns = ["Score", "ID"]
df_out["ID"] = df_out["ID"].astype(int)
df_out["Sentence"] = [sentence_refs[idx] for idx in df_out["ID"]]
return df_out
# 実行例
# 1つ目を正例として想定して入れた
article_titles = [
"Vector Searchによる検索",
"Firebase Authのリダイレクトログインを使っている人は今年の6月までに対応しないと大変ですよという注意喚起",
"データ分析基盤まとめ(随時更新)",
]
search_nearest_sentence(article_titles, ["似た文書をベクトル検索で探し出したい ~SentenceTransformersとFaissで効率的にベクトル検索~"])
# => 以下のpandas dataframeが返ってくる
# Score ID Sentence
# 0.278252 0 Vector Searchによる検索
クエリ文書では「ベクトル検索」と日本語なのに対し、リファレンスではVector searchと英語であるにもかかわらず検索を行えていることがわかる
補足説明
ベクトル化を行うモデルについて
Transformerの様々なモデルが公開されています。Massive Text Embedding Benchmark（MTEB）というリーダーボードなどで最新のモデルが比較されていますので、参考にしてください。
なお、日本語特化のモデルはあまり見たことがないため、基本的には多言語化モデルを使う必要があります。（今回の関数では multilingual-e5-base という多言語モデルを使用しています）
faissでのインデックス化について
上記のコードのインデックス化（faiss.IndexFlatL2(len(embeddings_refrence[0]))）のところ。
今回は（もともとKaggleのコードでIndexFlatL2が使われていたことから）IndexFlatL2を使用したが、他のインデックス化の方法も考えられる。以下の公式ページに詳しい。
Discussion