概要
前回 文章を数値化し faiss にデータを保存しました
今回は実際に faiss に検索をかけて文章の類似度を計算して似ている文章を見つけたいと思います
環境
- macOS 15.2
- transformers 4.47.0
- faiss-cpu 1.9.0-post1
コード全体
from torch import Tensor
from transformers import AutoModel, AutoTokenizer
# このaverage_pool関数はbertの出力結果ではよく使われる手法、何をしているかの詳細は以下の説明に記載
def average_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
# faissに保存する文章、この文章と類似度検索をする
input_texts = [
"好きな食べ物は何ですか?",
"どこにお住まいですか?",
"朝の電車は混みますね",
"今日は良いお天気ですね",
"最近景気悪いですね",
"最近、出かけていないので、たまには外で食事でもどうですか？",
]
# モデルとトークナイザの取得、日本語なので日本語に特化したモデルを使用
tokenizer = AutoTokenizer.from_pretrained("cl-tohoku/bert-base-japanese")
model = AutoModel.from_pretrained("cl-tohoku/bert-base-japanese")
# 文章の数値化
inputs = tokenizer(input_texts, padding=True, truncation=True, return_tensors="pt")
# 文章のベクトル化、768 次元のベクトル情報が取得できる (参考: https://github.com/cl-tohoku/bert-japanese)
outputs = model(**inputs)
# bert の結果を更によくするためのおまじない
embeddings = average_pool(outputs.last_hidden_state, inputs["attention_mask"])
# faiss で扱えるように numpy 配列に変換
embeddings_np = embeddings.cpu().detach().numpy()
# なぜか先頭でimportすると segmentation fault になるのでここで import
import faiss
# 次元数(768)で初期化
index_flat_l2 = faiss.IndexFlatL2(embeddings_np.shape[1])
# ベクトル情報を追加、6文章(vectors)登録される
index_flat_l2.add(embeddings_np) # type: ignore
# クエリとなるテキストの定義、数値化、faiss 用ベクトル化
query_text = ["今日は晴れている"]
query_dict = tokenizer(
query_text, max_length=512, padding=True, truncation=True, return_tensors="pt"
)
query_output = model(**query_dict)
query_embeddings = average_pool(
query_output.last_hidden_state, query_dict["attention_mask"]
)
query_embeddings_np = query_embeddings.cpu().detach().numpy()
# 上位何件を取得するかの設定
k = 2
# 検索の実行
D, I = index_flat_l2.search(query_embeddings_np, k) # type: ignore
# 結果の表示
for i in range(k):
print(
f" Rank: {i+1}, Index: {I[0][i]}, Distance: {D[0][i]}, Text: '{input_texts[I[0][i]]}"
)
今回追加した faiss で検索をする部分は以下です
# クエリとなるテキストの定義、数値化、faiss 用ベクトル化
query_text = ["今日は晴れている"]
query_dict = tokenizer(
query_text, max_length=512, padding=True, truncation=True, return_tensors="pt"
)
query_output = model(**query_dict)
query_embeddings = average_pool(
query_output.last_hidden_state, query_dict["attention_mask"]
)
query_embeddings_np = query_embeddings.cpu().detach().numpy()
# 上位何件を取得するかの設定
k = 2
# 検索の実行、結果はD(distance)とI(index)の情報が返ってくる
D, I = index_flat_l2.search(query_embeddings_np, k) # type: ignore
# 結果の表示
for i in range(k):
print(
f" Rank: {i+1}, Index: {I[0][i]}, Distance: {D[0][i]}, Text: '{input_texts[I[0][i]]}"
)
最後に
faiss で文章の類似度計算をしてみました
文章の類似度を計算する手法はいろいろとありますが特に自分で何も計算することなく search するだけで勝手に似たような文章を持ってきてくれるので簡単に使えます
今回は IndexFlatL2 を使いましたが他にも「IndexFlatIP」「IndexHNSWFlat」などがあります
参考: https://github.com/facebookresearch/faiss/wiki/Faiss-indexes
データの規模やマシンスペックによって適宜切り替えれば OK です
使い方によって多少データの加工が必要ですがほぼ同じように使えます
流れとしては以下で可能だとわかりました
- transformers で文章を数値化、ベクトル化
- 結果を faiss に格納
- クエリを作成し検索を実施
bert のテンソルベクトル情報を faiss で扱えるように numpy形式のベクトル情報に変換するところが少しクセがありますがそこ以外は数値化した情報をそのまま faiss に渡すだけなので簡単に使えました
また数値化する際に今回は cl-tohoku/bert-base-japanese を使いましたがモデルとトークナイザを変えるだけでも結果は変わってくるかなと思います
0 件のコメント:
コメントを投稿