LoginSignup
11
9
@shimajiroxyz

【langchain】retrieverで検索文章とLLMに渡す文章を別にする方法およびretrieverの保存・読込方法の検討

Last updated at Posted at 2023-10-15

はじめに

langchainで検索拡張生成（RAG）を実装するときに、検索用の文章とLLMに渡す用の文章を分ける方法を整理しました。

使えそうなretrieverの候補として、MultiVectorRetrieverとParentDocumentRetrieverがあるので、それらの使い分けを理解するための整理を行いました。

環境

langchainのバージョンは0.0.314です。

背景

検索拡張生成（RAG）のシンプルな実装は、検索してヒットした文章をそのままLLMに渡すというものです。しかし検索に適した文章の粒度とLLMに渡すのに適した文章の粒度が同じである保証はありません。そのため、検索用の文章とLLMに渡す用の文章を別々にしたいという要求が生じます。

一般には、検索用の文章よりも、LLMに渡す文章を長くしたいケースが多いです。LLMに文章を渡すときは、周辺情報も含めた文章を渡した方が精度が上がりやすいからです。逆に、検索用の文章は長すぎると、特徴が平均化されてしまうため、LLMに渡す文章よりは短い文章が適していることが多いです。

したがって、短めの文章で検索して、それを含むLLMに渡す用の（より長い）文章を取得するという使い方が想定されます。この用途で使えそうなものを、langchainのretrieverから探すと、MultiVectorRetrieverとParentDocumentRetrieverが見つかりました。
それぞれの使い方について整理します。

MultiVectorRetriever

