LlamaIndex v0.10のQueryEngineをカスタムしてみる

2024/05/08

こんにちは。少し投稿期間が空きましたが、今回はQueryEngineを自作することにトライします。

モチベーションとしては、より良い検索のためにハイブリッド検索をしたいとなった場合にこの過程が必要になるためです。

今回はこちらのドキュメントをベースに紹介していきます。

Usage Pattern - LlamaIndex

クエリエンジンの作成（復習）

クエリエンジンを作成するには、次の1コードで済みました。

query_engine = index.as_query_engine(
    response_mode="tree_summarize",
    verbose=True,
)

このようにindexから直接クエリエンジンを構築しました。

より詳細な制御をする場合、例えば、検索器の指定など、は各オブジェクトを明示的に構築します。

from llama_index.core import VectorStoreIndex, get_response_synthesizer
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core.query_engine import RetrieverQueryEngine

# インデックスの構築
index = VectorStoreIndex.from_documents(documents)

# 検索器の設定
retriever = VectorIndexRetriever(
    index=index,
    similarity_top_k=2,
)

# 応答合成器の設定
response_synthesizer = get_response_synthesizer(
    response_mode="tree_summarize",
)

# クエリエンジンのアセンブル(構築)
query_engine = RetrieverQueryEngine(
    retriever=retriever,
    response_synthesizer=response_synthesizer,
)

# クエリ
response = query_engine.query("What did the author do growing up?")
print(response)

カスタムクエリエンジンの定義

カスタムクエリエンジンの定義も可能です。

• `CustomQueryEngine`をサブクラス化

• 必要な属性を定義

• `Response`オブジェクトか文字列を返す`custom_query`関数を実装

from llama_index.core.query_engine import CustomQueryEngine
from llama_index.core.retrievers import BaseRetriever
from llama_index.core import get_response_synthesizer
from llama_index.core.response_synthesizers import BaseSynthesizer

class RAGQueryEngine(CustomQueryEngine):
    """RAGクエリエンジン。"""
    
    retriever: BaseRetriever
    response_synthesizer: BaseSynthesizer
    
    def custom_query(self, query_str: str):
        nodes = self.retriever.retrieve(query_str)
        response_obj = self.response_synthesizer.synthesize(query_str, nodes)
        return response_obj

公式ドキュメントではCustom Query Engine guideで詳細を確認するように指示されています。

ひとまず、ここまでのクラス作成までをまとめた実装コードがこちらです。

Google Colab

驚いたことにこれだけに回答のクオリティがかなり向上しています。なぜかは実行しただけでは分かりません。

Custom Query Engineの定義

ここでは、カスタムクエリエンジンの定義についてさらに詳しくみていきます。

Defining a Custom Query Engine - LlamaIndex

セットアップ

設定

サンプルデータをロードして、インデックス付けを行います。

%pip install llama-index-llms-openai
!pip install llama-index

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

データのダウンロード

今回はインデックス作成済みのデータを読み込みます。

!git clone https://github.com/ken20c/corolla.git

すでにインデックス作成済みのため、ドキュメントの読み込みは不要です。

# ドキュメント読み込み
# documents = SimpleDirectoryReader("/path/to/directory")

# index = VectorStoreIndex.from_documents(documents)
retriever = index.as_retriever()

カスタムクエリエンジンの構築

RAGパイプラインを動かすカスタムクエリエンジンを構築します。最初に検索を実行し、次に合成を実行します。

`CustomQueryEngine`を定義するには、いくつかの初期化パラメータを属性として定義して、`custom_query`関数を実装するだけです。

デフォルトでは、custom_query は応答合成器が返すResponseオブジェクトを返すことができますが、単に文字列を返すこともできます、これは、下に示すオプションの1と2です。

from llama_index.core.query_engine import CustomQueryEngine
from llama_index.core.retrievers import BaseRetriever
from llama_index.core import get_response_synthesizer
from llama_index.core.response_synthesizers import BaseSynthesizer

Opetion 1(RAGQueryEngine)

class RAGQueryEngine(CustomQueryEngine):
    """RAGクエリエンジン"""

    retriever: BaseRetriever
    response_synthesizer: BaseSynthesizer

    def custom_query(self, query_str: str):
        nodes = self.retriever.retrieve(query_str)
        response_obj = self.response_synthesizer.synthesize(query_str, nodes)
        return response_obj

Opetion2(RAGStringQueryEngine)

# Option 2: 文字列を返す (we use a raw LLM call for illustration)

from llama_index.llms.openai import OpenAI
from llama_index.core import PromptTemplate

qa_prompt = PromptTemplate(
    "コンテキスト情報は以下のとおりです.\n"
    "---------------------\n"
    "{context_str}\n"
    "---------------------\n"
    "事前知識ではなく、与えられたコンテキスト情報に基づき、 "
    "クエリに回答してください。\n"
    "クエリ: {query_str}\n"
    "回答: "
)


class RAGStringQueryEngine(CustomQueryEngine):
    """RAG String Query Engine."""

    retriever: BaseRetriever
    response_synthesizer: BaseSynthesizer
    llm: OpenAI
    qa_prompt: PromptTemplate

    def custom_query(self, query_str: str):
        nodes = self.retriever.retrieve(query_str)

        context_str = "\n\n".join([n.node.get_content() for n in nodes])
        response = self.llm.complete(
            qa_prompt.format(context_str=context_str, query_str=query_str)
        )

        return str(response)

やってみよう

サンプルデータを用いてやってみましょう。

Trying Option 1 (RAGQueryEngine)

synthesizer = get_response_synthesizer(response_mode="compact")
query_engine = RAGQueryEngine(
    retriever=retriever, response_synthesizer=synthesizer
)

response = query_engine.query("カローラスポーツはどんな車ですか？")

次のような応答が返ってきました。

カローラスポーツは、フロントロアグリルやフロントフォグモール、リヤバンパーロアモールなどの外装装備が特徴的な車です。さらに、異なるタイヤサイズやホイールサイズ、安全装備などが標準装備されています。また、パーキングサポートブレーキやドライブレコーダーなどの機能も備えています。

Option 2 (RAGStringQueryEngine)をやってみる

llm = OpenAI(model="gpt-3.5-turbo")

query_engine = RAGStringQueryEngine(
    retriever=retriever,
    response_synthesizer=synthesizer,
    llm=llm,
    qa_prompt=qa_prompt,
)

response = query_engine.query("カローラスポーツはどのような車ですか？日本語で教えてください")

print(str(response))

responseオブジェクトを出力すると、テキストのみ入っていることがわかります。

Response(
response='カローラスポーツは、フロントロアグリル（ブラック）＆フロントフォグモール（ブラック）やリヤバンパーロア（艶有り・ブラック）などの外装装備が特徴的な車です。また、タイヤやホイールのサイズも異なるグレードがあります。価格やオプション装備については、販売店にお問い合わせください。', 
source_nodes=[], 
metadata=None
)

---

こちらを試したコードはこちらです。

Google Colab

以上となります。

まとめ

• CustomQueryEngineクラスを用いることで、自作のクエリエンジンが作成できる

• 自作クエリエンジンでは、検索器、生成器（LLM)、応答合成器、クエリなどをカスタマイズすることができる

次回は、これをもとにハイブリッド検索器を作成したいと思います。

それでは、ご覧いただきありがとうございました。