LlamaIndex の RecursiveRetriever を試す

「LlamaIndex」の「RecursiveRetriever」を試したので、まとめました。

1. RecursiveRetriever

「RecursiveRetriever」は、階層的な文書表現 (ネストされた表/図を含むPDFなど) に対して再帰的検索を実行できるリトリバーです。文書内の階層構造を活用することで、複雑な文書のRAGシステムの性能を向上させることができます。

ユースケースは、次のとおりです。

(1) PDF内に埋め込まれた表・画像と残りのテキストのモデリング
(2) 多数の異質な文書のモデリング - 最初に要約によってインデックスを作成し、その内容にリンク。
(3) 複雑なデータに対する検索精度の向上 - 合成のためのテキストチャンクを、検索方法（より小さなチャンク、メタデータ、要約など）から切り離し、より良い検索を実現。
(4) コードブロックのモデリング + ASTのトラバース

クエリ対象となるノードリストに、「テキストチャンク」だけでなく「インデックスへのリンク」を含めることで、階層構造を表現しています。

Recursive Retriever + Query Engine Demo - LlamaIndex 🦙 0.8.17

2. セットアップ

Colabでのセットアップ手順は、次のとおりです。

(1) パッケージのインストール。

# パッケージのインストール
!pip install llama-index

(2) 環境変数の準備。
以下のコードの <OpenAI_APIのトークン> にはOpenAI APIのトークンを指定します。(有料)

import os

# 環境変数の設定
openai.api_key = "<OpenAI_APIのトークン>"

(3) ログレベルの設定。

import logging
import sys

# ログレベルの設定
logging.basicConfig(stream=sys.stdout, level=logging.DEBUG, force=True)

3. ドキュメントの準備

Colabでのドキュメントの準備手順は、次のとおりです。

(1) Wikipediaからのデータ読み込み。
Wikipediaから「魔法少女まどか☆マギカ」「ぼっち・ざ・ろっく!」の情報を取得し、テキストファイルに保存します。

from pathlib import Path
import requests

# Wikipediaからのデータ読み込み
wiki_titles = ["魔法少女まどか☆マギカ", "ぼっち・ざ・ろっく!"]
for title in wiki_titles:
    response = requests.get(
        "https://ja.wikipedia.org/w/api.php",
        params={
            "action": "query",
            "format": "json",
            "titles": title,
            "prop": "extracts",
            # 'exintro': True,
            "explaintext": True,
        },
    ).json()
    page = next(iter(response["query"]["pages"].values()))
    wiki_text = page["extract"]

    data_path = Path("data")
    if not data_path.exists():
        Path.mkdir(data_path)

    with open(data_path / f"{title}.txt", "w") as fp:
        fp.write(wiki_text)

dataフォルダに「魔法少女まどか☆マギカ.txt」「ぼっち・ざ・ろっく!.txt」が出力されます。

(2) ドキュメントの読み込み。

from llama_index import SimpleDirectoryReader

# ドキュメントの読み込み
anime_docs = []
for wiki_title in wiki_titles:
    docs = SimpleDirectoryReader(input_files=[f"data/{wiki_title}.txt"]).load_data()
    docs[0].doc_id = wiki_title
    anime_docs.extend(docs)

4. 子階層のクエリエンジンの準備

Colabでの子階層のクエリエンジンの準備手順は、次のとおりです。

(1) QAテンプレートの準備。
LlamaIndexの英語版のQAプロンプトの直訳になります。

from llama_index.llms.base import ChatMessage, MessageRole
from llama_index.prompts.base import ChatPromptTemplate

# QAシステムプロンプト
TEXT_QA_SYSTEM_PROMPT = ChatMessage(
    content=(
        "あなたは世界中で信頼されているQAシステムです。\n"
        "事前知識ではなく、常に提供されたコンテキスト情報を使用してクエリに回答してください。\n"
        "従うべきいくつかのルール:\n"
        "1. 回答内で指定されたコンテキストを直接参照しないでください。\n"
        "2. 「コンテキストに基づいて、...」や「コンテキスト情報は...」、またはそれに類するような記述は避けてください。"
    ),
    role=MessageRole.SYSTEM,
)

