RAGのチャンクサイズと検索手法(Vector Index / Summary Index)について考察してみた。Llama Index, LLM
こんにちは!株式会社IZAI、エンジニアチームです。
今回は、Llama IndexというOSSライブラリを使ってRAGを実装し、その検索精度に与える各パラメータの影響を考察してみます。
そもそもRAGとは
RAG(Retrieval Augmented Generation)とはLLMにデータを入力する際に、 外部の知識ベースからその入力内容に合った情報を探し出し、その情報をLLMに渡すことで回答精度を上げる手法です。 RAG(Retrieval Augmented Generation)を導入する目的はFine-tuningと似ていて、LLMを最新情報に対応させることです。Fine-tuningとの違いは、情報を直接参照するかどうかといった点にあります。RAGの場合、情報源から直接情報を取得し、それに基づいて回答を生成しますが、Fine-tuningはモデルの内部パラメータ自体を更新するという違いがあります。
ただし、RAGは情報源から取得する情報量が多すぎるとトークン制限に引っかかる問題があります。そのため回答に必要な情報を見つけ出してくること(Retrieval)が重要で、本記事ではRetrievalのチャンクサイズや検索手法によって回答精度がどのように変化するか確認していきます。
Llama Indexとは
Llama IndexとはLLMと外部データベース等の接続を容易にするライブラリです。LLM開発ではLangChainと並んでよく使われています。
今回はこのLlama Indexを使ってRAGを実装してみます。
実装
題材
今回は、JAISTのスパコンの(スーパーコンピュータ)の説明書を外部知識としてRAG実装していきます。文字数は約17,000文字あり、人間が全て目を通すのは大変です。
1. ライブラリのインポート
# 使用するライブラリ
from llama_index.core.extractors import TitleExtractor
from llama_index.core.ingestion import IngestionPipeline
from llama_index.core.node_parser import SentenceSplitter
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.readers.file import DocxReader
2. Loading
文字通りデータを読み込むフェーズです。情報を検索するためにデータを小分け(チャンク)にしていきます。
# 情報源をロード
loader = DocxReader()
documents = loader.load_data(file_path)
# Documents形式をNode形式へ変換 & IndexingフェーズでVector Indexingするためにembedding
chunk_size = 1000
chunk_overlap = 100
pipeline = IngestionPipeline(
transformations=[
SentenceSplitter(),
OpenAIEmbedding(model='text-embedding-3-small'),
]
)
nodes = pipeline.run(documents=documents)
3. Indexing
データの中からマッチする内容を検索するために小分けにされている各NodeにIndexを付与していきます。
# Indexing
index = VectorStoreIndex(nodes)
4. Quering
LLMに質問します。
# Quering
chat_engine = index.as_chat_engine()
response_chat = chat_engine.chat('〇〇とはなんですか?')
# 〇〇はスパコンの名前
これでRAGの構築が完了しました。Llama Indexを用いることで少ないコードでRAGが完成します。
5. チャンクサイズ検証用コード
ここから、チャンクサイズによる実験です。
ここではLoadingフェーズでデータを小分けにする際のチャンクサイズが検索精度にどのように影響するか調べます。
chunk_size_list = [100, 500, 900, 1000, 3000, 5000, 7000, 8000]
for i in chunk_size_list:
# Documents形式をNode形式へ変換 & IndexingフェーズでVector Indexingするためにembedding
chunk_size = i
chunk_overlap = int(chunk_size*0.1)
pipeline = IngestionPipeline(
transformations=[
SentenceSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap),
OpenAIEmbedding(model='text-embedding-3-small'),
]
)
nodes = pipeline.run(documents=documents)
# Indexing
index = VectorStoreIndex(nodes)
chat_engine = index.as_chat_engine()
response_chat = chat_engine.chat('〇〇とはなんですか?')
print('chunk_size : {} responce : {}'.format(chunk_size, response_chat))
結果
chunk_size : 100 responce : 〇〇はファイル名です。
chunk_size : 500 responce : 〇〇は、高速アクセス用ストレージや高速ストレージ一時貸出(プロジェクト向け)などの情報を含む文書の名前です。
chunk_size : 900 responce : 〇〇は、高速アクセス用ストレージや高速ストレージ一時貸出(プロジェクト向け)など、計算機からアクセスできる特定のストレージサービスを指す用語です。
chunk_size : 1000 responce : 〇〇は、計算機リソースやイメージ管理に関する情報を提供する文書のタイトルです。
chunk_size : 3000 responce : 〇〇は、計算機リソースを利用してジョブを実行する際に使用されるシステムの名称です。このシステムを使用することで、ジョブの投入や実行、リソースの確保などが行われます。
chunk_size : 5000 responce : 〇〇はGPU搭載システムで、ログインしてインタラクティブジョブを実行したり、バッチ実行したり、計算タスクにコンテナを利用したりすることができるシステムです。システムへのアクセス方法や、インタラクティブまたはバッチモードでのジョブ実行方法、DockerイメージをSingularityイメージに変換して計算ジョブで使用する方法についての指示が提供されています。
chunk_size : 7000 responce : 〇〇は、MathematicaなどのGUIを使用する場合に役立つツールです。Mathematicaは数学や科学技術分野で使用されるソフトウェアであり、〇〇はそのGUIを操作する際に役立つツールのようです。
chunk_size : 8000 responce : 〇〇は、GPUコンピューティング向けに設計された高性能コンピューティングシステムです。
全体を通して間違った説明や遠回りな説明が目立つ結果となりました。強いて言うならchunk_size=8000の回答が最もらしいですがRAGとしてはもっと小さいchunk_sizeで回答して欲しいですね。
この結果から、今回のケースではチャンクをそのままembeddingモデルでベクトル情報に変換しただけでは、検索(retrieval)がうまくいかないことがわかります。
そこで次に、Document Summary Index (SummaryIndex)を使用して再度試してみまます。
SummaryIndex はその名の通り、チャンクごとに要約を生成しその埋め込みによって検索する手法で、ハイコンテクストな情報をもとに検索できる。無関係の情報(ノイズ)の影響を受けにくいという特徴があります。
chunk_size : 100 responce : 〇〇 is a high-performance computing (HPC) system at JAIST designed to achieve large-scale and high-speed computations. It consists of various parallel architectures to provide high processing capabilities. Users affiliated with the university can utilize these machines for free 24 hours a day. The system offers benefits such as access to large-scale computing with hundreds of CPU cores and hundreds of gigabytes to terabytes of memory, batch job scheduling, and a wide range of software and libraries for research purposes.
chunk_size : 500 responce : 〇〇は、高い処理能力を持つ並列計算機群であり、大規模・高速計算を実現することを目的として導入されています。利用者は24時間無料で利用でき、数百のCPUコアや大容量メモリを利用することができます。また、バッチジョブを登録しておくことでシステムが自動的に実行し、豊富なソフトウェア・ライブラリも利用可能です。
chunk_size : 900 responce : 〇〇はJAISTが保有する大規模計算機システムの名称です。このシステムは大規模・高速計算を実現するために様々な並列アーキテクチャを持った並列計算機群を導入しており、構成員であれば24時間無料で利用することができます。利用者は数百のCPUコアや大容量メモリを利用でき、バッチジョブを登録しておけばシステムが自動的に実行するなど、多くのメリットがあります。
chunk_size : 1000 responce : 〇〇は、大規模・高速計算を実現するための高性能な並列計算機群です。利用者は24時間無料で利用でき、バッチジョブを登録しておくことでシステムが自動的に実行することができます。さらに、豊富なソフトウェアやライブラリも提供されており、研究に利用される様々なソフトウェアが利用可能です。
chunk_size : 3000 responce : 〇〇 is an HPC system designed to achieve large-scale and high-speed computing. It consists of a variety of parallel architectures to provide high processing capabilities. Users affiliated with the university can utilize these machines for free 24 hours a day. The system offers benefits such as large-scale computing with hundreds of CPU cores and large memory capacities, batch job capabilities for efficient execution, and a wide range of software and libraries including Matlab, Materials Studio, and Gaussian09 for research purposes.
chunk_size : 5000 responce : 〇〇は大規模で高速な計算向けのHPCシステムです。高い処理能力を持つ複数の並列アーキテクチャを備えています。大学関連のユーザーは24時間365日無料でこれらのマシンにアクセスできます。このシステムは、数百のCPUコアを使用した大規模な計算、効率的なリソース利用のためのバッチジョブスケジューリング、研究用のMatlab、Materials Studio、Gaussian09などのさまざまなソフトウェアを提供しています。
chunk_size : 7000 responce : 〇〇は、高い処理能力を持つ並列計算機群であり、大規模・高速計算を実現するために導入されたHPCシステムです。利用者は24時間無料で利用でき、数百のCPUコアや大容量メモリを利用することができます。さらに、バッチジョブや豊富なソフトウェア・ライブラリを利用することができます。
chunk_size : 8000 responce : 〇〇は大規模で高速な計算を行うためのHPCシステムであり、高い処理能力を持つさまざまな並列アーキテクチャを特徴としています。大学関連のユーザーは、無料で24時間365日アクセスでき、数百のCPUコア、大規模なメモリ容量、バッチジョブ機能、Matlab、Materials Studio、Gaussian09などのさまざまなソフトウェアやライブラリを利用することができます。研究目的で利用されています。
Summary Indexの結果
SummaryIndexを使った結果、全体として正しい回答が得られました。
chunk_size=500程度でも正しい検索結果を得られていることから、十分RAGとして実用的であることがわかります。
考察
VectorIndexの精度が悪い理由
仮説を立ててみます。
情報源とした説明書の前処理/整理が足りていない
Splitterとの相性が悪かった
今回使用したembeddingモデル(OpenAIのtext-embedding-3-small)との相性が悪かった →試しにtext-embedding-3-largeを使ってみたらさらに悪化しました。(「I don’t know」しか言わなくなりました。。。)
こちらも機会があれば検証していきたいと思います。
チャンクサイズはいくつがベストなのか
チャンクサイズは小さいほどピンポイントで情報を絞り込めるはずです。ただし, その分回答のために参照するノードが多くなるので、検索にかかる処理時間は大きくなります。また、retrievalの精度という意味ではchunk_sizeを小さくするほど不利ということになります。
ただし、チャンクサイズを大きくすると、チャンク内に含まれる不要な情報を参照することになり精度低下を招き、トークン使用量も定数倍で大きくなっていきます。
パフォーマンス・精度・コストの兼ね合いで検索対象の文書の大きさや特性、求める回答の精度や形式によってチャンクサイズや検索手法の正解は異なりそうです。
chunk_size : 100
responce : 〇〇は大規模で高速な計算を行うための高性能コンピューティング(HPC)システムです。このシステムは高い処理能力を持つさまざまな並列アーキテクチャを特徴としています。大学のユーザーは、24時間無料でこれらの計算機にアクセスすることができます。システムは数百のCPUコアを使用した大規模な計算、大容量のメモリ、バッチジョブスケジューリング、研究目的のさまざまなソフトウェアやライブラリなどの利点を提供しています。並列計算機の使用方法に関する詳細な手順は提供されているドキュメントに記載されています。
retrieval result : 204
chunk_size : 500
responce : I dont have any information on what "〇〇" is.
retrieval result : 32
chunk_size : 900
responce : 〇〇は大規模で高速な計算に使用される高性能コンピューティングシステムです。異なる並列アーキテクチャを備え、高い処理能力を提供しています。大学のユーザーは24時間無償でこれらのマシンにアクセスできます。システムには多数のCPUコア、広範囲のメモリ容量、効率的なリソース利用のためのバッチジョブスケジューリング、Matlab、Materials Studio、Gaussian09などの研究ニーズに対応するさまざまなソフトウェアやライブラリが備わっています。
retrieval result : 18
chunk_size : 1000
responce : 〇〇は大規模で高速な計算を行うための高性能コンピューティング(HPC)システムです。さまざまな並列アーキテクチャを組み込んでおり、高い処理能力を提供しています。大学に所属するユーザーは、24時間無料でこれらのマシンにアクセスできます。システムには数百のCPUコア、広範な計算のための大容量メモリ、効率的な実行のためのバッチジョブ機能、Matlab、Materials Studio、Gaussian09などの幅広いソフトウェアやライブラリを利用できる利点があります。
retrieval result : 16
chunk_size : 3000
responce : 〇〇はHPCシステムの名前であり、大規模・高速計算を実現するための並列計算機群です。このシステムは、様々な並列アーキテクチャを持ち、数百のCPUコアや大容量メモリを利用できる特徴があります。利用者は24時間無料で利用可能であり、バッチジョブや豊富なソフトウェア・ライブラリを利用することができます。
retrieval result : 5
chunk_size : 5000
responce : I dont have any information on 〇〇.
retrieval result : 3
chunk_size : 7000
responce : 〇〇は大規模で高速な計算向けに設計されたHPCシステムであり、高い処理能力を持つさまざまな並列アーキテクチャを提供しています。大学に所属するユーザーは24時間無料でこれらのマシンにアクセスできます。このシステムは、数百のCPUコアと大規模なメモリ容量を備えた大規模な計算、効率的なリソース利用のためのバッチジョブスケジューリング、Matlab、Materials Studio、Gaussian09などのさまざまなソフトウェアやライブラリを利用できることなどの利点を提供しています。Matlab、MathematicaなどのGUIを使う場合に役立ちます。
retrieval result : 3
chunk_size : 8000
responce : 〇〇 is an HPC system designed for large-scale and high-speed computing. It consists of various parallel architectures to provide high processing capabilities. Users affiliated with the university can utilize these machines for free 24 hours a day. The system offers benefits such as large-scale computing with hundreds of CPU cores and large memory capacities, batch job scheduling for automatic execution based on resource availability, and a wide range of software and libraries including Matlab, Materials Studio, and Gaussian09 for research purposes.
retrieval result : 2
chunk_size_list = [100, 500, 900, 1000, 3000, 5000, 7000, 8000]
q = '〇〇とはなんですか?'
for i in chunk_size_list:
# Documents形式をNode形式へ変換 & IndexingフェーズでVector Indexingするためにembedding
chunk_size = i
chunk_overlap = int(chunk_size*0.1)
pipeline = IngestionPipeline(
transformations=[
SentenceSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap),
]
)
nodes = pipeline.run(documents=documents)
# Indexing
index = SummaryIndex(nodes)
chat_engine = index.as_chat_engine()
response_chat = chat_engine.chat(q)
retriever = index.as_retriever()
reference_node = retriever.retrieve(q)
print('chunk_size : {}\\nresponce : {}\\nretrieval result : {}\\n'.format(chunk_size, response_chat, len([node.text for node in reference_node])))
参考
[1] 生成AIのトレンド「LLM+RAG」を解説
[2] RAGフレームワーク LlamaIndexの概要を整理してみる
[3] 公式ドキュメント