Zenn
♾️

【LangChain】ParentDocumentRetrieverのデータを永続化する方法

2024/02/24に公開
Python
faiss
OpenAI
LangChain
tech

概要

ParentDocumentRetrieverのデータを永続化するのにハマったので、備忘録として残します。

結論から言うと、こちら のページの内容を以下のように書き換えれば行けそうです

  • retriver構築時
    • InMemoryByteStoreをLocalFileStore + create_kv_docstoreに置き換える。
    • vectorstoreもsave_loacalでローカルに保存する。
  • retriver使用時
    • ローカルに保存したstoreとvectorstoreを使用してParentDocumentRetrieverを再構築する。

Retrieverについて(知ってる人は読み飛ばしてください。)

LLMを使ったチャットボットを作る際には、参考文献を根拠にLLMに回答を生成させたいことがあります。こういったケースでは、ユーザーの質問に関連した文書をローカルからベクトル検索で抽出してLLMのプロンプトに付与する、いわゆる検索拡張生成(RAG)という手法を使うことが多いです。
この文書検索の機能をRetrieverといい、Langchainではさまざまな実装のRetrieverが提供されています。

この中でも、ParentDocumentRetireverはローカルの文章を細切れ(チャンク)にしてそれぞれベクトル化し、チャンクと元文章とIDで紐づけることで、

  • 検索時 : ユーザーからの質問と類似するチャンクを検索
  • LLMへの情報入力時: チャンクと同じIDの元の文章を使用
    ということができるようになっています。

詳細は以下を参照してください

https://qiita.com/shimajiroxyz/items/facf409b81f59bb68775

https://python.langchain.com/docs/modules/data_connection/retrievers/parent_document_retriever

今回の記事の動機

公式ドキュメントでは、Retrieverのstoreの実装にInMemoryByteStoreというものを使っています。簡単に試すだけならこれでもいいのですが、この実装だとRetrieverのstoreが永続化されません。

Retrieverの構築をチャットボットの起動の度に行うのはできれば避けたいです。チャットボットを起動する度に文書のチャンク化やベクトル化処理が毎回走ってしまい、アプリが立ち上がるのに時間がかかってしまいます。またOpenAI APIなどを使用している場合には、アプリの起動の度にAPIを叩いてしまい、コストの増加につながります。そのため、構築したRetrieverの情報を永続化してローカルの文書に変化がない限りは使い続けるのが理想です。

上記のような需要は結構あると思うので、今回の記事ではRetrieverのデータをどう永続化する手順を紹介していこうと思います。なお、今回はParentDocumentRetrieverに絞って解説しますが、MultivectorRetireverなどstoreを使う他のRetrieverでも応用できるかと思います(未検証)

Retrieverを構築→ローカルに情報を保存

Retrieverを構築してデータをローカルに保存するまでの手順は以下の通りです。
保存するデータは以下2つです。

  • storeの情報
    • チャンク化する前の文書をid付きで保存
  • vector storeの情報
    • チャンク化した文章をEmbeddingsでベクトル化したもの

各種ライブラリを読み込み

from langchain.retrievers import ParentDocumentRetriever
from langchain.storage import LocalFileStore
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings
from langchain.schema import Document
from langchain.storage._lc_store import create_kv_docstore

テキストを用意

texts = [
    "センチメンタル ドライブ 作詞作曲 三鷹アサ 初めてバイクに乗った日に　貴方の体温知りました 頭は無機質だったけど　体は熱いと知りました 思い返してみれば　貴方が助けてくれる時 いつも私の体は冷たくて　冷たくて冷たくて　いつも一人になるけれど　いつも貴方は助けてくれました",
    "ワシの持っている車ににとるなあ・・・これワシの車じゃないか？ウヌは盗人なのか？ワシのじゃ・・・",
    "顏が違えば兄さんじゃないって思ってるの・・・？顏がちがうのは・・・！嫌だろ！",
    "アアアア・・・！！こっちも光ん力で回復ウ・・・！！",
    "今からぁ・・・浮きましゅ・・・　浮いてる．．．す．．．すげぇ．．．浮いた・・・！浮いたぞ・・・！"
]
docs = [Document(t) for t in texts]

storeの構築

チャンク化前の文書を保存するストアをローカルに"test_store"という名前で作ります。

