LlamaIndexを完全に理解するチュートリアル その1:処理の概念や流れを理解する基礎編(v0.7.9対応)

LlamaIndexブラックボックス化していませんか?きちんと理解してカスタマイズの幅を広げましょう!
2023.07.17

こんちには。

データアナリティクス事業本部 インテグレーション部 機械学習チームの中村です。

・LlamaIndexを完全に理解するチュートリアル その1を最新版で書き直しました。
・差分が大きいため、旧記事は残しつつ、新しい記事として投稿します。

本記事から数回に分けて、LlamaIndexの中身を深堀していく記事を書いていこうと思います。

本記事の内容

初回は基礎ということで、サンプルコードを例にLlamaIndexで登場する処理を深堀して、カスタマイズのヒントになる知識を得ていきます。

基礎編とはなっていますが、結構長めの記事となっています。

基礎が最も大事ということで、LlamaIndexをただ単に使用するだけではブラックボックス化してしまいがちな概念や処理の流れを一通り洗い出しています。

これ以降の記事では、この基礎をベースにしてどのようなユースケースとカスタマイズが考えられるのかを記事にしていこうと思いますが、初回から重たいので完全理解するよりは、今後の拠り所とする感じにして頂けますと幸いです。

LlamaIndexとは

LlamaIndexは大規模言語モデル(LLM)と外部データ(あなた自身のデータ)を接続するためのインターフェースを提供するプロジェクトとなっています。

LLMのカスタマイズするためのパラダイムには主に2種類があります。

  • LLMをFine-tuningする
  • 入力プロンプトにコンテキストを埋め込む

LlamdaIndexは後者を、より性能よく、効率よく、安価に行うために様々なデータ取り込みやインデックス化を実施することが可能です。

公式ドキュメントは以下となっています。(日々更新されているので、現在の最新のv0.7.9を参照しておきます)

環境準備

実行環境

Google Colaboratoryを使います。ハードウェアアクセラレータは無し、ラインタイム仕様も標準です。

ローカル環境で実行されても良いと思います。

Pythonのバージョンは以下です。

!python --version
Python 3.10.11

また、llama-indexは以下のようにインストールします。

!pip install llama-index

インストールされた主要なライブラリは以下です。

!pip freeze | grep -e "openai" -e "llama-index" -e "langchain"
langchain==0.0.234
llama-index==0.7.9
openai==0.27.8

このやり方にこだわる必要はないのですが、環境変数をpython-dotenvを使って読み込みます。

!pip install python-dotenv

OPEN_AI_KEYを.envに書き込みます。

!echo 'OPENAI_API_KEY="あなたのOPENAIのAPIキー"' >.env

この前提として、OPENAI_API_KEYは以下などに従って準備する必要があります。(API利用する分で従量課金となりますのでご注意ください)

環境変数をロードします。

from dotenv import load_dotenv
load_dotenv()

データの準備

適当なテキストデータを準備しておきます。今回は3,000文字程度のテキスト2つにしておきました。

./data/
  sample_001.txt
  sample_002.txt

動かしてみる

ListIndexを使ったサンプル

最も基本的なサンプルとして、ListIndexを使ったコードを準備しました。

from llama_index import SimpleDirectoryReader
from llama_index import ListIndex

documents = SimpleDirectoryReader(input_dir="./data").load_data()

list_index = ListIndex.from_documents(documents)

query_engine = list_index.as_query_engine()

response = query_engine.query("機械学習に関するアップデートについて300字前後で要約してください。")

for i in response.response.split("。"):
    print(i + "。")

おおむね以下の流れで処理されます。

  • Readerの処理
    • データソース(この場合はローカルフォルダ)をList[Document]に変換
  • Indexの作成
    • List[Document]からインデックス(この場合ListIndex)を作成
  • QueryEngineの作成
    • インデックスのas_query_engineでQueryEngineを作成
    • 0.6以降はAPIが新しくなったため、as_query_engineを使用する必要があります

Readerについて

Readerはデータソースに応じて多数のものがLlamaIndexライブラリ内に準備されています。

また、LlamaHubという形では更に多くのが提供されているようです。

サンプルでは、以下のようにローカルのテキストファイルを読み込むSimpleDirectoryReaderを使用しています。

documents = SimpleDirectoryReader(input_dir="./data").load_data()

Indexについて

Indexの仕組みは既に多くの記事で紹介がありますが、公式には以下に記載があります。

このページでは、4つのインデックス構造が紹介されています。

  • List Index(ListIndex)
    • 単にNodeのリストを保持
    • クエリ時は先頭から順次処理し、それぞれの出力を合成
  • Vector Store Index(VectorStoreIndex)
    • 各Nodeに対応する埋め込みベクトルと共に順序付けせずに保持
    • 埋め込みベクトルを使用してNodeを抽出し、それぞれの出力を合成
  • Tree Index(TreeIndex)
    • ノードをツリー構造にして保持
    • クエリ時はRootから探索して、使用するノードを決め、その出力を合成
  • Keyword Table Index(KeywordTableIndex、RAKEKeywordTableIndex、SimpleKeywordTableIndex)
    • 各Nodeからキーワードを抽出し、キーワードに対するNodeをマッピングして保持
    • クエリ時はクエリのキーワードを使ってNodeを選択し、それぞれのノードの出力を合成
    • 3種類のIndexがあり、それぞれノードのキーワード抽出方法が異なる(Retrieverも同様な3種類があるが、独立に設定可能)

これ以外にもIndexは、Knowledge Graph IndexやSQL Indexなど色々とありますが、ここでは省略します。

・元々GPTListIndexのように、GPTがprefixで付いていましたがv0.7.9時点ではそれらが削除されています

Retrieverについて

Indexはその種類によってはRetrieverModeを選択することで、Nodeの選択方法を変えることができます。

