こんにちは、CCCMKホールディングス TECH LABの三浦です。
歳を重ねていくと、なんとなく一年の中でのこの時期は特に体調を崩しやすいな、ということが分かってきます。私にとっては今の時期がちょうどその時期で、今年もやっぱり風邪を引いてしまいました。来年はもう少し自分の"勘"を信じて風邪を引かないようにしようと思います。
さて、最近大量のドキュメントからその中に含まれる重要なトピックスだけを抜き出すことが出来ないかな、と考える機会がありました。文章量が多いドキュメントは一度にLLMに取り込むことが難しいため、なんらかの形でドキュメントを細かく分ける必要があります。一番簡単なアイデアはドキュメントを文字数を基準に分割し、分割したドキュメントの塊ごとに重要なトピックを抜き出し、それらを集約して文章全体の重要なトピックを作る、というアプローチです。
このアプローチでは最初の重要トピックを抜き出すときにドキュメント内の連続する文章だけを対象にしている点で上手くいかないことがあります。ドキュメント内のあちこちで触れられているトピックなのかどうかを判定することが出来ないためです。
この課題に対処出来そうなアプローチが、RAG(Retrieval-augmented generation)に関する論文の中で提案されていて、とても面白そうだったので試してみることにしました。
今回のブログでは、最初に論文の内容をまとめ、論文で提案されているアプローチを試す方法までをご紹介したいと思います。
参照論文
今回参照した論文はこちらの論文です。
Title: RAPTOR: Recursive Abstractive Processing for Tree-Organized Retrieval
Author: Parth Sarthi, Salman Abdullah, Aditi Tuli, Shubh Khanna, Anna Goldie, Christopher D. Manning
Submit: Submitted on 31 Jan 2024
arXiv: https://arxiv.org/abs/2401.18059v1
RAPTORの処理の流れ
まず論文で提案されているRAPTORという手法の処理の内容についてまとめてみます。RAPTORでは最初にドキュメントの木構造化を行い、その後に検索を実行します。
ドキュメントの木構造化
クラスタの構成
RAPTORは最初にドキュメントをトークン数を基準に細かく分割します。その後、分割されたテキスト(チャンク)ごとに埋め込みモデルによって埋め込みベクトルの生成を行います。ここまでは一般的なRAGと同じ手続きです。
RAPTORはこの後、ドキュメント全体から似た意味を持つチャンクの塊を生成します。これは埋め込みベクトルに対するクラスタリングによって実現されます。
クラスタリングは"Gaussian Mixture Models (GMMs)"が使用されます。GMMsはソフトクラスタリングの手法のため、所属確率のしきい値の設定によっては同じチャンクが複数のクラスタに所属するケースも発生します。GMMsを行う前に最適なクラスタ数の探索も行っていて、こちらは"Bayesian Information Criterion(BIC)"を基準に決定します。
また、GMMsは高次元ベクトルを対象にした場合に上手くいかないことがあるそうで、RAPTORではまずUMAPによる次元削減を行って次元削減後のベクトルを対象にGMMsによるクラスタリングを行います。UMAPでの次元数は論文の中には記載が恐らくなかったと思いますが、GitHubの実装を見ると、次元数10がデフォルト値として使用されています。
クラスタの要約
クラスタリング後、次に各クラスタに所属するチャンクの内容をLLMを使って要約します。さらに要約文に対しても埋め込みモデルによって埋め込みベクトルを計算しておきます。
クラスタ要約文のクラスタリング
クラスタごとに生成された要約文を、さらにクラスタリングします。クラスタリングの手続きは、ドキュメントのチャンクに対するものと同様です。要約文のクラスタ生成後、さらにそのクラスタの要約文を生成します。以降同じ処理を繰り返していくことで、ドキュメントの木構造が構築されます。論文に掲載された、以下の図のようになります。
検索処理
構築されたドキュメントの木構造に対するクエリ方法として、 "tree traversal"と"collapsed tree"の2つの方法が提案されています。"tree traversal"は木構造の上位からクエリの埋め込み表現と近似したノードを決まった分検索し、ノードの子ノードに対してさらに同様に近似したものを検索し・・・というように、木構造に沿って関連するノードを取得する方法です。一方"collapsed tree"の方は木構造を無視して全てのノードに対して近似したものを検索する方法です。2つの方法の比較結果によると、"collapsed tree"の方が良い結果になることが多いようです。
自分のドキュメントを木構造化したい
ここからは実際にRAPTORを自分の保有するドキュメントで試す方法について紹介します。最もすぐに試すことが出来る方法は、LangChainのCookbookに含まれるNotebookを利用することだと思います。
私も今回このNotebookを利用させて頂いて自分の保有するドキュメントに対する木構造化を行うことが出来ました。ただいくつか変更が必要な箇所があったので、その点だけここではご紹介したいと思います。
umapのインポート
Notebookではumap-learnというライブラリを使ってUMAPの次元削減を実行しています。恐らくライブラリ側の仕様変更なのかな、と思うのですが、Notebookのまま実行すると、"module 'umap' has no attribute 'UMAP'"といったエラーが表示されてしまいます。
この対応のため、import umapをimport umap.umap_ as umapに変更しました。
要約文生成プロンプトの変更
NotebookはLangChainの"LangChain Expression Language"に関するドキュメントを読み込むことを想定した作りになっていて、特にクラスタごとの要約文を生成するためのプロンプトがその用途に特化したものになっています。具体的にはembed_cluster_summarize_texts関数の内部でそのプロンプトが定義されているので、こちらを目的の用途に合ったものに変更する必要がありました。
実行してみる
まず、ドキュメントの読み込みを行います。このブログの自分の記事を対象にしました。
from langchain_community.document_loaders import UnstructuredURLLoader import nltk nltk.download('punkt') nltk.download('averaged_perceptron_tagger') urls = [ "https://techblog.cccmkhd.co.jp/entry/2024/10/01/145518", "https://techblog.cccmkhd.co.jp/entry/2024/09/24/164833", "https://techblog.cccmkhd.co.jp/entry/2024/08/02/083413" ] loader = UnstructuredURLLoader(urls) data = loader.load()
チャンクを生成します。日本語の文章なので、句読点"。"も区切り文字として指定しました。論文だとトークンサイズ100でチャンキングをしていたので、それに近い200をchunk_sizeに設定しました。ただやってみた感じだともっと大きいサイズでも良いかな、という印象を持ちました。
from langchain_text_splitters import RecursiveCharacterTextSplitter text_splitter = RecursiveCharacterTextSplitter( separators=[ "\n\n", "\n", " ", ".", ",", "\u3002", # Ideographic full stop("。"に該当) ], chunk_size=200 ) chunks = text_splitter.split_documents(data)
使用するモデルを読み込みます。LLMはOllamaで起動している"gemma2:27b"を使用し、埋め込みモデルはHugging Faceで公開されている"intfloat/multilingual-e5-large"を使用しました。
from langchain_huggingface import HuggingFaceEmbeddings from langchain_ollama.llms import OllamaLLM model = OllamaLLM(model="gemma2:27b") # EmbeddingModel model_name = "intfloat/multilingual-e5-large" model_kwargs = {'device': 'cuda'} encode_kwargs = {'normalize_embeddings': False} embd = HuggingFaceEmbeddings( model_name=model_name, model_kwargs=model_kwargs, encode_kwargs=encode_kwargs )
そしたらNotebookで定義されているrecursive_embed_cluster_summarizeを呼び出し、ドキュメントの木構造化が実行できます。
ちなみに私はDatabricksのNotebookでこの処理を実行したのですが、クラスタ数探索時のGMMsの実行結果が全てMLflowにモデルを含めて記録されてしまうことを避けるため、autologを無効にしました。
import mlflow mlflow.sklearn.autolog(disable=True) texts = [chunk.page_content for chunk in chunks] result = recursive_embed_cluster_summarize(texts)
デフォルトだと3階層の木構造が生成されます。生成された結果は辞書型になっていて、それぞれのキーに階層番号が、値にその階層のノードに対するクラスタリングの結果と各クラスタの要約文が格納されています。
たとえば一番上位の3階層目に含まれるノードとクラスタリングの結果と、各クラスタの要約文を確認することが出来ます。
まとめ
今回はRAGのテクニックである"RAPTOR"で用いられているドキュメントの木構造化について、その方法と実際に試してみるための手順について調べたことをご紹介させて頂きました。 ドキュメントを木構造化出来ると、たとえばマインドマップのように可視化してドキュメント全体を確認することが出来るので、とても良いな、と思いました！
