ベクトル検索ライブラリ Faiss を試す

npaka

2023年1月5日 20:34

ベクトル検索ライブラリ「Faiss」を試したので、使い方をまとめました。

1. Faiss

「Faiss」は、Facebookがリリースしたベクトル検索ライブラリです。

2. テキストを埋め込みに変換

「埋め込み」は、意味的類似性を示すベクトル表現です。2つのベクトル間の距離は、その関連性を表し、小さな距離は高い関連性、大きな距離は低い関連性を示します。

一般的に次のような用途に使用されます。

・検索 : 検索結果がクエリ文字列との関連性でランク付けされる
・クラスタリング : テキストを類似性によってグループ化
・レコメンデーション : 関連するテキストを含む項目を推奨
・異常検出 : 関連性の少ない外れ値を特定
・ダイバーシティ測定 : 類似性分布を分析
・分類 : テキストを最も類似したラベルで分類

Colabでの実行手順は、次のとおりです。

(1) メニュー「編集→ノートブックの設定」の「ハードウェアアクセラレータ」で「GPU」を選択。

(2) パッケージのインストールとAPIキーの準備
以下のコードの <OpenAI_APIのトークン> にはOpenAI APIのトークンを指定します。(有料)

# パッケージのインストール
!pip install openai

# APIキーの設定
import openai
openai.api_key = "<OpenAI_APIのトークン>"

(3) テキストを埋め込みに変換。
OpenAI APIが生成する埋め込みの次元は1536になります。

# テキストを埋め込みに変換
texts = ["これはテストです。"]
response = openai.Embedding.create(input=texts, model="text-embedding-ada-002")
embeds = [record['embedding'] for record in response['data']]

# 確認
print(len(embeds[0]))
print(embeds)

1536
[[-0.0019341740990057588, -0.011006913147866726,... ]]

3. 近傍探索

入力テキストと意味の近いものを、対象テキストから探します。

(1) パッケージのインストール。

# パッケージのインストール。
!pip install faiss-gpu

(2) テキストを埋め込みに変換。

# 入力テキストを埋め込みに変換
in_texts = ["今日は雨振らなくてよかった"]
response = openai.Embedding.create(input=in_texts, model="text-embedding-ada-002")
in_embeds = [record['embedding'] for record in response['data']]

# 対象テキストをベクトル表現に変換
target_texts = [
    "好きな食べ物は何ですか?",
    "どこにお住まいですか?",
    "朝の電車は混みますね",
    "今日は良いお天気ですね",
    "最近景気悪いですね"]
response = openai.Embedding.create(input=target_texts, model="text-embedding-ada-002")
target_embeds = [record['embedding'] for record in response['data']]

(3) インデックスの生成と対象テキストの追加。
1536は埋め込みの次元です。埋め込みをインデックスに追加する際は、numpyのfloat32に変換します。

import faiss 
import numpy as np

# インデックスの生成
index = faiss.IndexFlatL2(1536)
print(index.is_trained)

# 対象テキストの追加
index.add(np.array(target_embeds).astype('float32'))
print(index.ntotal)

True
5

インデックスの検索アルゴリズムの種類は、次のとおりです。

・IndexFlatL2 : L2ノルム
・IndexFlatIP : コサイン類似度
・IndexIVFFlat : 高速化アルゴリズム

(4) 近傍探索の実行。

# 近傍探索の実行。
D, I = index.search(np.array(in_embeds).astype('float32'), 1)

# 確認
print(D)
print(I)
print(I[0][0])
print(target_texts[I[0][0]])

[[0.1780048]]
[[3]]
3
今日は良いお天気ですね

"今日は雨振らなくてよかった" に一番近い対象テキストは "今日は良いお天気ですね" であることがわかりました。