Index種類 Retriever Mode 説明
List Index ListRetrieverMode.DEFAULT すべてのノードを抽出
List Index ListRetrieverMode.EMBEDDING 埋め込みベクトルを使って抽出
List Index ListRetrieverMode.LLM LLMを使ってノードを選択
Vector Store Index 一意 埋め込みベクトルを使って抽出
Tree Index TreeRetrieverMode.SELECT_LEAF プロンプトを使ってLeafノードを探索して抽出
Tree Index TreeRetrieverMode.SELECT_LEAF_EMBEDDING 埋め込みベクトルを使ってLeafノードを探索して抽出
Tree Index TreeRetrieverMode.ALL_LEAF 全てのLeafノードを使いクエリ固有のツリーを構築して応答
Tree Index TreeRetrieverMode.ROOT ルートノードのみを使って応答
Table Index KeywordTableRetrieverMode.DEFAULT GPTを使ってクエリのキーワード抽出を行う
Table Index KeywordTableRetrieverMode.SIMPLE 正規表現を使ってクエリのキーワード抽出を行う
Table Index KeywordTableRetrieverMode.RAKE RAKEキーワード抽出器を使ってクエリのキーワード抽出を行う

なおRetrieverModeは、後述するas_query_engineでQueryEngineを作成する際に与えることができます。

・v0.7.9時点ではListRetrieverModeにListRetrieverMode.LLMが追加されています

Contextについて

IndexとRetrieverは密接に関連しているものですが、それとは別に依存する処理クラスをContextとして与えます。

このContextは具体的には、Storage ContextとService Contextの2種類です。

冒頭のサンプルでは、デフォルトで動作しているためContextが見えないのですが、Contextを明示的に書くと以下のようになります。

from llama_index import StorageContext
from llama_index.storage.docstore import SimpleDocumentStore
from llama_index.storage.index_store import SimpleIndexStore
from llama_index.vector_stores import SimpleVectorStore
from llama_index import ServiceContext
from llama_index.node_parser import SimpleNodeParser
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index import LLMPredictor
from llama_index.indices.prompt_helper import PromptHelper
from llama_index.logger.base import LlamaLogger
from llama_index.callbacks.base import CallbackManager

# Storage Contextの作成
storage_context = StorageContext.from_defaults(
    docstore=SimpleDocumentStore()
    , vector_store=SimpleVectorStore()
    , index_store=SimpleIndexStore()
)

# Service Contextの作成
llm_predictor = LLMPredictor()
service_context = ServiceContext.from_defaults(
    node_parser=SimpleNodeParser()
    , embed_model=OpenAIEmbedding()
    , llm_predictor=llm_predictor
    , prompt_helper=PromptHelper.from_llm_metadata(llm_metadata=llm_predictor.metadata)
    , llama_logger=LlamaLogger()
    , callback_manager=CallbackManager([])
)

# Index作成時にContextを入れる
list_index = ListIndex.from_documents(
    documents
    , storage_context=storage_context
    , service_context=service_context
)

# これ以降は同じ
query_engine = list_index.as_query_engine()

response = query_engine.query("機械学習に関するアップデートについて300字前後で要約してください。")

for i in response.response.split("。"):
    print(i + "。")

・v0.7.9時点では`PromptHelper.from_llm_predictor`は、`PromptHelper.from_llm_metadata`に変更されています

このようにIndexは、StorageContextとServiceContextという処理クラスに依存しています。

Indexとは独立してこれらのContextをカスタマイズできます。

Storage Contextについて

Storage Contextは3つのストアで構成されます。

  • Vector Store
  • Document Store
  • Index Store

Storage Context全体は、以下のようにJSONファイルにダンプすることが可能です。

import json

with open("storage_context.json", "wt") as f:
    json.dump(list_index.storage_context.to_dict(), f, indent=4)

構造は以下のようになっています。

{
    "vector_store": {
        // ...
    },
    "doc_store": {
        // ...
    },
    "index_store": {
        // ...
    },
    "graph_store": {
        // ...
    }
}

・v0.7.9時点ではgraph_storeもStorage Contextに追加されています

Vector Storeについて

Vector Storeについて見ていきましょう。

Vector Storeはベクトルデータが格納されるストアで、以下でJSONにダンプできます。

with open("vector_store.json", "wt") as f:
    json.dump(list_index.storage_context.vector_store.to_dict(), f, indent=4)

ListIndexはデフォルトではvector_storeを使わないため、空欄となっています。

{
    "embedding_dict": {},
    "text_id_to_doc_id": {}
}

今回のサンプルのようなデフォルトの場合、Vector StoreはSimpleVectoreStoreとなっています。

list_index.storage_context.vector_store
<llama_index.vector_stores.simple.SimpleVectorStore at 0x1fdbee37b80>

このSimpleVectorStoreは、InMemoryなストアとなっています。

Vector Storeは、その他多数のサードパーティーのストアに対応しており、以下で一覧の確認が可能です。

Document Storeについて

Document Storeについても見ていきましょう。

Document Storeはテキストデータが格納されるストアで、以下でJSONにダンプできます。

with open("docstore.json", "wt") as f:
    json.dump(list_index.storage_context.docstore.to_dict(), f, indent=4)
{
    "docstore/metadata": {
        // ...
    },
    "docstore/data": {
        // ...
    }
}

Document Storeはテキストデータが格納されています。docstore/metadataとdocstore/dataから構成されます。

docstore/metadataには、doc_idとそれぞれのdoc_idに対するdoc_hashが含まれています。