1つのdoc（=LLMに渡す文章）に複数の埋め込みベクトル(=検索用の文章）を対応づけるためのRetrieverです。

挙動確認

チュートリアルを参考に、挙動を確かめます。

ライブラリを読み込みます。

# https://python.langchain.com/docs/modules/data_connection/retrievers/parent_document_retriever
import uuid
from langchain.retrievers import MultiVectorRetriever
from langchain.vectorstores.faiss import FAISS
from langchain.embeddings import OpenAIEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.storage import InMemoryStore
from langchain.docstore import InMemoryDocstore
from langchain.document_loaders import TextLoader
from langchain.schema import Document
from langchain.schema.embeddings import Embeddings

チュートリアルではvectorstoreにChromaを使っていますが、今回は筆者が使い慣れたFAISSで試します。
FAISSはChromaと違って空のvectorstoreを用意する標準機能がないため、空のFAISS vectorstoreを得るための関数を作っておきます。

参考:【langchain】FAISSで空のvectorstoreを作る方法

def get_empty_faiss_vectorstore(embedding: Embeddings, dim = None, **kwargs):
    dummy_text, dummy_id = "1", 1
    if not dim: # 次元数が未知の場合
        dummy_emb = embedding.embed_query(dummy_text)
    else: #次元数が既知の場合
        dummy_emb = [0]*dim
    vectorstore = FAISS.from_embeddings([(dummy_text, dummy_emb)], embedding, ids=[dummy_id], **kwargs)
    vectorstore.delete([dummy_id])
    return vectorstore

空のvectorstoreとdocstoreでMultiVectorRetrieverを初期化します。
id_keyはデフォルトが"doc_id"なので今回の場合は省略してもいいのですが、わかりやすさのために明示的に与えています。
docstoreにはInMemoryStoreのインスタンスを与えます。langchainにはInMemoryDocstoreというクラスもあるのですが、それを与えるとエラーになります。名前が似ていますが、InMemoryStoreはほぼdict、InMemoryDocstoreはdict型のインスタンス変数とそれへの要素追加、削除用のメソッドを持つクラスです。MultiVectorRetrieverのdocstore引数にはdictとして使えるものを渡す必要があるので、InMemoryDocstoreだとエラーになってしまいます。

ID_KEY = "doc_id"

# The vectorstore to use to index the child chunks
vectorstore = get_empty_faiss_vectorstore(OpenAIEmbeddings(), 1536)
# The storage layer for the parent documents
store = InMemoryStore()
retriever = MultiVectorRetriever(
    vectorstore=vectorstore, 
    docstore=store,
    id_key = ID_KEY 
)

空のMultiVectorRetrieverを作ったら、retrieverのインスタンス変数として保持されているvectorstoreとdocstoreにsubdoc、docを加えます。

まずサンプルのdocとsubdocを用意します。本記事の文脈に沿えば、docがLLMに渡す文章、subdocが検索用の文章です。今回は、2つのdocを用意し、それぞれを開業で分割した文章をsubdocとしています。目的に応じて、subdocをdocと全く独立に定義することも可能です。

doc_greeting = """おはよう
こんにちは
さようなら"""
subdocs_greeting = doc_greeting.splitlines()

doc_subject = """国語
算数
理科
社会"""
subdocs_subject = doc_subject.splitlines()

doc_idとmetadataを作ります。doc_idをid_keyのキーに対応するvalueとしてmetadataに含めることで、subdocからdocを参照するのに使います。

docs = [doc_greeting, doc_subject]
doc_ids = [str(uuid.uuid4()) for _ in docs]
subdocs = [subdocs_greeting, subdocs_subject]
flatten_subdocs = subdocs_greeting + subdocs_subject

metadatas = []
for subdoc, doc_id in zip(subdocs, doc_ids):
    metadatas += [{ID_KEY: doc_id} for _ in subdoc]

retrieverのインスタンス変数であるvectorstore、docstoreにそれぞれsubdoc、docsを追加します。

retriever.vectorstore.add_texts(flatten_subdocs, metadatas)
retriever.docstore.mset(list(zip(doc_ids, docs)))

検索を実行します。subdocではなくdocを検索結果として得ることに成功しています。

retriever.get_relevant_documents("国語")

['国語\n算数\n理科\n社会', 'おはよう\nこんにちは\nさようなら']

ソースコード

MultiVectorRetrieverのソースコードを読んで、どのように検索結果とdocの対応付けが行われているか確認します。

_get_relevant_documentsという関数で以下の処理が行われています。

        sub_docs = self.vectorstore.similarity_search(query, **self.search_kwargs)
        # We do this to maintain the order of the ids that are returned
        ids = []
        for d in sub_docs:
            if d.metadata[self.id_key] not in ids:
                ids.append(d.metadata[self.id_key])
        docs = self.docstore.mget(ids)
        return [d for d in docs if d is not None]
  • vectorstore.similarity_searchでsubdocを検索
  • subdocのmetadataからdoc_idを取得
  • docstoreからdoc_idに対応するdocを取得

という流れです。シンプルですね。

一点注意があるとすると、retrieverに文書の取得数として与えるkはself.vectorstore.similarity_searchに与えるsearch_kwargsとして使われています。MultiVectorRetrieverの実装では、このあとに親docの選定が行われるので、実際に取得できる数はkより少なくなります。最終結果をk件取得したい場合は、多めの値をkに与える必要があります。

ParentDocumentRetriever

挙動確認

ParentDocumentRetrieverはMultiVectorRetrieverを継承したクラスです。
そのため、MultiVectorRetrieverと類似の方法で初期化、使用できます。

まずライブラリのimportと空のvectorstoreを作る関数を用意します。

# https://python.langchain.com/docs/modules/data_connection/retrievers/parent_document_retriever
from langchain.retrievers import ParentDocumentRetriever
from langchain.vectorstores.faiss import FAISS
from langchain.embeddings import OpenAIEmbeddings
from langchain.storage import InMemoryStore
from langchain.schema.embeddings import Embeddings
from langchain.schema import Document
from langchain.text_splitter import CharacterTextSplitter


def get_empty_faiss_vectorstore(embedding: Embeddings, dim = None, **kwargs):
    dummy_text, dummy_id = "1", 1
    if not dim: # 次元数が未知の場合
        dummy_emb = embedding.query(dummy_text)
    else: #次元数が既知の場合
        dummy_emb = [0]*dim
    vectorstore = FAISS.from_embeddings([(dummy_text, dummy_emb)], embedding, ids=[dummy_id], **kwargs)
    vectorstore.delete([dummy_id])
    return vectorstore

vectorstoreとdocstoreが空の状態でParentDocumentRetrieverを初期化します。
MultiVectorRetrieverと異なるのは、child_splitterを与えている点です。
child_splitterはdoc（parent_doc)をsubdocに分割するためのsplitterです。
今回は、改行で区切るだけのシンプルなsplitterとしています。
ちなみに、parent_splitterも初期化時のオプション引数として与えることができます。parent_splitterは文字通りparent_docをsplitするためのsplitterです。LLMに渡すという用途を考えた時に、parent_docが長すぎるとLLMのプロンプトに収まらないことがあり得るため、それを防ぐために、parent_docを適切な長さに分けるためのsplitterです。