# QAプロンプトテンプレートメッセージ
TEXT_QA_PROMPT_TMPL_MSGS = [
    TEXT_QA_SYSTEM_PROMPT,
    ChatMessage(
        content=(
            "コンテキスト情報は以下のとおりです。\n"
            "---------------------\n"
            "{context_str}\n"
            "---------------------\n"
            "事前知識ではなくコンテキスト情報を考慮して、クエリに答えます。\n"
            "Query: {query_str}\n"
            "Answer: "
        ),
        role=MessageRole.USER,
    ),
]

# チャットQAプロンプト
CHAT_TEXT_QA_PROMPT = ChatPromptTemplate(message_templates=TEXT_QA_PROMPT_TMPL_MSGS)

(2) インデックスとクエリエンジンの作成。

from llama_index import VectorStoreIndex

indexes = []
query_engines = []

for document in anime_docs:
    # インデックスの作成
    index = VectorStoreIndex.from_documents(
        [anime_doc],
    )
    indexes.append(index)

    # クエリエンジンの作成
    query_engine = index.as_query_engine(
        similarity_top_k=3,
        text_qa_template=CHAT_TEXT_QA_PROMPT,
    )
    query_engines.append(query_engine)

(3) まどか☆マギカの質問応答。

# まどか☆マギカの質問応答
response = query_engines[0].query("まどか☆マギカの主題歌を歌っている歌手は？")
print(response)

まどか☆マギカの主題歌を歌っている歌手は、クラリスとKalafinaです。

(4) ぼっち・ざ・ろっく!の質問応答

# ぼっち・ざ・ろっく！の質問応答
response = query_engines[1].query("ぼっち・ざ・ろっく！の作者の名前は？")
print(response)

ぼっち・ざ・ろっく！の作者の名前ははまじあきです。

5. 親階層のクエリエンジンの準備

Colabでの親階層のクエリエンジンの準備手順は、次のとおりです。

(1) 親階層のノードリストとクエリエンジン辞書の準備。
IndexNodeには、textに要約、index_idにインデックスIDを指定します。

from llama_index.schema import IndexNode

# 要約
summaries = [
    "「魔法少女まどか☆マギカ」は、新しいアプローチで魔法少女もののジャンルを刷新したダーク・ファンタジー作品です。物語は魔法少女たちの運命や苦悩、絶望を描いており、彼女たちが魔女と戦う姿が中心です。作品は複数の解釈がありますが、明確な情報は提供されていません。また、複数の会社が関与しており、新房昭之監督とシャフト制作のタッグによるオリジナル作品として注目されています。作品は多くのファンを獲得し、様々な賞を受賞しました。劇場版『新編』叛逆の物語も高い評価を受けています。さまざまなイベントやコラボレーションも行われました。",
    "提供されたテキストには、『ぼっち・ざ・ろっく！』という4コマ漫画の情報が含まれています。物語は主人公の後藤ひとりがバンド活動を通じて成長していく様子を描いています。彼は対人コミュニケーションが苦手であり、自己肯定感の低さやコンプレックスを抱えていますが、バンド活動を通じて成長していきます。彼はギターの演奏技術が高く、作詞も行っていますが、バンドでの実力を存分に発揮することができていません。また、ひとり以外のメンバーには虹夏、郁代、リョウというキャラクターがいます。彼らもバンド活動を通じて成長し、夢を追い求めています。",
]

# ノードリスト
nodes = [
    IndexNode(text=summaries[0], index_id=anime_docs[0].doc_id),
    IndexNode(text=summaries[1], index_id=anime_docs[1].doc_id),
]

