langchainとDatabricksで(私が)学ぶRAG : Parent Document Retrieverを使ったRAG

OpenAI DevDayが衝撃的でしたね。OSSを焼野原にする勢いかと思ったのですが、langchain-aiもOpenGPTsをいきなり出したりなど、プロプライエタリとOSSの切磋琢磨が止まらない。

導入

私が学ぶRAGの実質2回目です。準備編はこちら。

今回はParent Document Retrieverを使ったRAGを実践します。

langchainの公式Docは以下です。

これは何？

上の公式Docよりざっくり和訳。

検索用に文書を分割するとき、しばしば相反する望みがあります。

埋め込みが最も正確に効果を発揮できるようにするために、小さな文書に分割したいかもしれません。なぜなら、あまりに文章が長いと、埋め込みが効果を失ってしまうからです。
一方、各チャンクのコンテキストが保持されるような、十分な意味を持つ長さのドキュメントを持ちたいとも思うでしょう。
ParentDocumentRetrieverは、小さなデータの塊を分割して保存することで、そのバランスを取っています。検索時には、まず小さなチャンクを取得し、次にそれらのチャンクの親IDを調べて、より大きなドキュメントを返します。

親ドキュメントとは、小さなチャンクの元となったドキュメントのことです。これは生文書全体であったり、大きなチャンクであったりします。

ベクトル検索するときのチャンクは、ある程度小さくないと精度よく検索できないことがあります。
しかし、LLMに渡すチャンクが細かすぎると、十分な参照情報とならないことがあります。
ParentDocumentRetrieverは、これを解消するために検索用のチャンクとLLMに渡すチャンクや文書の塊を分ける手法の一種になります。

より具体的には、検索対象の文章をある程度意味を保持したサイズのチャンクに分割し、それを親ドキュメントとしてさらに細かい検索用チャンク（子チャンク）に分割、そして子チャンクに対して埋め込み(Embedding)を実行してベクトルストアに保持する考え方です。

手前みそですが、以下の記事でやったRAGもこの考え方に従って実装しています。

今回はlangchainが提供するParentDocumentRetrieverを公式Doc通りに実践してみます。

Step0. モジュールインストール

使うモジュールをインストールします。
今回はベクトルストアにChromaを利用。

%pip install -U -qq transformers accelerate autoawq=="0.1.5" langchain chromadb=="0.4.15"
%pip install databricks-feature-engineering 

dbutils.library.restartPython()

Step1. Document Loading

準備編で保管した特徴量を取得します。
読み込んだデータは、docsという名前のビューから参照できるようにします。

from databricks.feature_engineering import FeatureEngineeringClient

fe = FeatureEngineeringClient()

feature_name = "training.llm.sample_doc_features"
df = fe.read_table(name=feature_name)

df.createOrReplaceTempView("docs")

Step2. Splitting

ごくごく単純なText Splitterを使って長文データをチャンキングします。

今回は親ドキュメントと子ドキュメントそれぞれのために、2種のSplitterを用意します。

• 親ドキュメントは2,000文字でオーバーラップ0
• 子ドキュメントは200文字でオーバーラップ40文字

にします。

利用するデータは契約書なので、本来であれば親ドキュメントは、セクションごとなど分け方を工夫するべきです。今回は簡易化のために単純に文字数で分割します。

from typing import Any
import pandas as pd
from pyspark.sql.functions import pandas_udf
from langchain.text_splitter import RecursiveCharacterTextSplitter


class JapaneseCharacterTextSplitter(RecursiveCharacterTextSplitter):
    """句読点も句切り文字に含めるようにするためのスプリッタ"""

    def __init__(self, **kwargs: Any):
        separators = ["\n\n", "\n", "。", "、", " ", ""]
        super().__init__(separators=separators, **kwargs)


parent_text_splitter = JapaneseCharacterTextSplitter(chunk_size=2000, chunk_overlap=0)
child_text_splitter = JapaneseCharacterTextSplitter(chunk_size=200, chunk_overlap=40)

Step3. Storage

チャンクデータに埋め込み(Embedding)を行い、ベクトルストアへデータを保管します。

まずはEmbedding用のモデルをロード。
以下のモデルをダウンロードしたものを利用します。

import torch
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS

device = "cuda" if torch.cuda.is_available() else "cpu"
embedding_path = "/Volumes/training/llm/model_snapshots/models--intfloat--multilingual-e5-large"

embedding = HuggingFaceEmbeddings(
    model_name=embedding_path,
    model_kwargs={"device": device},
)

Chromaを使ってベクトルストアを作成します。
この際に、親ドキュメント用のSplitterと子ドキュメント用のSplitterも合わせて渡します。

from langchain.schema.document import Document
from langchain.vectorstores import Chroma
from langchain.storage import InMemoryStore
from langchain.retrievers import ParentDocumentRetriever