ID_KEY = "doc_id"

# The vectorstore to use to index the child chunks
vectorstore = get_empty_faiss_vectorstore(OpenAIEmbeddings(), 1536)
# The storage layer for the parent documents
store = InMemoryStore()
# The splitter to split parent documents to smaller chunks
child_splitter = CharacterTextSplitter(separator="\n", chunk_size=1, chunk_overlap=0)
retriever = ParentDocumentRetriever(
    vectorstore=vectorstore, 
    docstore=store,
    child_splitter = child_splitter,
    id_key = ID_KEY 
)

サンプルのdocを定義してadd_documentsで加えます。docはDocumentクラスである必要があります。
subdocの作成やdoc_idの紐付けはadd_documentの内部で実行されます。このため、MultiVectorRetrieverよりもすっきりしたコードになります。

doc_greeting = """おはよう
こんにちは
さようなら"""

doc_subject = """国語
算数
理科
社会"""

docs = [Document(page_content=doc) for doc in [doc_greeting, doc_subject]]

retriever.add_documents(docs)

検索を実行してみます。

retriever.get_relevant_documents("挨拶")

[Document(page_content='おはよう\nこんにちは\nさようなら'),
 Document(page_content='国語\n算数\n理科\n社会')]

上手くできました。

ソースコード

ParentDocumetRetrieverのソースコードを確認してみます。

ParentDocumentRetrieverはMultiVectorRetrieverを継承したクラスです。
新たに定義されたメソッドはadd_documentsだけです。

add_documentsで以下の処理が書かれています。

        if self.parent_splitter is not None:
            documents = self.parent_splitter.split_documents(documents)
        if ids is None:
            doc_ids = [str(uuid.uuid4()) for _ in documents]
            if not add_to_docstore:
                raise ValueError(
                    "If ids are not passed in, `add_to_docstore` MUST be True"
                )
        else:
            if len(documents) != len(ids):
                raise ValueError(
                    "Got uneven list of documents and ids. "
                    "If `ids` is provided, should be same length as `documents`."
                )
            doc_ids = ids

        docs = []
        full_docs = []
        for i, doc in enumerate(documents):
            _id = doc_ids[i]
            sub_docs = self.child_splitter.split_documents([doc])
            for _doc in sub_docs:
                _doc.metadata[self.id_key] = _id
            docs.extend(sub_docs)
            full_docs.append((_id, doc))
        self.vectorstore.add_documents(docs)
        if add_to_docstore:
            self.docstore.mset(full_docs)

流れは以下のとおりです。

  • parent_splitterが存在していれば、documentsを分割したもので置き換える
  • 各documentにdoc_idを付与する
  • 各documentをchild_splitterで分割してsubdocにし、subdocのmetadataにdoc_idを加える
  • vectorstore、docstoreにsubdoc、documentsを追加する

要するに、MultiVectorRetrieverではクラスの外部で実施する必要があった、subdocとdocの対応付けを、クラスの内部でやってくれるようになったのがParentDocumentRetrieverです。

考察

2つのRetrieverの使い分け

docをTextSplitterで実現可能な程度に単純なルールで分割したものをsubdocとしてよいのであれば、ParentDocumentRetrieverを使うのが簡単です。

subdocがdocの部分文字列とは限らない場合や分割のルールが複雑な場合は、MultiVectorRetrieverを使う必要があります。

どちらのRetrieverもできないこと

BM25やTF-IDFなど単語ベースの検索とMultiVectorRetriever(ParentDocumentRetriever含む）を組み合わせる標準的なやり方は用意されていません。
MultiVectorRetrieverは初期化時にvectorstoreを与える必要がありますが、BM25やTF-IDFで検索をするvectorstoreを作ることが難しいからです（一応、embedding_functionを独自に作って代入することで、実現不可能ではないとは思われます）

MultiVectorRetrieverを継承したクラスを定義して、初期化時に、vectorstoreではなくretrieverをうけとるようにできれば、この問題は解消されます。変更点は、_get_relevant_documents内で実行される、vectorstore.similarity_searchをretriever.get_relevant_documentsに変更するだけなので、そこまで難しくないです。機会があればやり方を検討してみようと思います。検索とLLMに渡す文章を別にしたいという要求は埋め込み検索だけで生じるわけではないので、vectorstoreに限らない一般的なretrieverを前提とした作りの方が使いやすい気はします。

Retrieverの保存と読込

MultiVectorRetriever（やParentDocumentRetriever）を一回作ったら、それをローカルに保存して使いまわしたくなることがあると思います。しかし、vectorstoreと違い、retriever（BaseRetriever)を保存や読込する標準のAPIは用意されておらず、retrieverの種類に応じた保存方法を考える必要があります。retreiverによってはsave_localなどのメソッドが用意されていることもありますが、MultiVectorRetrieverに関してはそういったものはなさそうです。