doc_idは以下の単位で割り当てられています。

  • ソースとなっているDocumentオブジェクトそれぞれ(今回の場合、2つのテキストファイル)
  • 上記を分割したノードそれぞれ(こちらには元となる文書の情報がref_doc_idで与えられる)
{
    "docstore/metadata": {
        // sample_001.txt
        "0adc4efe-8a04-46f1-a449-06980c947ae2": {
            "doc_hash": "ba7b36c4988a208804a431094c0bf0bcf95065812172378e241eb262052bfed1"
        },
        // sample_002.txt
        "c59647ac-5d2e-4fbd-8bc8-b56445d71634": {
            "doc_hash": "4e983b11b700f9a9140c4564d4b3a5143ec07ccade7c913ad7ce519ed0eee61f"
        },
        // sample_001.txtのノード1
        "58cab56b-5c6c-4bc7-b7da-58b2c9aebfce": {
            "doc_hash": "9652f6608521a9243eb3c16e796a2946ee136a5724e572e6a724da6bc4a2fd35",
            "ref_doc_id": "0adc4efe-8a04-46f1-a449-06980c947ae2"
        },
        // sample_001.txtのノード2
        "7fa08709-8e79-4a07-8fd4-c190596e2801": {
            "doc_hash": "627240e5f3ac51d9bd6c529773a598cc792d1c3434402340f8b857afa71d5a42",
            "ref_doc_id": "0adc4efe-8a04-46f1-a449-06980c947ae2"
        },
        // ... 以降略 ...
    }
}

・v0.7.9時点ではref_doc_idが追加され、参照関係が分かりやすくなっています

docstore/dataは各ノードとなっており、テキストの中身の情報やその他の情報が含まれます。

{
    "docstore/data": {

        // ... 中略 ...
        "7fa08709-8e79-4a07-8fd4-c190596e2801": {
            "__data__": {
                "id_": "7fa08709-8e79-4a07-8fd4-c190596e2801",
                "embedding": null,
                "metadata": {},
                "excluded_embed_metadata_keys": [],
                "excluded_llm_metadata_keys": [],
                "relationships": {
                    "1": {
                        "node_id": "0adc4efe-8a04-46f1-a449-06980c947ae2",
                        "node_type": null,
                        "metadata": {},
                        "hash": "ba7b36c4988a208804a431094c0bf0bcf95065812172378e241eb262052bfed1"
                    },
                    "2": {
                        "node_id": "58cab56b-5c6c-4bc7-b7da-58b2c9aebfce",
                        "node_type": null,
                        "metadata": {},
                        "hash": "9652f6608521a9243eb3c16e796a2946ee136a5724e572e6a724da6bc4a2fd35"
                    },
                    "3": {
                        "node_id": "a528dd85-f616-4855-9926-1034c946825c",
                        "node_type": null,
                        "metadata": {},
                        "hash": "a76edcd17cbbdc1332c2f4c1db97645664d23689cedb7855b4a7ba966aa196af"
                    }
                },
                "hash": "627240e5f3ac51d9bd6c529773a598cc792d1c3434402340f8b857afa71d5a42",
                "text": "<本文>",
                "start_char_idx": 904,
                "end_char_idx": 1803,
                "text_template": "{metadata_str}\n\n{content}",
                "metadata_template": "{key}: {value}",
                "metadata_seperator": "\n"
            },
            "__type__": "1"
        },

        // ... 中略 ...
    }
}

・v0.7.9時点ではノードオブジェクトのプロパティが大きく変わっています

metadataには様々なkey-valueを含むことができ、元ファイル名などの情報も格納できます。

start_char_idxとend_char_idxで、元のファイルのどこからどこまでを分割したものか確認することができます。

relationshipsは関連のあるノードとの関係性を表しており、元ファイルとなっている文書のdoc_idや、テキストの分割前の前後関係にあるノードのdoc_idなどの情報が格納されています。

以下のNodeRelationshipオブジェクトに基づいています。

class NodeRelationship(str, Enum):
    """Node relationships used in `BaseNode` class.

    Attributes:
        SOURCE: The node is the source document.
        PREVIOUS: The node is the previous node in the document.
        NEXT: The node is the next node in the document.
        PARENT: The node is the parent node in the document.
        CHILD: The node is a child node in the document.

    """

    SOURCE = auto()
    PREVIOUS = auto()
    NEXT = auto()
    PARENT = auto()
    CHILD = auto()

typeはノードのタイプを表しており、通常はTEXTですが、以下のObjectTypeオブジェクトに基づいています。

class ObjectType(str, Enum):
    TEXT = auto()
    IMAGE = auto()
    INDEX = auto()
    DOCUMENT = auto()

今回のサンプルのようなデフォルトの場合、Document StoreはSimpleDocumentStoreとなっています。

list_index.storage_context.docstore
<br />

このSimpleDocumentStoreは、InMemoryなストアとなっています。

Document Storeは、その他様々なストアに対応しており、以下で一覧の確認が可能です。

<iframe class="hatenablogcard" style="width:100%" frameborder="0" scrolling="no" src="https://hatenablog-parts.com/embed?url=https://gpt-index.readthedocs.io/en/v0.7.9/core_modules/data_modules/storage/docstores.html"></iframe>

Index Store

Index Storeについても見ていきましょう。

Index Storeはインデックスに関する情報が格納されるストアで、以下で構造を確認することができます。

list_index.storage_context.index_store.get_index_struct().to_dict()
{'index_id': '7ecde63e-a514-4612-9c77-22a81324f0f7',
 'summary': None,
 'nodes': ['58cab56b-5c6c-4bc7-b7da-58b2c9aebfce',
  '7fa08709-8e79-4a07-8fd4-c190596e2801',
  'a528dd85-f616-4855-9926-1034c946825c',
  '3b12a8f8-7136-476b-b134-d05e510e7c7b',
  'b8f01272-4379-4c80-af1d-14d459234a71',
  '0c38c52a-6fed-4090-9f1e-5cc3eba71a0f',
  'ebc9a7c0-f254-4e7f-88ba-d035030dce4b',
  '7cd71da3-677e-4091-9d7c-053c637e850d']}