pdf = spark.table("docs").select("page_content").toPandas()
docs = [Document(page_content=d) for d in pdf["page_content"]]

vectorstore = Chroma(
    vectorstore=vectorstore,
    docstore=store,
    child_splitter=child_text_splitter,
    parent_splitter=parent_text_splitter,
)

テストしてみます。

# サンプル
retriever.get_relevant_documents("この契約において知的財産権はどのような扱いなのか？")

出力

[Document(page_content='（実施計画書 （仕様書）の遵守）  \n第１条 乙は、 本契約に明記されていると否とを問わず、 関係法令諸規則 （要綱等を含む。 ）\nを遵守し、 別紙１の実施計画書（仕様書） （以下「仕様書」という。） に従って委託業\n務を実施しなければならない。  \n２ 乙は、自らの責任において 委託業務 を遂行するものとし、 第三者の権利処理 （第三者\nが所有し、又は管理する知的財産権の実施許諾や動産・不動産の使用許可の取得等を含\nむ。）が必要な場合には乙の費用及び責任で行うものとする。 甲の指示により、委託者\n名を明示して業務を行う場合も同様とする。  \n３ 甲は、委託業務及び納入物に関して、約定の委託金額以外の支払義務を負わない。 本\n契約終了後の納入物の利用についても同様とする。 委託金額には委託業務の遂行に必要\nな諸経費並びに消費税及び地方消費税を含む。  \n \n（納入物の提出）  \n第２条 乙は、委託業務についての納入物（以下単に「納入物」という。）を完了期限ま\nでに甲に提出しなければならない。納入物の所有権は、第１３条第１項の検査後、納入\n物が甲に引き渡されたときに、乙から甲に移転する。  \n２ 乙は、納入物を 文書で作成する場合は、国等による環境物品等の調達の推進等に関す\nる法律（平成 １２年法律第１００号）第６条第１項の規定に基づき定められた環境物品\n等の調達の推進に関する基本方針（ 令和５年２月２４日変更閣議決定）によ る紙類の印\n刷用紙及び役務の印刷の基準を満たすこととし、様式第１により作成した印刷物基準実\n績報告書を納入物とともに甲に提出しなければならない。  \n \n（契約保証金）  \n第３条 甲は、本契約に係る乙が納付すべき契約保証金の納付を全額免除する。  \n \n （知的財産 権の帰属及び 使用） \n第４条 本契約の締結時に乙が既に所有又は管理していた 知的財産権（以下「 乙知的財産\n権」という。）を 乙が納入物に使用した場合には、甲は、当該乙知的財産権を、仕様書\n記載の「目的」のため、仕様書の「納入物」の項 に記載した利用方法に従い、本契約終\n了後も期間の制限なく、また追加の対価を支払うことなしに 自ら使用し、又は第三者に\n使用させることができる 。ただし、仕様書に明確な利用方法等が定められていない場合\nには、甲は、仕様書記載の「目的」のために甲が相当と認める方法で自ら使用し、第三\n者に使用させることができる。 なお、本契約において納入物の「使用（利用）」には、\n納入物の改良・改変をはじめとして、あらゆる使用（利用）態様を含む 。また、本契約\nにおいて「知的財産権」とは、知的財産基本法第２条第２項所定の知的財産権をいい、\n知的財産権を受ける権利及びノウハウその他の秘密情報を含む。  \n２ 乙は、 納入物に第三者の知的財産権を利用する場合には、 第１条第２項 の規定に従い、\n乙の費用及び責任において当該第三者から本契約の履行及び本契約終了後の甲による\n納入物の利用 に必要な書面の許諾を得なければならない。なお、第三者より当該許諾に\n条件を付された場合には（以下「第三者の許諾条件」という。）、乙は、納入物に第三\n者の知的財産権を利用する前に、甲に対して第三者の 許諾条件を書面で速やかに通知し\nなければならない。甲は、当該第三者の許諾条件に同意できない場合には、本契約の 解\n約又は変更を含め、乙に対して協議を求めることができる。甲が当該条件に同意した場\n合、乙は、委託業務の遂行及び納入物の作成に 当たって第三者の許諾条件を遵守するこ\nとにつき全責任を負う。   \n３ 甲は、第三者の許諾条件を遵守することを条件として、 本契約終了後も 期間の制限な\nしに、納入物の利用 に必要な範囲で、前項の第三者の知的財産権を 自由かつ対価の追加\n支払なしに 使用し、又は第三者に使用させることができる。  \n４ 委託業務の遂行中に納入物に関して乙（甲の同意を得て一部を再委託する場合は再委\n託先を含む。） が新たに知的財産権（以下 「新規知的財産権」という。）を取得した場\n合には、乙は、その詳細を書面にしたものを納入物に添付して甲に提出するものとする。\n新規知的財産権は 約定の委託金額以外の追加支払なしに、納入物の引渡しと同時に乙か\nら甲に譲渡され、甲単独に帰属する。  \n５ 前項の規定にかかわらず、著作権等について は第２８条の定めに従う。  \n６ 乙は、本契約終了後であっても、知的財産権の取 扱いに関する本契約 の約定を 自ら遵\n守し、及び第７条第１項 の再委託先に遵守させ ることを 約束する。 \n７ 委託業務又は納入物に関して、第三者の知的財産権の侵害 に関する紛争 その他第三者')]

