Faiss Readerを使って、クエリに類似したノードのみのインデックスを作成する

データアナリティクス事業本部 インテグレーション部 機械学習チームの鈴木です。

Faiss Readerを使い、クエリに類似したテキストをFaissを使って調べ、類似したもののみを使ってインデックスを作成する方法を確認したのでご共有します。

本記事の内容

以下のLlamaIndexの『Faiss Reader』のドキュメントを参考にしました。

Faissに登録したEmbedding済みのテキストのベクトルとクエリのベクトルの類似度を計算し、指定した個数の類似するテキストのみを使ったインデックスの作成を試しました。

Faissに登録するベクトルの作成は、以下のドキュメントを参考にしました。GPTListIndexを使ってノードに分割し、そのテキスト部分を取り出した後、OpenAIのモデルを使ってベクトル化しました。

この記事では、llama_index.embeddings.openai.OpenAIEmbeddingのデフォルトのモデル（記事執筆時点だとtext-ada-embedding-002）を使いました。

クライアント実行環境

実行環境はGoogle Colaboratoryを使いました。ハードウェアアクセラレータ無し、ランタイム仕様は標準としました。

Pythonのバージョンは以下でした。

!python --version
# Python 3.10.12

また、ライブラリは以下のようにインストールしました。

# インデックス作成で使用
!pip install llama-index
!pip install python-dotenv

# 保存環境でSimpleWebPageReaderにて使用
!pip install html2text

Faissはpip install faiss-cpuにてfaiss-cpu 1.7.4として公開されているwheelをインストールしました。

# Faiss
!pip install faiss-cpu

# Collecting faiss-cpu
#   Downloading faiss_cpu-1.7.4-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (17.6 MB)
#      ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 17.6/17.6 MB 72.3 MB/s eta 0:00:00
# Installing collected packages: faiss-cpu
# Successfully installed faiss-cpu-1.7.4

Faissのレポジトリにて案内されているインストール方法についてはINSTALL.mdをご確認ください。

インストールされたライブラリのバージョンは以下でした。

!pip freeze | grep -e "openai" -e "llama-index" -e "langchain"
# langchain==0.0.220
# langchainplus-sdk==0.0.19
# llama-index==0.6.37
# openai==0.27.8

環境変数の設定

まず、必要な値を.envファイルに書き込みます。

!echo 'OPENAI_API_KEY="<トークン>"' >.env

load_dotenv()で環境変数を読み込みました。

from dotenv import load_dotenv
load_dotenv()

やってみる

1. ライブラリのインポート

以下のようにライブラリをインポートしました。

import faiss
from llama_index import Document
from llama_index import GPTListIndex
from llama_index import SimpleWebPageReader
from llama_index import ServiceContext
from llama_index.callbacks import CallbackManager, LlamaDebugHandler
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.readers.faiss import FaissReader
import numpy as np

2. ノードのベクトルの作成とFaissへの登録

続いて、SimpleWebPageReaderで機械学習チームで提供しているサービスの紹介ページを読み込みました。

documents = SimpleWebPageReader(html_to_text=True).load_data([
    "https://classmethod.jp/services/machine-learning/",
    "https://classmethod.jp/services/machine-learning/data-assessment/",
    "https://classmethod.jp/services/machine-learning/recommend/"
])

GPTListIndexを使ってノードに分割した後、OpenAIのモデルを使ってベクトルにしました。特にid_to_text_mapはFaissReaderで期待される形にしておきました。

list_index = GPTListIndex.from_documents(documents)

# 実装時点でデフォルトはtext-ada-embedding-002
embed_model = OpenAIEmbedding()

docs = []
id_to_text_map = {}
for i, (_, node) in enumerate(list_index.storage_context.docstore.docs.items()):
    text = node.get_text()
    docs.append(embed_model.get_text_embedding(text))
    id_to_text_map[i] = text
docs = np.array(docs)

text-ada-embedding-002から出力されるベクトル長を指定して、Faissにベクトルを登録しました。

# dimensions of text-ada-embedding-002
d = 1536
index = faiss.IndexFlatL2(d)
index.add(docs)

3. Faissからの類似するノードの取り出し

クエリをベクトル化し、kでFaissで何個の似たベクトルを探すか指定して、FaissReaderで類似するテキストを取得しました。

# クエリとFaissから取り出すノード数の設定
query_text = "機械学習支援サービスの内容ついて教えてください。"
k = 2

# ベクトル化
query = embed_model.get_text_embedding(query_text)
query=np.array([query])

# Faissからのノードの取り出し
reader = FaissReader(index)
documents = reader.load_data(query=query, id_to_text_map=id_to_text_map, k=k)

kで指定した通り、2つ取り出せました。

len(documents)
# 2

4. 取り出したノードでのインデックスの作成

Faissで確認した類似したノードを使って、GPTListIndexを作成しました。

# デバッグ用
llama_debug_handler = LlamaDebugHandler()
callback_manager = CallbackManager([llama_debug_handler])
service_context = ServiceContext.from_defaults(callback_manager=callback_manager)

# GPTListIndexの作成
index = GPTListIndex.from_documents(documents, service_context=service_context)

# **********
# Trace: index_construction
#     |_node_parsing ->  0.010935 seconds
#       |_chunking ->  0.00549 seconds
#       |_chunking ->  0.00408 seconds
# **********

作成したインデックスを使って回答の生成を実行しました。

query_engine = index.as_query_engine()
response = query_engine.query(query_text)
# **********
# Trace: query
#     |_query ->  15.361931 seconds
#       |_retrieve ->  0.000708 seconds
#       |_synthesize ->  15.36102 seconds
#         |_llm ->  15.336748 seconds
# **********

以下のような回答を得られました。

for i in response.response.split("。"):
    print(i + "。")

クラスメソッドの機械学習支援サービスでは、AWSが提供している機械学習サービスを活用して、お客様のご要望に沿ったBIツールのご紹介・導入支援を行います。

AWSの機械学習サービスには、Amazon SageMaker、Amazon Personalize、Amazon Forecast、Amazon Comprehend、Amazon Rekognitionなどがあります。

また、データ診断を行い、AWSデータ&アナリティクスコンピテンシー認定を取得したエンジニアが支援します。

さらに、レコメンドシステムプランを提供しています。

考察

前回、同じテキストを使ってGPTVectorStoreIndexで回答生成を試した際は、Retrieveにかかる時間が0.14秒程度でした。

今回Faissを使った場合は、単純な比較はできないものの、GPTListIndexによるRetrieveにかかる時間はノード数が減った影響もあり0.0007秒程度で、FaissReaderによる類似度計算も非常に高速だったため、十分に速くなったと言ってもよさそうでした。

※ 初回のFaissReaderによるドキュメント取得にかかった目安時間

少ないテキスト数で検証を行ったのでそれほど明らかな差は出ませんでしたが、類似度計算の対象が多い場合はもっと速くなることが予想されました。

最後に

FaissReaderを使って、Faissに登録したノードの埋め込みベクトルから、クエリに類似したものを取り出してGPTListIndexを使って回答を生成する例をご紹介しました。

Data Connectorsのモジュールガイドの一覧にはほかのベクターストアのReaderもありました。

Retrieveにかかる時間が短縮できそうなので、試してみたいなと思います。

参考になりましたら幸いです。