MultiVectorRetrieverを保存する場合は、vectorstoreとdocstoreをそれぞれ保存するのが良いと思います。以下ではvectorstoreの保存にsave_localを用いていますが、vectorstoreの種類によってはsave_localメソッドをもっていない可能性もあるため、必要に応じて保存処理を置き換えてください。

from langchain.retrievers import MultiVectorRetriever
from pathlib import Path
import json
DEFAULT_VECTORSTORE_DIR = "vectorstore"
DEFAULT_DOCSTORE_JSON = "docstore.json"
def save_multi_vector_retriever(retriever: MultiVectorRetriever, save_dir: str
                                , *
                                , vectorstore_dir = DEFAULT_VECTORSTORE_DIR
                                , docstore_json = DEFAULT_DOCSTORE_JSON):
    # vectorstoreがsave_localメソッドを持つ場合のみ使える関数。

    vectorstore = retriever.vectorstore
    docstore_dict = {k: dict(document) for k, document in retriever.docstore.store.items()}

    save_dir: Path = Path(save_dir)
    save_dir.mkdir(exist_ok=True)
    if not hasattr(vectorstore, "save_local"):
        raise AttributeError("vectorstore does not have 'save_local' method")
    vectorstore.save_local(save_dir.joinpath(vectorstore_dir))
    with open(save_dir.joinpath(docstore_json), "w") as f:
        json.dump(docstore_dict, f, ensure_ascii=False, indent=2)
save_multi_vector_retriever(retriever, "output")

なお、pickleでバイナリとして保存するのも手軽ですが、OpenAIEmbeddingを使っている場合、APIキーも一緒に保存されてしまうのでお勧めしません。

読み込み時もvectorstoreとdocstoreをそれぞれ読み込んでインスタンスを復元します。以下のコードでは、vectorstoreのクラスがload_localメソッドを持つことを前提としています。当てはまらない場合は、適宜処理を書き換えてください。またload_localを用いる場合には、embeddingを合わせて外から与える必要があることに注意してください。

from langchain.storage import InMemoryStore
from langchain.schema import Document
def load_multi_vector_retriever(input_dir, embedding
                                ,*
                                , vectorstore_dir = DEFAULT_VECTORSTORE_DIR
                                , docstore_json = DEFAULT_DOCSTORE_JSON
                                , vectorstore_class = FAISS
                                , retriever_class = MultiVectorRetriever
                                , **kwargs
                                )->MultiVectorRetriever:
    

    input_dir = Path(input_dir)
    vectorstore = vectorstore_class.load_local(input_dir.joinpath(vectorstore_dir), embedding)
    with open(input_dir.joinpath(docstore_json)) as f:
        docstore_dict = json.load(f)
    
    docstore_dict = {k: Document(**v) for k,v in docstore_dict.items()}
    store = InMemoryStore()
    store.mset(docstore_dict.items())

    return retriever_class(
        vectorstore=vectorstore
        , docstore = store
        , **kwargs
    )
embedding = OpenAIEmbeddings()
loaded_retriever = load_multi_vector_retriever("output", embedding)

保存も読込もだいぶ複雑な処理になってしまいました。MultiVectorRetrieverには保存・読込のAPIが存在しないので、上記のようにせざるを得ないと思うのですが、もしよりスマートな方法をご存知の方がいたら教えてください。

おわりに

検索とLLMに渡す文章を別にしたいという要求に応えるために使えるlangchainのretrieverの使い方、違いを整理しました。subdocの検索方法がvectorstoreに基づくものに限られるなど、少し発展途上な感はあります。

今回はlangchainに注目しましたが、RAGのためのライブラリとしてはllama_indexも有名です。llama_indexでも同様の機能があるはず、かつ、検索に関してはllama_indexのほうが洗練されている印象があるので、llama_indexの実装も機会があれば見てみたいです。

11
9
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
11
9