Go to Qiita Advent Calendar 2024 Top
LoginSignup
2
8

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

@daikumatan(長尾 太介 (Daisuke Nagao))

FAISSによるベクトル検索 Hands-On

Posted at

FAISS Hands-On

テストコードを動かして以下の4つの検索手法をなんとなく理解することを目的とします

  • IndexFlatL2
  • IndexFlatIP
  • IndexIVFFlat
  • IndexHNSW

事前準備

pip install faiss-cpu

Embedding

FAISSに格納するデータをembeddingする

embedding用のコードを以下にしめす。後にFAISSに格納するembedding対象の日本語は以下の5つである

  • "好きな食べ物は何ですか?",
  • "どこにお住まいですか?",
  • "朝の電車は混みますね",
  • "今日は良いお天気ですね",
  • "最近景気悪いですね",
# 必要なライブラリをインポートします。
import torch.nn.functional as F
from torch import Tensor
import faiss  # FAISSライブラリをインポート
import numpy as np  # NumPyライブラリをインポート
from transformers import AutoTokenizer, AutoModel

# 最後の隠れ層の状態を平均プーリングする関数を定義します。
def average_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
    # print(attention_mask) # 情報があるところが1, paddingが0で表示
    # print(attention_mask[..., None]) # 次元を一つ追加して3次元にしている
    # print(attention_mask[..., None].bool()) # 1をTrue, 0をFalseへ変換
    # print(~attention_mask[..., None].bool()) # TrueとFalseを入れ替え (paddingがTrueになる)
    
    # last_hidden_statesの値におけるpaddingの値が0.0に置き換えられる
    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)

    # last_hidden.sum(dim=1)でシーケンシャル方向に和を取る
    # つまり、[<batch_size>, <sequence_length>, <隠れ層の次元数>] => [<batch_size>, <隠れ層の次元数>] に変換
    # attention_mask.sum(dim=1)[..., None]は、padding除いたToken数の和。この値で割って平均化する
    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]

# 解析したいテキストのリストを定義します。
input_texts = [
    "好きな食べ物は何ですか?",
    "どこにお住まいですか?",
    "朝の電車は混みますね",
    "今日は良いお天気ですね",
    "最近景気悪いですね",
    #"最近、出かけていないので、たまには外で食事でもどうですか？"
]

# 使用するモデルとトークナイザーの事前学習済みのパスを指定してロードします。
tokenizer = AutoTokenizer.from_pretrained('intfloat/multilingual-e5-large')
model = AutoModel.from_pretrained('intfloat/multilingual-e5-large')

# テキストをモデルが処理できる形式にトークン化します。
batch_dict = tokenizer(input_texts, max_length=512, padding=True, truncation=True, return_tensors='pt')

# batch_dict内の各テンソルのサイズを確認する
for i, (key, tensor) in enumerate(batch_dict.items()):
    print(f"[{i}][batch_dict] {key}: {tensor.size()}")

# モデルを実行して出力を取得します。
outputs = model(**batch_dict)

# batch_dict内の各テンソルのサイズを確認する
for i, (key, tensor) in enumerate(outputs.items()):
    print(f"[{i}][outputs] {key}: {tensor.size()}")

# `average_pool`関数を使って、最終的な埋め込みベクトルを計算します。
# print("attention_mask:", batch_dict['attention_mask'].sum(dim=1))
# print("outputs.last_hidden_state.sum(dim=1).shape:", outputs.last_hidden_state.sum(dim=1).shape)
# print("last_hidden_states:", outputs.last_hidden_state.sum(dim=1))
embeddings = average_pool(outputs.last_hidden_state, batch_dict['attention_mask'])

# PyTorchのTensorからNumPy配列に変換します。
embeddings_np = embeddings.cpu().detach().numpy()
print("embeddings_np:", embeddings_np.shape)

# わかりやすさのための表示
for i, vec in enumerate(embeddings_np[:]):
    print(f"----{i}----")
    print(input_texts[i])
    print(vec)

以下のような結果を得て、Inputした文章が1024次元のベクトル化されていることが分かる