ListIndexなので、各ノードがリストで格納されていることが分かります。

今回のサンプルのようなデフォルトの場合、Index StoreはSimpleIndexStoreとなっています。

list_index.storage_context.index_store
<br />

このSimpleIndexStoreは、InMemoryなストアとなっています。

Index Storeは、その他様々なストアに対応しており、以下で一覧の確認が可能です。

<iframe class="hatenablogcard" style="width:100%" frameborder="0" scrolling="no" src="https://hatenablog-parts.com/embed?url=https://gpt-index.readthedocs.io/en/v0.7.9/core_modules/data_modules/storage/index_stores.html"></iframe>

Document Storeの詳細

ここで、Document Storeの各ノードの情報を見ていきます。

以下のコードで、各ノードの長さや元テキストファイルに対する始点、終点を見ることができます。

for doc_id, node in list_index.storage_context.docstore.docs.items():
    node_dict = node.__dict__
    print(f'{doc_id=}, len={len(node_dict["text"])}, start={node_dict["start_char_idx"]}, end={node_dict["end_char_idx"]}')
doc_id='58cab56b-5c6c-4bc7-b7da-58b2c9aebfce', len=903, start=0, end=903
doc_id='7fa08709-8e79-4a07-8fd4-c190596e2801', len=899, start=904, end=1803
doc_id='a528dd85-f616-4855-9926-1034c946825c', len=942, start=1798, end=2740
doc_id='3b12a8f8-7136-476b-b134-d05e510e7c7b', len=245, start=2747, end=2992
doc_id='b8f01272-4379-4c80-af1d-14d459234a71', len=735, start=0, end=735
doc_id='0c38c52a-6fed-4090-9f1e-5cc3eba71a0f', len=795, start=736, end=1531
doc_id='ebc9a7c0-f254-4e7f-88ba-d035030dce4b', len=832, start=1532, end=2364
doc_id='7cd71da3-677e-4091-9d7c-053c637e850d', len=616, start=2365, end=2981

<div style="background-color: #FBE1CC;"> <p stype="margin: 16px 0;"> ・ノードのプロパティが大きく変わっているため、アクセス方法も変わっています(非公開APIなので致し方ない) </p> </div>

ある程度のサイズで分割されているものの、一定ではなくオーバーラップもしている形のようです。

これらの分割の挙動をカスタマイズするには、以降で述べるService ContextのNodeParser、より具体的にはNodeParserに与えるTextSplitterをカスタマイズする必要があります。

Service Contextについて

次に、Service Contextについて見ていきましょう。

Service Contextは以下のページに引数として与えられるクラスから、その構成を見ることができます。

<iframe class="hatenablogcard" style="width:100%" frameborder="0" scrolling="no" src="https://hatenablog-parts.com/embed?url=https://gpt-index.readthedocs.io/en/v0.7.9/api_reference/service_context.html"></iframe>

これによればService Contextには以下が含まれていそうです。

  • NodeParser : テキストをチャンクに分割してノードを作成する
  • Embeddings : テキストを埋め込みベクトルに変換する
  • LLMPredictor : テキスト応答(Completion)を得るための言語モデル(LLM)処理クラス
  • PromptHelper : LLM側のトークン数制限に合うようテキストを分割する
  • CallbackManager : 様々な処理のstart, endでコールバックを設定
  • LlamaLogger : LLMへのクエリのログを取得するのに使用

なお以下のように属性一覧を確認すると、具体的なオブジェクトが確認できます。

list_index.service_context.__dict__
{'llm_predictor': ,
 'prompt_helper': ,
 'embed_model': ,
 'node_parser': ,
 'llama_logger': ,
 'callback_manager': }

NodeParser

NodeParserは、テキストをチャンクに分割してノードを作成する部分を担っています。

以下により、今回のサンプルのようなデフォルトの場合は、SimpleNodeParserが設定されていることが分かります。

list_index.service_context.node_parser
<br />

以下を確認すると、現在NodeParserはこのSimpleNodeParser一種類のようです。

<iframe class="hatenablogcard" style="width:100%" frameborder="0" scrolling="no" src="https://hatenablog-parts.com/embed?url=https://gpt-index.readthedocs.io/en/v0.6.8/reference/node_parser.html#llama_index.node_parser.NodeParser"></iframe>

SimpleNodeParserの属性は以下で確認ができます。

list_index.service_context.node_parser.__dict__
{'callback_manager': ,
 '_text_splitter': ,
 '_include_metadata': True,
 '_include_prev_next_rel': True,
 '_metadata_extractor': None}

具体的な分割処理は、TextSplitterが担っています。

TextSplitterには以下のような属性が含まれており、カスタマイズのヒントになりそうです。

list_index.service_context.node_parser._text_splitter.__dict__
{'_separator': ' ',
 '_chunk_size': 1024,
 '_chunk_overlap': 20,
 'tokenizer': <bound method Encoding.encode of >,
 '_backup_separators': ['\n'],
 'callback_manager': }

<div style="background-color: #FBE1CC;"> <p stype="margin: 16px 0;"> ・オーバーラップのデフォルト値が20に変更されています </p> </div>

Embeddings

Embeddingsは、テキストを埋め込みベクトルに変換する部分を担っています。

以下により、今回のサンプルのようなデフォルトの場合は、OpenAIEmbeddingが設定されていることが分かります。

list_index.service_context.embed_model
<br />

以下を確認すると、EmbeddingsにはOpenAIEmbeddingの他にもLangchainEmbeddingなどが使えそうです。

<iframe class="hatenablogcard" style="width:100%" frameborder="0" scrolling="no" src="https://hatenablog-parts.com/embed?url=https://gpt-index.readthedocs.io/en/v0.7.9/core_modules/model_modules/embeddings/modules.html"></iframe>