分かりづらいですが、内部的に子ドキュメントに対して検索した後、得られるドキュメントとしては親ドキュメントの内容が取得されます。
これによって、LLMに十分意味のある（かもしれない）コンテキストを渡すことが出来ます。

Step4. Chain creation

RAGを実行するChainを作成します。前回とほとんど同じ。

Prompt Template

簡単なチャットテンプレートを準備。

from langchain.prompts import ChatPromptTemplate
from langchain.prompts.chat import (
    AIMessagePromptTemplate,
    HumanMessagePromptTemplate,
)

template = """次のcontextの内容のみを使い、なるべく平易な文章を使って日本語で質問に回答してください。
{context}

Question: {question}
"""
prompt = ChatPromptTemplate.from_messages(
    [
        HumanMessagePromptTemplate.from_template(template),
        AIMessagePromptTemplate.from_template(""),
    ]
)

LLM

事前にダウンロードしておいた、TheBloke兄貴が変換したCALM2-chatのAWQモデルを利用します。

from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers_chat import ChatHuggingFaceModel

model_path = "/Volumes/training/llm/model_snapshots/models--TheBloke--calm2-7B-chat-AWQ"

generator = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_path)

chat_model = ChatHuggingFaceModel(
    generator=generator,
    tokenizer=tokenizer,
    human_message_template="USER: {}\n",
    ai_message_template="ASSISTANT: {}",
    temperature=0.1,
    max_new_tokens=1024,
)

Chain

これまで作成した構成要素を組み合わせてChainを作成します。
最近のlangchainはLangChain Expression Language (LCEL)の利用が推奨されており、こちらで記載します。

from langchain.schema.output_parser import StrOutputParser
from langchain.schema.runnable import RunnablePassthrough

chain = (
   {"context": retriever, "question": RunnablePassthrough()} 
   | prompt 
   | chat_model 
   | StrOutputParser()
)

Step5. Run

準備は整いましたので、ストリーミング出力で実行してみます。

for s in chain.stream("この契約において知的財産権はどのような扱いなのか？"):
    print(s, end="", flush=True)

出力

この契約において、知的財産権は以下のように扱われます。

第4条（知的財産権の帰属及び使用）では、乙が納入物に使用した知的財産権を、仕様書記載の目的のために、仕様書の納入物項目に記載された利用方法に従って、本契約終了後も期間の制限なく、追加の対価を支払うことなしに、自ら使用し、または第三者に使用させることができます。ただし、仕様書に明確な利用方法等が定められていない場合には、甲は、仕様書記載の目的のために、甲が相当と認める方法で自ら使用し、第三者に使用させることができます。また、本契約において、納入物の「使用（利用）」には、納入物の改良・改変をはじめとして、あらゆる使用態様が含まれます。

ただし、乙は、納入物に第三者の知的財産権を利用する場合には、第1条第2項の規定に従い、乙の費用及び責任において当該第三者から本契約の履行及び本契約終了後の甲による納入物の利用に必要な書面の許諾を得なければなりません。また、乙は、納入物に第三者の知的財産権を利用する場合には、仕様書記載の目的のために、甲が相当と認める方法で自ら使用し、第三者に使用させることができる場合でも、第三者から書面の許諾を得なければなりません。

なお、乙が新たに知的財産権を取得した場合には、その詳細を書面にしたものを納入物に添付して甲に提出する必要があります。また、乙は、本契約終了後も、自ら知的財産権の遵守責任を負い、再委託先に遵守させることを約束します。

まとめ

Parent Document Retrieverを使ったRAGを実践してみました。

今回はまとめて親ドキュメントと子ドキュメントをSplitしましたが、個別に分けるなども可能です。
また、考え方的に「検索するチャンク」と「LLMに渡すチャンク」を分けるというものなので、実装の仕方含めて様々なバリエーションがあります。

このような考え方は以下の記事でSmall-to-Big Retrievalとして紹介されています。

他のバリエーションとして、検索するチャンクの前後のチャンクにも関連した内容が含まれるという考え方から、LLMには検索対象チャンク＋前後チャンクを結合して渡すというようなSentence Window Retrievalなどがあります。

LangChain Templateには、MongoDBを使ったParent Document Retrievalのテンプレートも公開されています。
参考にしてください！

次回はSelf-Queryingの予定です。