[0][batch_dict] input_ids: torch.Size([5, 10])
[1][batch_dict] attention_mask: torch.Size([5, 10])
[0][outputs] last_hidden_state: torch.Size([5, 10, 1024])
[1][outputs] pooler_output: torch.Size([5, 1024])
embeddings_np: (5, 1024)
----0----
好きな食べ物は何ですか?
[ 0.6108155  -0.2260578  -0.10909441 ... -0.62022376 -0.83260953
  0.28850102]
----1----
どこにお住まいですか?
[ 0.99075234  0.6782665  -0.68046105 ... -0.7541518  -1.3915895
 -0.14331307]
----2----
朝の電車は混みますね
[ 1.1466048   0.16233358 -0.02154386 ...  0.07610887 -0.13563985
  0.2248166 ]
----3----
今日は良いお天気ですね
[ 0.9829717   0.06835324 -0.34694287 ...  0.00633434 -0.7837953
 -0.28379193]
----4----
最近景気悪いですね
[ 0.26057956  0.5884724  -0.49848753 ... -0.19543631 -0.8918747
 -1.4221884 ]

query用文章のembedding

同様にquery用の文章をembeddingしておきます。クエリ用の日本語は "今日は雨が振らなくてよかった"です。

# クエリテキストを定義します。
query_text = ["今日は雨が振らなくてよかった"]

# クエリテキストをトークナイズし、モデルが処理できる形式にします。
query_dict = tokenizer(query_text, max_length=512, padding=True, truncation=True, return_tensors='pt')

# モデルを実行して出力を取得します。
query_output = model(**query_dict)

# `average_pool`関数を使って、クエリの埋め込みベクトルを計算します。
query_embeddings = average_pool(query_output.last_hidden_state, query_dict['attention_mask'])

# クエリの埋め込みベクトルをNumPy配列に変換します。
query_embeddings_np = query_embeddings.cpu().detach().numpy()

FAISSへのデータ格納

今回は以下の4つの方法でデータを格納した。

  • IndexFlatL2
    • ベクトル間のユークリッド距離（L2距離）を使用して類似性を計測します。全データ点に対して線形探索を行うため、計算は正確ですが、大規模データセットでは計算コストが高くなります。
  • IndexFlatIP
    • いわゆるcos類似度でありベクトル間の内積を用いて類似性を評価します。この指標は類似度が高いほど大きな値となるため、類似性の高いベクトルを効果的に特定できますが、大きなデータセットには適していません。
  • IndexIVFFlat
    • 先にデータをクラスタリングして索引を作成し、検索時には近似的な検索を行います。これにより、IndexFlatL2やIndexFlatIPと比較して高速な検索が可能ですが、精度は若干落ちます。
  • IndexHNSW
    • 階層的ナビゲータブル小世界（Hierarchical Navigable Small World）グラフを使用し、近似最近傍探索を非常に高速に行います。大規模なデータセットに適しており、距離計算の回数を劇的に削減することができます。
# ベクトルの次元数を設定
dimension = embeddings_np.shape[1]

# ベクトルを正規化
normalized_embeddings_np = embeddings_np / np.linalg.norm(embeddings_np, axis=1, keepdims=True)

# IndexFlatL2インデックスの初期化とデータ追加
index_flat_l2 = faiss.IndexFlatL2(dimension)
index_flat_l2.add(embeddings_np)  # 正規化なしで追加
print("Number of vectors in the IndexFlatL2:", index_flat_l2.ntotal)

# IndexFlatIP（内積を使用するインデックス）の初期化と正規化されたデータの追加
index_flat_ip = faiss.IndexFlatIP(dimension)
index_flat_ip.add(normalized_embeddings_np)  # 正規化されたデータを追加
print("Number of vectors in the IndexFlatIP:", index_flat_ip.ntotal)

# IndexIVFFlat（量子化されたインデックス）の初期化
nlist = 4  # クラスタ数
quantizer = faiss.IndexFlatL2(dimension)  # 量子化器としてIndexFlatL2を使用
index_ivf_flat = faiss.IndexIVFFlat(quantizer, dimension, nlist, faiss.METRIC_L2)
index_ivf_flat.train(embeddings_np)
index_ivf_flat.add(embeddings_np)
print("Number of vectors in the IndexIVFFlat:", index_ivf_flat.ntotal)