LangchainEmbeddingを使うことで、Hugging Faceのモデルなどより広範なモデルに対応することができるようです。

<iframe class="hatenablogcard" style="width:100%" frameborder="0" scrolling="no" src="https://hatenablog-parts.com/embed?url=https://gpt-index.readthedocs.io/en/v0.7.9/core_modules/model_modules/embeddings/usage_pattern.html#embedding-model-integrations"></iframe>

OpenAIEmbeddingの属性は以下で確認ができます。

list_index.service_context.embed_model.__dict__
{'_total_tokens_used': 0,
 '_last_token_usage': None,
 '_tokenizer': <bound method Encoding.encode of >,
 'callback_manager': ,
 '_text_queue': [],
 '_embed_batch_size': 10,
 'deployment_name': None,
 'query_engine': ,
 'text_engine': ,
 'openai_kwargs': {}}

このことから、デフォルトではtext-embedding-ada-002が使われることが分かります。

今回のサンプルのListIndexでは、埋め込みベクトルは使われないのですが、デフォルトの設定はこのようになっていることを確認できました。

LLMPredictor

LLMPredictorはテキスト応答(Completion)を得るための言語モデルの部分を担っています。

以下により、LLMPredictorというクラスが設定されていることが分かります。

list_index.service_context.llm_predictor
<llama_index.llm_predictor.base.LLMPredictor at 0x7f076a4e1db0>

LlamaIndexのLLMクラスは、OpenAI、Hugging Face、LangChainのどれであっても、LLMモジュールを定義するための統一されたインターフェイスを提供しています。

・元々は、LangChainのLLMChainクラスのラッパーでしたが独立したモジュールとして整備されました

以下にそのモジュールの一覧があり、OpenAI以外にもAnthropicやHugging Face、PaLMなどを使用することも可能です。

LLMPredictorの属性は以下で確認ができます。

list_index.service_context.llm_predictor.__dict__
{'_llm': OpenAI(model='text-davinci-003', temperature=0.0, max_tokens=None, additional_kwargs={}, max_retries=10),
 'callback_manager': <llama_index.callbacks.base.CallbackManager at 0x1fdbee5b040>}

LLMとしては、OpenAIクラスが使用されており、こちらが処理の本体を担っています。

type(list_index.service_context.llm_predictor.llm)
llama_index.llms.openai.OpenAI

OpenAIクラスの属性は以下で確認ができます。

list_index.service_context.llm_predictor.llm.__dict__
{'model': 'text-davinci-003',
 'temperature': 0.0,
 'max_tokens': None,
 'additional_kwargs': {},
 'max_retries': 10}

これらを確認すると、text-davinci-003という言語モデルがデフォルトで使われることが分かります。

PromptHelper

PromptHelperは、トークン数制限を念頭において、テキストを分割するなどの部分を担っています。

NodeParserと同じようなイメージですが、こちらはLLM側のトークン数制限に合うようにするというのが主な用途となります。

以下により、PromptHelperが設定されていることが分かります。

list_index.service_context.prompt_helper
<llama_index.indices.prompt_helper.PromptHelper at 0x1fdbee5b160>

以下にその説明があります。

PromptHelperの属性は以下で確認ができます。

list_index.service_context.prompt_helper.__dict__
{'context_window': 4097,
 'num_output': 256,
 'chunk_overlap_ratio': 0.1,
 'chunk_size_limit': None,
 '_tokenizer': <bound method Encoding.encode of <Encoding 'gpt2'>>,
 '_separator': ' '}

・プロパティ名や設定方法がratioになるなどその方法がかわっています

CallbackManager

LlamaIndexの様々な処理のstart, endでコールバックを設定することができます。

CallbackManagerにCallbackHandlerを設定することで、各CallbackHandlerのon_event_start, on_event_endが発火します。

on_event_startとon_event_endには、それぞれ処理のタイムであるCBEventTypeとpayloadが与えられます。

CBEventTypeの一覧は以下のようになっています。

  • CBEventType.CHUNKING : テキスト分割処理の前後
  • CBEventType.NODE_PARSING : NodeParserの前後
  • CBEventType.EMBEDDING : 埋め込みベクトル作成処理の前後
  • CBEventType.LLM : LLM呼び出しの前後
  • CBEventType.QUERY : クエリの開始と終了
  • CBEventType.RETRIEVE : ノード抽出の前後
  • CBEventType.SYNTHESIZE : レスポンス合成の前後
  • CBEventType.TREE : サマリー処理の前後

準備されているCallbackHandlerは、LlamaDebugHandlerのみとなっています。

CallbackManagerの属性は以下で確認ができます。

list_index.service_context.callback_manager.__dict__
{'handlers': [],
 '_trace_map': defaultdict(list,
             {'root': ['e1cf000e-c0e0-49f1-9e9d-a82b345457da'],
              'e1cf000e-c0e0-49f1-9e9d-a82b345457da': ['8b3ecbb6-997c-47db-8fa6-dfd79a2105aa',
               '05002f75-32c4-4277-a7b7-9ef4ffdb7ba6'],
              '05002f75-32c4-4277-a7b7-9ef4ffdb7ba6': ['720e67ef-11a8-454a-8392-f2047a4daaf4',
               '39e28128-d528-4878-b514-369ffa0b3b34',
               '38601252-b008-43e9-9b0d-9d8bb6f80049']}),
 '_trace_event_stack': ['root'],
 '_trace_id_stack': []}

今回のサンプルのようなデフォルトでは、handlersに何も設定されていないようです。

LlamaLogger

LlamaLoggerはあまりドキュメントに記載がないのですが、主にLLMへのクエリのログを取得するのに使用されるようです。

Loggerを設定すると、クエリ実行後にログを取得することができます。

list_index.service_context.llama_logger.get_logs()

Query Engineについて