# 前のstoreが残っていたら削除
if os.path.exists("test_store"):
    shutil.rmtree("test_store")

# チャンクに細切れにする前の元文章を保存するストアを構築
fs = LocalFileStore("test_store")
store = create_kv_docstore(fs)

公式のサンプルではInMemoryByteStore()というインメモリ型のストアを使っていますが、これを永続化して別コードで読み込む際にはLocalFileStore + create_kv_docstoreを使います。

なお、LocalFileStoreをcreate_kv_docstoreを介さずに使おうとすると、保存したいドキュメントをbyteに変換しなければいけないので不便です。セットで使いましょう。
https://github.com/langchain-ai/langchain/issues/9345
https://stackoverflow.com/questions/77385587/persist-parentdocumentretriever-of-langchain

vector storeを構築

# openAIのEmbeddingsを読み込み
embeddings = OpenAIEmbeddings(model="text-embedding-3-large")

テキストをベクトル化するライブラリを使って、vector storeを作成します。今回はFAISSを使用しました。

ParentDocumentRetrieverの構築

ParentDocumentRetrieverに

  • store
# 文書を細切れにするためのsplitter
child_splitter = RecursiveCharacterTextSplitter(chunk_size=10, chunk_overlap = 0)

# FAISSで空のベクトルストアを初期化
# 3072次元(text-embedding-3-largeの埋め込み次元数)のダミーのベクトルを追加し、即削除
vectorstore = FAISS.from_embeddings([("test", [0]*3072)], embeddings, ids = [1])
vectorstore.delete([1])

# Retrieverを作成
retriever = ParentDocumentRetriever(
    vectorstore=vectorstore,
    docstore=store,
    child_splitter=child_splitter,
    id_key = "doc_id"
)

ドキュメントをEmbeddingを使ってベクトル化・vector storeをローカルに保存

ParentDocumentRetrieverを定義すると、add_documents()を呼び出すだけで、

  • ドキュメントを細かいチャンクに分割してそれぞれベクトル化
  • 分割後のチャンクと元文章をIDで紐づけてstoreに保存
    という2つを同時にやってくれます。

また、コードの最後でFAISSのsave_loaclを呼び出して、vector storeをローカルに保存します。

# add documentすると文書のチャンク分けとベクトル化が走る
retriever.add_documents(docs)

# add documentでvector storeに埋め込みが追加されたはずなので，ローカルに保存
vectorstore.save_local("test_vector_store")

# LocalFileStoreとvectorStoreの中身を確認
# 上手くいっていれば、
keys = list(retriever.docstore.yield_keys())
print(keys[0])
print("元の文章 : ", retriever.docstore.mget([keys[0]]))
print(
    "分割後のチャンク : ",
    [
        item.page_content for _, item in vectorstore.docstore._dict.items() 
        if item.metadata["doc_id"] == keys[0]
    ]
)

ローカルに情報を読み込み→Retrieverを再構築

先ほどのコードで作成したRetrieverを別のコードで再度呼び出すにはローカルに保存したvectorstoreとstoreを読み込み、ParentDocumentRetrieverに参照を渡せばよいです。

from langchain.retrievers import ParentDocumentRetriever
from langchain.storage import LocalFileStore
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings
from langchain.schema import Document
from langchain.storage._lc_store import create_kv_docstore
import pickle

# openAIのEmbeddingsを読み込み
embeddings = OpenAIEmbeddings(model="text-embedding-3-large")

# The storage layer for the parent documents
fs = LocalFileStore("test_store")
store = create_kv_docstore(fs)

# It should create documents smaller than the parent
child_splitter = RecursiveCharacterTextSplitter(chunk_size=10, chunk_overlap = 0)

# retrieverを作成
retriever = ParentDocumentRetriever(
    vectorstore=FAISS.load_local("test_vector_store", embeddings),
    docstore=store,
    child_splitter = child_splitter,
    id_key = "doc_id",
    search_kwargs = {"k": 1}
)

このコードで作ったRetrieverを使って検索してみると、しっかりとクエリに関連した文章がヒットしてます。

# retrieverで検索
print(retriever.get_relevant_documents("浮遊する俺"))

# 結果
#  [Document(page_content='\n    今からぁ・・・浮きましゅ・・・\n    浮いてる．．．す．．．すげぇ．．．\n    浮いた・・・！浮いたぞ・・・！\n    ')]

Discussion