# IndexHNSW（階層的ナビゲーティブ小世界グラフを使用するインデックス）の初期化とデータ追加
index_hnsw = faiss.IndexHNSWFlat(dimension, 10)  # ここで10はHNSWのグラフの近隣接続数
index_hnsw.add(embeddings_np)
print("Number of vectors in the IndexHNSW:", index_hnsw.ntotal)

Queryの実行

# 上位3つの近いベクトルを検索するための設定（k=3）
k = 2

# クエリの処理（以前のコードからのクエリテキストの部分を続ける）
query_embeddings_np = query_embeddings.cpu().detach().numpy()
# クエリの埋め込みベクトルも正規化
normalized_query_embeddings_np = query_embeddings_np / np.linalg.norm(query_embeddings_np, axis=1, keepdims=True)

# 各インデックスに対して検索を実行し、結果を取得
D_l2, I_l2 = index_flat_l2.search(query_embeddings_np, k)
D_ip, I_ip = index_flat_ip.search(normalized_query_embeddings_np, k) # 正規化されたクエリを使用
D_ivf, I_ivf = index_ivf_flat.search(query_embeddings_np, k)
D_hnsw, I_hnsw = index_hnsw.search(query_embeddings_np, k)

# 結果を表示する関数
def display_results(D, I, input_texts, index_type):
    print(f"Results for {index_type}:")
    for i in range(k):
        print(f"  Rank: {i+1}, Index: {I[0][i]}, Distance: {D[0][i]}, Text: '{input_texts[I[0][i]]}")

# 各インデックスタイプの結果を表示
display_results(D_l2, I_l2, input_texts, "IndexFlatL2")
display_results(D_ip, I_ip, input_texts, "IndexFlatIP")
display_results(D_ivf, I_ivf, input_texts, "IndexIVFFlat")
display_results(D_hnsw, I_hnsw, input_texts, "IndexHNSW")

以下のような結果をえます。この例では"今日は雨が振らなくてよかった"に文章として一番近く感じる、"今日は良いお天気ですね"が、どの手法でも選ばれていることが分かります。

IndexFlatIPに関してはベクトルの長さを正規化しているため必ず0から1になる値が返ってきています。

Results for IndexFlatL2:
  Rank: 1, Index: 3, Distance: 202.58700561523438, Text: '今日は良いお天気ですね
  Rank: 2, Index: 4, Distance: 267.9387512207031, Text: '最近景気悪いですね
Results for IndexFlatIP:
  Rank: 1, Index: 3, Distance: 0.8827050924301147, Text: '今日は良いお天気ですね
  Rank: 2, Index: 4, Distance: 0.8428647518157959, Text: '最近景気悪いですね
Results for IndexIVFFlat:
  Rank: 1, Index: 3, Distance: 202.58700561523438, Text: '今日は良いお天気ですね
  Rank: 2, Index: -1, Distance: 3.4028234663852886e+38, Text: '最近景気悪いですね
Results for IndexHNSW:
  Rank: 1, Index: 3, Distance: 202.58700561523438, Text: '今日は良いお天気ですね
  Rank: 2, Index: 4, Distance: 267.9387512207031, Text: '最近景気悪いですね

まとめ

いろいろ調べていると、以下のようにまとめることができそうです

  • 小規模データで何も考えず、とりあえず高精度に使いたい場合は IndexFlatL2
  • データが正規化されているなら IndexFlatIP が最も精度が高い
  • 大規模データで高速処理が必要なら IndexHNSWは最も期待できる

Appendix

IndexHNSWのチューニングパラメータ

IndexHNSWの検討すべきパラメータとしては以下があります。必要に応じてチューニングしましょう。

  • M(Default: 16):
    • 各ノードが保持する近傍の数。大きいほど精度は向上しますが、構築とメモリコストが増加します
  • efConstruction (Default: 200):
    • インデックス構築時の動的なパラメータで、探索の深さを定義します
  • ef(Default: 200):

    • 検索時の動的なパラメータで、より広い範囲を探索します。大きいほど検索精度が向上しますが、検索時間が増加します

    • 検討すべき処理: パフォーマンスと精度のバランスを取るためにMとefを調整します。大規模なデータセットでは、これらのパラメータが検索性能に大きな影響を与えるため、実際のデータ特性に応じたチューニングが必要になります

2
8
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
2
8

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?