Query Engineはサンプルのように、インデックスクラスからas_query_engineで作成されるクラスで、Storage ContextやService Context以外の設定をここで実施できます。

サンプルでは以下が該当しています。

query_engine = list_index.as_query_engine()

Indexによりインスタンス化されるQuery Engineは異なり、以下にQuery Engineの一覧があります。

今回Indexとして挙げたList Index, Vector Index, Tree Index, Keyword Table Indexは、ベーシックなRetriever Query EngineがQuery Engineとして生成されます。

以下で、Query Engineのクラスを確認できます。

query_engine
<llama_index.query_engine.retriever_query_engine.RetrieverQueryEngine at 0x7fb0b041b220>

公式ドキュメントであまり明記されていないのですが、as_query_engineで設定できるものの例として以下が挙げられます。

  • Retriever Mode : retriever_mode
    • Indexのところで述べた通り、Retrieverを切り替え可能
  • Node Postprocessor : node_postprocessors
    • Node抽出後の後処理
    • キーワードフィルタや前後のチャンク取得などをここで実現可能
  • Response Synthesizer : response_synthesizer
    • LLMからのレスポンスを合成する処理を行います
  • Response Mode : response_mode
    • レスポンス合成のモード
    • Response Synthesizerが未指定の場合に有効
  • Prompt Templates : text_qa_template, refine_template, simple_template
    • レスポンス合成に必要な各種プロンプトのテンプレート
    • Response Modeにより使用するPrompt Templateは異なる
    • Response Synthesizerが未指定の場合に有効

・Response Synthesizerの直接指定が可能となっています
・Response Modeで指定するAPIも残してあるようですが、今後非推奨となるかもしれません

RetrieverQueryEngineの属性は以下で確認ができます。

query_engine.__dict__
{'_retriever': <llama_index.indices.list.retrievers.ListIndexRetriever at 0x1fdbecc4640>,
 '_response_synthesizer': <llama_index.response_synthesizers.compact_and_refine.CompactAndRefine at 0x1fdbee5b790>,
 '_node_postprocessors': [],
 'callback_manager': <llama_index.callbacks.base.CallbackManager at 0x1fdbee5b040>}

Node Postprocessorはデフォルトでは指定されないことが分かります。

最終的に応答を合成する処理はResponse Synthesizerが担っています。

Response Synthesizerには、ResponseMode.COMPACTに相当する処理が設定されています。

更に、Response Synthesizerの属性は以下です。

query_engine._response_synthesizer.__dict__
{'_service_context': ServiceContext(llm_predictor=<llama_index.llm_predictor.base.LLMPredictor object at 0x000001FDBEE39880>, prompt_helper=<llama_index.indices.prompt_helper.PromptHelper object at 0x000001FDBEE5B160>, embed_model=<llama_index.embeddings.openai.OpenAIEmbedding object at 0x000001FDBEE5B0A0>, node_parser=<llama_index.node_parser.simple.SimpleNodeParser object at 0x000001FDBEE37FD0>, llama_logger=<llama_index.logger.base.LlamaLogger object at 0x000001FDBEE5B1C0>, callback_manager=<llama_index.callbacks.base.CallbackManager object at 0x000001FDBEE5B040>),
 '_callback_manager': <llama_index.callbacks.base.CallbackManager at 0x1fdbee5b040>,
 '_streaming': False,
 '_text_qa_template': <llama_index.prompts.base.Prompt at 0x1fdb2aaba60>,
 '_refine_template': <llama_index.prompts.base.Prompt at 0x1fdb2ae9280>,
 '_verbose': False}

・ResponseBuilderは実装の見直し時に削除されており、登場しなくなっています

またtext_qa_templaterefine_templateはPromptクラスが設定されており、それぞれの詳細を以下で確認できます。

query_engine._response_synthesizer._text_qa_template.__dict__
{'prompt': PromptTemplate(input_variables=['context_str', 'query_str'], output_parser=None, partial_variables={}, template='Context information is below.\n---------------------\n{context_str}\n---------------------\nGiven the context information and not prior knowledge, answer the question: {query_str}\n', template_format='f-string', validate_template=True),
 'prompt_selector': PromptSelector(default_prompt=PromptTemplate(input_variables=['context_str', 'query_str'], output_parser=None, partial_variables={}, template='Context information is below.\n---------------------\n{context_str}\n---------------------\nGiven the context information and not prior knowledge, answer the question: {query_str}\n', template_format='f-string', validate_template=True), conditionals=[]),
 'partial_dict': {},
 'prompt_kwargs': {},
 'prompt_type': <PromptType.QUESTION_ANSWER: 'text_qa'>,
 'output_parser': None,
 '_original_template': 'Context information is below.\n---------------------\n{context_str}\n---------------------\nGiven the context information and not prior knowledge, answer the question: {query_str}\n',
 'metadata': {}}
query_engine._response_synthesizer._refine_template.__dict__
{'prompt_selector': PromptSelector(default_prompt=PromptTemplate(input_variables=['context_msg', 'existing_answer', 'query_str'], output_parser=None, partial_variables={}, template="The original question is as follows: {query_str}\nWe have provided an existing answer: {existing_answer}\nWe have the opportunity to refine the existing answer (only if needed) with some more context below.\n------------\n{context_msg}\n------------\nGiven the new context, refine the original answer to better answer the question. If the context isn't useful, return the original answer.", template_format='f-string', validate_template=True), conditionals=[(<function is_chat_model at 0x000001FDB2A71AF0>, ChatPromptTemplate(input_variables=['context_msg', 'query_str', 'existing_answer'], output_parser=None, partial_variables={}, messages=[HumanMessagePromptTemplate(prompt=PromptTemplate(input_variables=['context_msg', 'existing_answer', 'query_str'], output_parser=None, partial_variables={}, template="We have the opportunity to refine the above answer (only if needed) with some more context below.\n------------\n{context_msg}\n------------\nGiven the new context, refine the original answer to better answer the question: {query_str}. If the context isn't useful, output the original answer again.\nOriginal Answer: {existing_answer}", template_format='f-string', validate_template=True), additional_kwargs={})]))]),
 'prompt': PromptTemplate(input_variables=['context_msg', 'existing_answer', 'query_str'], output_parser=None, partial_variables={}, template="The original question is as follows: {query_str}\nWe have provided an existing answer: {existing_answer}\nWe have the opportunity to refine the existing answer (only if needed) with some more context below.\n------------\n{context_msg}\n------------\nGiven the new context, refine the original answer to better answer the question. If the context isn't useful, return the original answer.", template_format='f-string', validate_template=True),
 'partial_dict': {},
 'prompt_kwargs': {},
 'prompt_type': <PromptType.REFINE: 'refine'>,
 'output_parser': None,
 '_original_template': None,
 'metadata': {}}