# クエリエンジン辞書
query_engine_dict = {
    anime_docs[0].doc_id: query_engines[0],
    anime_docs[1].doc_id: query_engines[1],
}

(2) VectorStoreIndexとVectorRetrieverの準備。
VectorStoreIndexには、先程準備したノードリストを設定します。

# VectorStoreIndexの準備
vector_index = VectorStoreIndex(nodes)

# VectorRetrieverの準備
vector_retriever = vector_index.as_retriever(
    similarity_top_k=1
)

(3) RecursiveRetrieverとクエリエンジンの準備。
RecursiveRetrieverには、先程準備したVectorRetrieverとクエリエンジン辞書を設定します。

from llama_index.retrievers import RecursiveRetriever

# RecursiveRetrieverの準備
recursive_retriever = RecursiveRetriever(
    "vector",
    retriever_dict={"vector": vector_retriever},
    query_engine_dict=query_engine_dict,
)

(4) クエリエンジンの準備。

from llama_index.query_engine import RetrieverQueryEngine
from llama_index.response_synthesizers import get_response_synthesizer

# ResponseSynthesizerの準備
response_synthesizer = get_response_synthesizer(
    # service_context=service_context,
    response_mode="compact",
    text_qa_template=CHAT_TEXT_QA_PROMPT,
)

# クエリエンジンの準備
query_engine = RetrieverQueryEngine.from_args(
    recursive_retriever, 
    response_synthesizer=response_synthesizer
)

(5) まどか☆マギカの質問応答。

# まどか☆マギカの質問応答
response = query_engines[0].query("まどか☆マギカの主題歌を歌っている歌手は？")
print(response)

まどか☆マギカの主題歌を歌っている歌手は、クラリスとKalafinaです。

ログを確認すると、最初に親階層のtop-k=1で「まどか☆マギカ」のノードが選択されています。

Top 1 nodes:
[Similarity score:             0.856092] 「魔法少女まどか☆マギカ」は、新しいアプローチで魔法少女もののジャンルを刷新したダーク・ファンタジー作品です。...

子階層のtop-k=3で質問応答を処理した後、最後に子階層から返された結果で以下の質問応答を行い、最終結果を出力します。

コンテキスト情報は以下のとおりです。
---------------------
Query: まどか☆マギカの主題歌を歌っている歌手は？
Response: まどか☆マギカの主題歌を歌っている歌手は、Kalafinaです。
---------------------
事前知識ではなくコンテキスト情報を考慮して、クエリに答えます。
Query: まどか☆マギカの主題歌を歌っている歌手は？
Answer:

(6) ぼっち・ざ・ろっく!の質問応答

# ぼっち・ざ・ろっく！の質問応答
response = query_engines[1].query("ぼっち・ざ・ろっく！の作者の名前は？")
print(response)

ぼっち・ざ・ろっく！の作者の名前ははまじあきです。

ログを確認すると、最初に親階層のtop-k=1で「ぼっち・ざ・ろっく！」のノードが選択されています。

Top 1 nodes:
[Similarity score:             0.863588] 提供されたテキストには、『ぼっち・ざ・ろっく！』という4コマ漫画の情報が含まれています。...

子階層のtop-k=3で質問応答を処理した後、最後に子階層から返された結果で以下の質問応答を行い、最終結果を出力します。

コンテキスト情報は以下のとおりです。
---------------------
Query: ぼっち・ざ・ろっく！の作者の名前は？
Response: ぼっち・ざ・ろっく！の作者の名前ははまじあきです。
---------------------
事前知識ではなくコンテキスト情報を考慮して、クエリに答えます。
Query: ぼっち・ざ・ろっく！の作者の名前は？
Answer:  

6. DocumentSummaryIndex

「DocumentSummaryIndex」を使うと、要約の自動生成も可能です。

OpenAI GPT-4/ChatGPT/LangChain 人工知能プログラミング実践入門