Retriever Mode

Indexのところで述べたためここでは割愛します。

Node Postprocessor

Retrieverにより抽出されたノードについて、後処理を行います。

以下に準備されているPostprocessorの一覧があります。

LLMRerankやTimeWeightedPostprocessorなど様々な処理が準備されています。

・Postprocessorの種類はかなり以前より増えています

Optimizer

現在は存在していません。

・Optimizerは、v0.7.9では廃止されています

Response Synthesizer

LLMからのレスポンスを合成する処理を行います。

・以前は内部クラスの扱いでしたが、v0.7.9時点では外部APIとして定義されています

今はデフォルト値が使用されていますが、明示的に作成する場合は以下のようなfactoryを使う処理が準備されています。

from llama_index.response_synthesizers import get_response_synthesizer
from llama_index.response_synthesizers import ResponseMode

response_synthesizer = get_response_synthesizer(response_mode=ResponseMode.COMPACT)

response_modeについては以下に記載があります。

以下に概要を記載しておきます。

Response Mode 概要
ResponseMode.REFINE 各ノードを順次調べて答えを作成しながら洗練させていく
ResponseMode.COMPACT 実際はCompactAndRefineという処理であり、ほぼRefineと同じだが、チャンクをLLMのコンテキスト長を最大限利用するようにrepackするためより高速
ResponseMode.TREE_SUMMARIZE 抽出されたノードを使いサマライズを繰り返すことでチャンクを圧縮した後にクエリを行う。バージョンにより多少動作が変わる
ResponseMode.SIMPLE_SUMMARIZE ノードを単純に結合してひとつのチャンクとしてクエリする。LLMコンテキスト長を超えるとエラーとなる
ResponseMode.GENERATION ノード情報を使わず単にLLMに問い合わせる
ResponseMode.ACCUMULATE 各ノードに対するLLMの結果を計算して結果を連結する
ResponseMode.COMPACT_ACCUMULATE ACCUMULATEとほぼ同じだが、チャンクをLLMのコンテキスト長を最大限利用するようにrepackするためより高速

デフォルトでは、ResponseMode.COMPACTが使用されます。

Prompt Templates

get_response_synthesizerには、Prompt Templateを与えることができ、これによりデフォルトのPrompt Templateをカスタマイズすることが可能です。

以下のように3種類のテンプレートが設定できます。

  • text_qa_template : シンプルに以下のようなコンテキストに対して回答をもとめるようなプロンプト
  • refine_template : 以前の回答を新しいコンテキストでREFINE(ブラッシュアップ)するプロンプト
  • simple_template : 単純にユーザのクエリをLLMにそのまま投げるプロンプト

以下が設定の例です。

from llama_index.response_synthesizers import get_response_synthesizer
from llama_index.response_synthesizers import ResponseMode
from llama_index.prompts.prompts import Prompt
from llama_index.prompts.prompt_type import PromptType
from llama_index.prompts.default_prompts import DEFAULT_TEXT_QA_PROMPT_TMPL, DEFAULT_SIMPLE_INPUT_TMPL, DEFAULT_REFINE_PROMPT_TMPL

response_synthesizer = get_response_synthesizer(response_mode=ResponseMode.COMPACT,
    text_qa_template=Prompt(DEFAULT_TEXT_QA_PROMPT_TMPL, prompt_type=PromptType.QUESTION_ANSWER),
    refine_template=Prompt(DEFAULT_REFINE_PROMPT_TMPL, prompt_type=PromptType.REFINE),
    simple_template=Prompt(DEFAULT_SIMPLE_INPUT_TMPL, prompt_type=PromptType.SIMPLE_INPUT)
    )

デフォルトのResponseMode.COMPACTでは、最初のノードに対してtext_qa_templateで動作し、その後は前の結果を使いながらrefine_templateで動作します。

simple_templateは使用しませんが、上記の例では汎用性のため設定しています)

デフォルトのプロンプトは以下で確認ができます。

この中から今回関連しそうなプロンプトを見ていきます。

text_qa_template

こちらはシンプルに以下のようなコンテキストに対して回答をもとめるようなプロンプトとなっています。

DEFAULT_TEXT_QA_PROMPT_TMPL = (
    "Context information is below.\n"
    "---------------------\n"
    "{context_str}\n"
    "---------------------\n"
    "Given the context information and not prior knowledge, "
    "answer the question: {query_str}\n"
)
DEFAULT_TEXT_QA_PROMPT = Prompt(
    DEFAULT_TEXT_QA_PROMPT_TMPL, prompt_type=PromptType.QUESTION_ANSWER
)

refine_template

こちらは多少複雑になっており、LLMがチャットモデルかどうかで挙動を変えるため、ConditionalPromptSelectorによりラップされています。

DEFAULT_REFINE_PROMPT_SEL_LC = PromptSelector(
    default_prompt=DEFAULT_REFINE_PROMPT.get_langchain_prompt(),
    conditionals=[(is_chat_model, CHAT_REFINE_PROMPT.get_langchain_prompt())],
)
DEFAULT_REFINE_PROMPT_SEL = RefinePrompt(
    langchain_prompt_selector=DEFAULT_REFINE_PROMPT_SEL_LC,
    prompt_type=PromptType.REFINE,
)

デフォルト(is_chat_modelがFalse)の場合、以下が使用されます。

DEFAULT_REFINE_PROMPT_TMPL = (
    "The original question is as follows: {query_str}\n"
    "We have provided an existing answer: {existing_answer}\n"
    "We have the opportunity to refine the existing answer "
    "(only if needed) with some more context below.\n"
    "------------\n"
    "{context_msg}\n"
    "------------\n"
    "Given the new context, refine the original answer to better "
    "answer the question. "
    "If the context isn't useful, return the original answer."
)
DEFAULT_REFINE_PROMPT = Prompt(
    DEFAULT_REFINE_PROMPT_TMPL, prompt_type=PromptType.REFINE
)

チャットモデルの場合は以下が使用されます。

CHAT_REFINE_PROMPT_TMPL_MSGS = [
    HumanMessagePromptTemplate.from_template(
        "We have the opportunity to refine the above answer "
        "(only if needed) with some more context below.\n"
        "------------\n"
        "{context_msg}\n"
        "------------\n"
        "Given the new context, refine the original answer to better "
        "answer the question: {query_str}. "
        "If the context isn't useful, output the original answer again.\n"
        "Original Answer: {existing_answer}"
    ),
]


CHAT_REFINE_PROMPT_LC = ChatPromptTemplate.from_messages(CHAT_REFINE_PROMPT_TMPL_MSGS)
CHAT_REFINE_PROMPT = RefinePrompt.from_langchain_prompt(CHAT_REFINE_PROMPT_LC)

OpenAIのAPIがLLMによって少し異なるため、吸収するための工夫がされていることが分かります。

simple_template

こちらはシンプルにコンテキスト等もなく、クエリをそのまま送るプロンプトとなっています。

DEFAULT_SIMPLE_INPUT_TMPL = "{query_str}"
DEFAULT_SIMPLE_INPUT_PROMPT = Prompt(
    DEFAULT_SIMPLE_INPUT_TMPL, prompt_type=PromptType.SIMPLE_INPUT
)

カスタマイズ用のサンプルコード

Query Engine関連の処理も明示的に書き直したサンプルコードは以下になりました。

(一部チャットモデルに対応するConditionalPromptSelector部分のみ煩雑なため省略しています)

from llama_index import StorageContext
from llama_index.storage.docstore import SimpleDocumentStore
from llama_index.storage.index_store import SimpleIndexStore
from llama_index.vector_stores import SimpleVectorStore
from llama_index import ServiceContext
from llama_index.node_parser import SimpleNodeParser
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index import LLMPredictor
from llama_index.indices.prompt_helper import PromptHelper
from llama_index.logger.base import LlamaLogger
from llama_index.callbacks.base import CallbackManager
from llama_index.indices.list.base import ListRetrieverMode
from llama_index.response_synthesizers import ResponseMode
from llama_index.response_synthesizers import get_response_synthesizer
from llama_index.prompts.prompts import Prompt
from llama_index.prompts.prompt_type import PromptType
from llama_index.prompts.default_prompts import DEFAULT_TEXT_QA_PROMPT_TMPL, DEFAULT_SIMPLE_INPUT_TMPL, DEFAULT_REFINE_PROMPT_TMPL

# Storage Contextの作成
storage_context = StorageContext.from_defaults(
    docstore=SimpleDocumentStore()
    , vector_store=SimpleVectorStore()
    , index_store=SimpleIndexStore()
)

# Service Contextの作成
llm_predictor = LLMPredictor()
service_context = ServiceContext.from_defaults(
    node_parser=SimpleNodeParser()
    , embed_model=OpenAIEmbedding()
    , llm_predictor=llm_predictor
    , prompt_helper=PromptHelper.from_llm_metadata(llm_metadata=llm_predictor.metadata)
    , llama_logger=LlamaLogger()
    , callback_manager=CallbackManager([])
)

# Index作成時にContextを入れる
list_index = ListIndex.from_documents(
    documents
    , storage_context=storage_context
    , service_context=service_context
)

# Response Synthesizeの作成
response_synthesizer = get_response_synthesizer(response_mode=ResponseMode.COMPACT,
    text_qa_template=Prompt(DEFAULT_TEXT_QA_PROMPT_TMPL, prompt_type=PromptType.QUESTION_ANSWER),
    refine_template=Prompt(DEFAULT_REFINE_PROMPT_TMPL, prompt_type=PromptType.REFINE),
    simple_template=Prompt(DEFAULT_SIMPLE_INPUT_TMPL, prompt_type=PromptType.SIMPLE_INPUT)
    )

# Query Engineをインスタンス化
query_engine = list_index.as_query_engine(
    retriever_mode=ListRetrieverMode.DEFAULT,
    node_postprocessors=[],
    response_synthesizer=response_synthesizer
)

response = query_engine.query("機械学習に関するアップデートについて300字前後で要約してください。")

for i in response.response.split("。"):
    print(i + "。")

まとめ

いかがでしたでしょうか。

初回から重たい内容となってしまいましたが、基礎として登場する用語やブラックボックス化してしまいがちな処理を洗い出せたかなと考えています。

これ以降の記事では、この基礎をベースにしてどのようなユースケースとカスタマイズが考えられるのかを記事にしていこうと思います。

本記事をベースにカスタマイズにぜひチャレンジして見てください。(そしてそれをシェアしてくださったらとても喜びます)

本記事が、今後LlamaIndexをお使いになられる方の参考になれば幸いです。