LlamaIndexを完全に理解するチュートリアル その５：TreeIndexを使ってその動作を確認してみる

こんちには。

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

「LlamaIndexを完全に理解するチュートリアル その５」では、TreeIndexを使ってみようと思います。

本記事で使用する用語は以下のその１で説明していますので、そちらも参照ください。

TreeIndexの概要

TreeIndexはここまでで説明してきたListIndexと異なり、一定以上の規模のデータセットに対して親ノードを持つツリー構造のIndexを構成します。

末端のノード(リーフノード)は、ListIndexの各ノードと変わりません。

その親ノードはリーフノードを要約したテキストを保持しています。

親ノードを作成する際の要約処理は、Indexを作成する際にLLMを使って行われます。

そのため、Tree Indexには要約処理をするプロンプトが内包されている点は注意が必要です。

また、その他にもそもそものリーフノード数が一定数以下の場合はツリー構造が構築されないため、その点も注意が必要になります。

以降、これらの詳細を使いながら見ていきます。

環境準備

その１と同様の方法で準備します。

今回は以下のバージョンとなっています。

• Python 3.10.2
• langchain==0.0.264
• llama-index==0.8.2.post1
• openai==0.27.8
• nltk==3.8.1 (どこかのタイミングで必要になったようです)

データの準備

今回は比較的大きめなテキストファイルを./dataにいくつか準備しておきました。

以下でテキスト長を確認してみました。1万～2万文字程度のテキストを5ファイル準備しています。

import pathlib

INPUT_DIR="./data"

for text_file in pathlib.Path(INPUT_DIR).glob("*.txt"):
    with open(text_file, "rt", encoding='utf-8') as f:
        lines = f.readlines()
        text = "\n".join(lines)
        print(f"file_name: {text_file}, str_len: {len(text)}")

file_name: ..\input\plain_text\developers-io-in-las-vegas-2022.txt, str_len: 21736
file_name: ..\input\plain_text\ml-webinar-2023-03.txt, str_len: 27210
file_name: ..\input\plain_text\nn-architecture-1.txt, str_len: 12240
file_name: ..\input\plain_text\nn-architecture-2.txt, str_len: 11496
file_name: ..\input\plain_text\showcase-2023-06.txt, str_len: 12849

TreeIndexの作成

TreeIndexをまずは作成してみます。処理状況がわかりやすいよう、LlamaDebugHandlerをCallbackManagerに設定しておきます。

from llama_index import SimpleDirectoryReader
from llama_index import TreeIndex
from llama_index import ServiceContext
from llama_index.callbacks import CallbackManager, LlamaDebugHandler

documents = SimpleDirectoryReader(
    input_dir=INPUT_DIR,
    file_metadata=lambda x: {"file_name": str(x)}
).load_data()

llama_debug_handler = LlamaDebugHandler()
callback_manager = CallbackManager([llama_debug_handler])
service_context = ServiceContext.from_defaults(callback_manager=callback_manager)

index = TreeIndex.from_documents(documents
    , service_context=service_context)

実行時のログは以下の通りです。

**********
Trace: index_construction
    |_node_parsing ->  0.252133 seconds
      |_chunking ->  0.080947 seconds
      |_chunking ->  0.05978 seconds
      |_chunking ->  0.039663 seconds
      |_chunking ->  0.030426 seconds
      |_chunking ->  0.030645 seconds
    |_tree ->  127.22388 seconds
      |_llm ->  4.957016 seconds
      |_llm ->  5.013251 seconds
      |_llm ->  7.724937 seconds
      |_llm ->  24.916 seconds
      |_llm ->  7.986544 seconds
      |_llm ->  5.645312 seconds
      |_llm ->  4.472437 seconds
      |_llm ->  18.001501 seconds
      |_llm ->  8.514929 seconds
      |_llm ->  7.078803 seconds
      |_llm ->  9.280937 seconds
      |_llm ->  19.016226 seconds
      |_llm ->  4.611049 seconds
    |_tree ->  36.47499 seconds
      |_llm ->  13.918979 seconds
      |_llm ->  22.556011 seconds
**********

インデックス作成時にもLLMが動作していることが分かります。

インデックス構造の確認

ツリー構造の親子関係を明らかにするためにノードの情報を確認してみます。

for k,v in index.docstore.docs.items():
    node = v
    document_id = k

    # 子ノードのdoc_idを取得
    children_doc_ids = index.index_struct.get_children(node)

    # 末端(リーフノード)は飛ばす
    if len(children_doc_ids) == 0:
        continue

    print(f"parent_doc_id: {document_id}")

    # 子ノードの情報を取得
    for node_num, children_doc_id in children_doc_ids.items():

        # ノード情報の取得
        children_node = index.docstore.get_node(children_doc_id)

        # ファイル名のメタ情報を取得
        file_name = children_node.metadata.get("file_name")

        print("\t" + f"child_doc_id: {children_doc_id} (source: [{file_name}])")

出力結果は長いのですが以下となります。

parent_doc_id: 4446c2d6-15c1-4017-9901-0d4c7ed7f12d
    child_doc_id: b3e8783b-0f8f-4cf4-a228-56b7fc5462bf (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
    child_doc_id: d4975234-83fb-43f6-a794-c5d71f828428 (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
    child_doc_id: 4d3a2d30-9807-471e-8817-f82465973d21 (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
    child_doc_id: 3e7bc15f-30f2-4709-8b96-8954b15ac817 (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
    child_doc_id: 9814244a-4a73-4ab8-82ff-1727599fe1d1 (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
    child_doc_id: 371bc341-1edb-44a8-91f9-be61f252b253 (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
    child_doc_id: dd88463f-e50f-45d5-9ae0-0e2812b2371b (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
    child_doc_id: 7d433c90-af6a-4bb4-9a2c-1bae4d834bb1 (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
    child_doc_id: 06fa7308-575d-4fd0-80b0-f3bf28d0e3b4 (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
    child_doc_id: 363af2fd-828d-4f66-8768-520a203a3c67 (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
parent_doc_id: c387422a-a310-4acb-bafe-1b769c1fe17a
    child_doc_id: a14bf333-352b-49c8-bd13-6c8f769ef1d0 (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
    child_doc_id: 7930a513-ac8c-48fc-9a56-19883b090a69 (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
    child_doc_id: a146e13d-ccb0-4c59-9d5f-9d4611f11554 (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
    child_doc_id: 51da6677-f1f3-4252-8b4f-71bccfa162d9 (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
    child_doc_id: c2d59f72-2447-4664-97e6-d506abf5bcc8 (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
    child_doc_id: 232838b1-c6dd-41e7-8791-de1c20a5aef5 (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
    child_doc_id: b8b0952e-f94b-4bb9-956a-3493f645884c (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
    child_doc_id: ce4c60ed-7e43-4de3-aa08-4da8a7a39ceb (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
    child_doc_id: da16e0a9-735b-4044-81a3-a7e4d07c1511 (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
    child_doc_id: 055e013f-9f8d-4823-8d18-91b841971ce5 (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
parent_doc_id: 3ba6c23e-d589-4c0a-850d-0ef2a9ec921b
    child_doc_id: d844a583-df3a-484f-adc3-301019f19ca7 (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
    child_doc_id: 14b50626-db88-44d2-8de7-a0f483b8edce (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
    child_doc_id: 1f92607d-ed22-44f4-a848-c4a442fc9649 (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
    child_doc_id: 250c57c9-27de-48e7-9057-2f55d3c2b497 (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
    child_doc_id: 4c7a7ef8-73ff-4867-ab61-b341f4570eaa (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
    child_doc_id: 0dd0337a-05c0-458d-ada8-9199650ae36e (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
    child_doc_id: 67038485-0b9b-4c3a-9b93-fd7000fbdf07 (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
    child_doc_id: 5fb91ddc-7c09-4f99-9006-84ac18810b4b (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
    child_doc_id: 9b8c9e18-300d-4d58-87a6-86cab161db5a (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
    child_doc_id: 706c7721-31d5-4ec6-aab2-56c430546a1b (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
parent_doc_id: 3a54e716-1bfb-40a2-a2a1-5a0523bb328d
    child_doc_id: b7ec8024-0d2f-4117-88c7-703c5f4fb8bb (source: [..\input\plain_text\developers-io-in-las-vegas-2022.txt])
    child_doc_id: baf19ab4-ec11-46c5-bbf0-eed6dc812cba (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: 3407b599-696d-4de2-af60-0b46afb94c95 (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: b4ad8a70-9b90-41ca-8746-90806c68a93e (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: f35da2d0-9b39-4fb1-b83c-6cd0b2dff1de (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: bea88963-2a3a-4a04-b8ae-03d9655bf31f (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: dc53d9a9-fcb7-40e8-9077-334e96eb0951 (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: 31159a4f-b3c5-4e9c-a49a-2919330ca744 (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: e33c86ab-3995-4340-92ea-db64e961b986 (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: 4ae8cd1b-b894-4dbd-8d2a-2b1487e707e4 (source: [..\input\plain_text\ml-webinar-2023-03.txt])
parent_doc_id: 2351ea4e-a89a-4c6d-8c38-6e5464dcdbc6
    child_doc_id: 58a051f7-8756-491b-8806-82dd700ac891 (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: da85f344-e88f-4b42-adbc-255889747cfb (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: 7a884910-06a3-4f1a-a125-8b222580a266 (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: 564719dc-92fc-483d-9136-09fd45ad1828 (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: 8509d3da-ae47-476f-a7c5-179de58e0ef7 (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: 419ae0cd-ff1c-4129-878f-5274ed93742b (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: bec49066-a3ca-4b5e-9101-11c52550cbec (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: c2a82c6d-44d6-4ca7-8143-2c5eb4dab05c (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: ced05e47-c708-4700-8732-1193b5f6d31d (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: 32a62f29-187b-4931-8033-a7ece0b5e045 (source: [..\input\plain_text\ml-webinar-2023-03.txt])
parent_doc_id: a08b4cce-ef99-4105-9af8-f30abcdc7267
    child_doc_id: d47d32d8-27c9-4352-b7ab-a0bbd81e0e2b (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: 37ab416e-c181-409d-919f-4fe19b1b9c0b (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: f8d579b9-cc2f-4bf3-bc7d-8db3acfe6a72 (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: c4a497cb-6838-4fb4-a9fb-3d08471e3262 (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: c80a13cc-34e6-4383-aa1b-ffb7a13627eb (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: 83a4364c-bbad-4908-b7ba-52281c6f1dd2 (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: d49dfa1c-d274-4f73-9283-517cfbaaa95f (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: 8d017f63-fb44-4b1b-a579-df00afb37dca (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: 40cca122-9a7a-4cc4-9c79-68641d1903e4 (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: fe6d4408-67e7-4cf1-8b53-fc7d7ddfd163 (source: [..\input\plain_text\ml-webinar-2023-03.txt])
parent_doc_id: 87445710-53d2-47bc-9909-ee9d8541e3f8
    child_doc_id: c74f35a3-0b4d-48d3-beb7-04a5c3f0174a (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: d1d37e85-bf58-410d-b3a1-63b8bec53f07 (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: d4670cb9-7cf8-407e-aa3e-d602428f0af6 (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: 6b2a9cb0-507a-40fd-bc3c-933a58a8e3ed (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: 83e75b7c-bda0-4d1d-8f79-398b84018963 (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: bdbbbb4f-d71b-4b1b-933d-058aa77f504d (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: dc98cd54-0f5e-4e8b-a0b3-db005eb87764 (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: fbd55f79-0937-4aa3-86d3-dc9cb06a450d (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: 78ee2162-fa59-4e0b-b506-e8539f8b8fbf (source: [..\input\plain_text\ml-webinar-2023-03.txt])
    child_doc_id: 482dda09-ef3a-4c32-bcc9-8897dac9dcad (source: [..\input\plain_text\nn-architecture-1.txt])
parent_doc_id: 7a0920a9-039d-4c54-a46f-047245b61cd7
    child_doc_id: 1d32d6f2-c563-4f1a-8616-80ff8b320fe8 (source: [..\input\plain_text\nn-architecture-1.txt])
    child_doc_id: adbf25d1-e249-4e56-a37b-6ec385b338ff (source: [..\input\plain_text\nn-architecture-1.txt])
    child_doc_id: 48876b16-6042-42eb-b0aa-77666dafaa95 (source: [..\input\plain_text\nn-architecture-1.txt])
    child_doc_id: 998193d2-7a8f-486d-ba05-a2ec367bfe9f (source: [..\input\plain_text\nn-architecture-1.txt])
    child_doc_id: ca75351d-ab2b-49b5-abe9-595db636c085 (source: [..\input\plain_text\nn-architecture-1.txt])
    child_doc_id: 168e5780-3826-4e04-b6c1-0dcb6937429e (source: [..\input\plain_text\nn-architecture-1.txt])
    child_doc_id: 9fa34aff-b2a0-45c7-8a9e-3189d34bf6a9 (source: [..\input\plain_text\nn-architecture-1.txt])
    child_doc_id: d2322a64-f44c-48ca-8f16-46248d852ec0 (source: [..\input\plain_text\nn-architecture-1.txt])
    child_doc_id: e1dd3efe-8dd7-4526-80be-941f3d976304 (source: [..\input\plain_text\nn-architecture-1.txt])
    child_doc_id: 854b2c74-fdab-4ef7-b86f-142bc781be3e (source: [..\input\plain_text\nn-architecture-1.txt])
parent_doc_id: b8f3d2cb-df1d-425a-b506-836ea1e3dcf2
    child_doc_id: 412dc45a-1244-4de5-8522-fb9f69cccfc0 (source: [..\input\plain_text\nn-architecture-1.txt])
    child_doc_id: 93e3944a-6344-47dd-9659-85103a739776 (source: [..\input\plain_text\nn-architecture-1.txt])
    child_doc_id: 99f5fab2-3b09-4b2e-9670-0615bb403ea0 (source: [..\input\plain_text\nn-architecture-1.txt])
    child_doc_id: 9ae94e11-12bd-4a8f-99aa-27d92b0ce56a (source: [..\input\plain_text\nn-architecture-1.txt])
    child_doc_id: 3164bef5-e668-42aa-8474-260c98b8cd12 (source: [..\input\plain_text\nn-architecture-1.txt])
    child_doc_id: e393f31f-2b56-4a17-bb93-05f2a4ad8b34 (source: [..\input\plain_text\nn-architecture-1.txt])
    child_doc_id: f0d34f71-c9e2-462e-b2a2-cba60c5b8d02 (source: [..\input\plain_text\nn-architecture-1.txt])
    child_doc_id: fd0512aa-90de-49c8-9980-3753802aaf99 (source: [..\input\plain_text\nn-architecture-2.txt])
    child_doc_id: 01172d25-beb8-4f6c-a762-a9636fbb3a04 (source: [..\input\plain_text\nn-architecture-2.txt])
    child_doc_id: 271e7cd2-2493-427b-b843-4743cf4840c6 (source: [..\input\plain_text\nn-architecture-2.txt])
parent_doc_id: f7649d66-2945-455b-b008-518ffd47065e
    child_doc_id: 46ca721b-a360-4360-8c5e-5331a8fa1f66 (source: [..\input\plain_text\nn-architecture-2.txt])
    child_doc_id: 74add42b-b53a-4e4f-ac1e-a2f3c0a7c3d8 (source: [..\input\plain_text\nn-architecture-2.txt])
    child_doc_id: a8acc3df-73d6-49fb-917f-7b1778d0f569 (source: [..\input\plain_text\nn-architecture-2.txt])
    child_doc_id: da25f8c9-0cf1-4fb7-8d1a-b987cb1d1ff3 (source: [..\input\plain_text\nn-architecture-2.txt])
    child_doc_id: dbc423f9-f5cf-459c-bfa6-64f61d6b8335 (source: [..\input\plain_text\nn-architecture-2.txt])
    child_doc_id: fa756403-dafa-4855-b90c-83deabda8855 (source: [..\input\plain_text\nn-architecture-2.txt])
    child_doc_id: 3c829469-4817-421d-91d6-3f06ca0ed438 (source: [..\input\plain_text\nn-architecture-2.txt])
    child_doc_id: 6b14e7af-ea11-4e9d-89a9-727fe96e88b1 (source: [..\input\plain_text\nn-architecture-2.txt])
    child_doc_id: 3f2f5075-7341-4cdc-9230-fc8d875ad1cf (source: [..\input\plain_text\nn-architecture-2.txt])
    child_doc_id: f884e754-f672-4fd3-b23d-985247948853 (source: [..\input\plain_text\nn-architecture-2.txt])
parent_doc_id: 8f8bcc04-fe23-44c0-86c3-fac6d33021eb
    child_doc_id: 8128de1f-c3a3-49b0-943f-224d800ad567 (source: [..\input\plain_text\nn-architecture-2.txt])
    child_doc_id: ee188ca1-0b0c-4f99-bc8f-f3732fa40f25 (source: [..\input\plain_text\nn-architecture-2.txt])
    child_doc_id: 06034687-de17-4844-b968-9d45c99090b9 (source: [..\input\plain_text\nn-architecture-2.txt])
    child_doc_id: cd58cf66-bf23-41ad-bc06-7753cb5b2252 (source: [..\input\plain_text\nn-architecture-2.txt])
    child_doc_id: 77a9b040-016e-48e4-888f-45ad49e8c809 (source: [..\input\plain_text\showcase-2023-06.txt])
    child_doc_id: c9b3e7ef-7c29-4899-b2d1-26426d2f4dcb (source: [..\input\plain_text\showcase-2023-06.txt])
    child_doc_id: ac782bf7-bdc6-4da2-ae10-d5fad510c407 (source: [..\input\plain_text\showcase-2023-06.txt])
    child_doc_id: e447525a-7198-4524-a207-8df6d04cc1ab (source: [..\input\plain_text\showcase-2023-06.txt])
    child_doc_id: 90d38265-4eda-4f2b-8787-349ade061dfe (source: [..\input\plain_text\showcase-2023-06.txt])
    child_doc_id: 0fbce75e-4a73-48ba-ad4e-6ce4e45faec4 (source: [..\input\plain_text\showcase-2023-06.txt])
parent_doc_id: 8de31b73-a303-4785-ad14-4fb0300d38db
    child_doc_id: 5f8526bc-26b9-4b80-8814-c603db2d9699 (source: [..\input\plain_text\showcase-2023-06.txt])
    child_doc_id: 35ab52a7-acb0-402e-b214-795035f42718 (source: [..\input\plain_text\showcase-2023-06.txt])
    child_doc_id: 1fe88304-3fa7-420c-a990-5c7cdd4c9424 (source: [..\input\plain_text\showcase-2023-06.txt])
    child_doc_id: e0de756e-346b-4cf0-989b-8adea158fceb (source: [..\input\plain_text\showcase-2023-06.txt])
    child_doc_id: 5d322e50-7be8-4ccb-86f4-75af9ed74038 (source: [..\input\plain_text\showcase-2023-06.txt])
    child_doc_id: a71a74f8-be03-4e18-b3ff-210a1c76fd05 (source: [..\input\plain_text\showcase-2023-06.txt])
    child_doc_id: c5b7cd43-691d-4bb4-9c3f-85f128a890ad (source: [..\input\plain_text\showcase-2023-06.txt])
    child_doc_id: 9ccda7e8-3331-4769-9f55-2749551d9b09 (source: [..\input\plain_text\showcase-2023-06.txt])
    child_doc_id: 0cd54c3a-326a-47dd-9efc-2f2fb5ca99c1 (source: [..\input\plain_text\showcase-2023-06.txt])
    child_doc_id: 82f1ff89-1e10-4be8-8e08-ef6a767f23c4 (source: [..\input\plain_text\showcase-2023-06.txt])
parent_doc_id: d2ad14b3-163c-4f87-815f-4a0f3a0ff20a
    child_doc_id: afe4bed0-d5b1-4e92-a4a4-0e0f5462a2c0 (source: [..\input\plain_text\showcase-2023-06.txt])
    child_doc_id: ee66ebc1-c455-4a8a-9c7d-63eb46de8c14 (source: [..\input\plain_text\showcase-2023-06.txt])
    child_doc_id: 98cd9fe7-5361-4802-900e-18c4460cbecc (source: [..\input\plain_text\showcase-2023-06.txt])
parent_doc_id: ebf1de83-a5fc-44fe-b2c2-2d59471ac817
    child_doc_id: 4446c2d6-15c1-4017-9901-0d4c7ed7f12d (source: [None])
    child_doc_id: c387422a-a310-4acb-bafe-1b769c1fe17a (source: [None])
    child_doc_id: 3ba6c23e-d589-4c0a-850d-0ef2a9ec921b (source: [None])
    child_doc_id: 3a54e716-1bfb-40a2-a2a1-5a0523bb328d (source: [None])
    child_doc_id: 2351ea4e-a89a-4c6d-8c38-6e5464dcdbc6 (source: [None])
    child_doc_id: a08b4cce-ef99-4105-9af8-f30abcdc7267 (source: [None])
    child_doc_id: 87445710-53d2-47bc-9909-ee9d8541e3f8 (source: [None])
    child_doc_id: 7a0920a9-039d-4c54-a46f-047245b61cd7 (source: [None])
    child_doc_id: b8f3d2cb-df1d-425a-b506-836ea1e3dcf2 (source: [None])
    child_doc_id: f7649d66-2945-455b-b008-518ffd47065e (source: [None])
parent_doc_id: c341d25a-d9c4-41c8-9235-dfc927b97672
    child_doc_id: 8f8bcc04-fe23-44c0-86c3-fac6d33021eb (source: [None])
    child_doc_id: 8de31b73-a303-4785-ad14-4fb0300d38db (source: [None])
    child_doc_id: d2ad14b3-163c-4f87-815f-4a0f3a0ff20a (source: [None])

気付いた点としては以下でした。

• ソースファイルが同一のリーフノードが必ずしも同じ親ノードを持つわけではない
• ルートノードは複数の場合がある
• 親子の階層構造は複数の場合もある

特に１点目は単純にファイルの並び順で違うソースファイルが同じ親ノードに要約されるため注意が必要かなと思います。

また２点目についてもルートノードといいつつ、ルートノードは複数あることがほとんどでるため、そちらも注意が必要です。

親ノードが保持する要約結果

親ノードにどのような要約結果が保持されているか見てみましょう。

for k,v in index.docstore.docs.items():
    node = v
    document_id = k

    # 子ノードのdoc_idを取得
    children_doc_ids = index.index_struct.get_children(node)

    # 末端(リーフノード)は飛ばす
    if len(children_doc_ids) == 0:
        continue

    print(f"parent_doc_id: {document_id}")

    print(f"\ttext: {node.text}")

parent_doc_id: 942fc862-c14a-4036-976a-a40ce2021af4
    text: The file contains a conversation about the Developers IEO in Las Vegas 2022 event. The speakers introduce themselves as members of ClasMethod and AWS. They discuss various topics including CDK updates, Amazon OpenSearch service, Amazon ECS Service Connect, and vulnerability testing tools. They mention the benefits of serverless support in OpenSearch and the improved communication between services in ECS. They also mention the potential use of Java in Lambda functions due to improved cold start performance. Additionally, they discuss the Inspector tool for vulnerability detection and the excitement about a new service called Wicker that offers secure communication.
parent_doc_id: d4cc8c81-0e5b-4c93-8dc7-2885181970e8
    text: The text files contain discussions about various topics related to developers and AWS services at the Developers IO conference in Las Vegas in 2022. The discussions touch on issues such as communication between VPCs, the use of transit gateways and multi-account setups, and the challenges of troubleshooting login issues and security configurations. There is also mention of upcoming features like ECS container runtime security and ML governance. The participants express interest in updates and announcements related to application development, ML projects, and AI explainability. The discussions also highlight the importance of services like data clean rooms and tools for managing ML projects efficiently.
parent_doc_id: de5b873a-75b1-48ef-a054-be9265f15865
    text: The session discussed the construction of machine learning models using geospatial data and satellite images for segmentation. AWS provides pre-trained models and large satellite image datasets can be accessed through the AWS Political Maker Studio. The session also covered the use of sharp leverage and game theory in table data and the introduction of natural language processing models. The speaker expressed excitement about the Sim Space Weaver service for spatial simulation. The service allows for the creation of virtual spaces for events like summer festivals. The speaker also mentioned the trend of focusing on the entire system from data sources to machine learning to ensure safe and cost-effective usage. AWS aims to address real-world challenges faced by developers in machine learning projects. The speaker highlighted the usefulness of the released features in practical scenarios and expressed interest in incorporating them into their work. The session also touched on the use of large-scale inference models and the increase in data sources and connections to Redshift. The speaker expressed enthusiasm for innovative and interesting features. The speaker looked forward to the final sessions and expected more exciting features to be announced. The speaker mentioned the challenges faced in the Chalk Talk sessions and the overwhelming response from participants. The speaker specifically mentioned the discussion on inference instance sizes and the TRUSTERS/TRN1 instance for large-scale inference.

...以降略...

ソースファイルは日本語のはずですが、要約結果は英語となっています。

これは、以降で述べるようにTreeIndexがデフォルトで使う要約用のプロンプトが英語であるためと考えられます。

親ノードを作るか否かの判断について

TreeIndexをインスタンス化する際に、num_childrenを指定可能となっており、デフォルトが10となっています。

ここを以下のように20に変更してみましょう。

from llama_index import SimpleDirectoryReader
from llama_index import TreeIndex
from llama_index import ServiceContext
from llama_index.callbacks import CallbackManager, LlamaDebugHandler

documents = SimpleDirectoryReader(
    input_dir=INPUT_DIR,
    file_metadata=lambda x: {"file_name": str(x)}
).load_data()

llama_debug_handler = LlamaDebugHandler()
callback_manager = CallbackManager([llama_debug_handler])
service_context = ServiceContext.from_defaults(callback_manager=callback_manager)

index = TreeIndex.from_documents(documents,
    num_children=20,
    service_context=service_context)

数を確認してみますと、20個が上限となっていることが分かります。

for k,v in index.docstore.docs.items():
    node = v
    document_id = k

    # 子ノードのdoc_idを取得
    children_doc_ids = index.index_struct.get_children(node)

    # 末端(リーフノード)は飛ばす
    if len(children_doc_ids) == 0:
        continue

    print(f"parent_doc_id: {document_id}, child_num: {len(children_doc_ids)}")

parent_doc_id: 67645929-7185-4c77-af5d-cd6125f5bcb4, child_num: 20
parent_doc_id: da03bf14-dbd3-4ef6-a810-a767374267ed, child_num: 20
parent_doc_id: ea158374-9e8a-4b81-a618-df21f010326f, child_num: 20
parent_doc_id: 141a3755-8a3c-4c23-bb1b-9eae5f05060c, child_num: 20
parent_doc_id: 90d16d72-4c07-4d53-a999-9b0c69d9062c, child_num: 20
parent_doc_id: 2564137e-a7f1-45e0-95cf-b5436e97f0f1, child_num: 20
parent_doc_id: 01a50446-d74f-4405-bbea-100b39d17de2, child_num: 3

親ノードが作られる方法

親ノードは冒頭で申し上げた通り、子ノードの要約されたテキストを保持します。

要約に使われるプロンプトは、TreeIndexのインスタンス化時に与えることができますが、デフォルトでは以下のようにDEFAULT_SUMMARY_PROMPTが使用されます。

DEFAULT_SUMMARY_PROMPTは以下のように英語となっています。

DEFAULT_SUMMARY_PROMPT_TMPL = (
    "Write a summary of the following. Try to use only the "
    "information provided. "
    "Try to include as many key details as possible.\n"
    "\n"
    "\n"
    "{context_str}\n"
    "\n"
    "\n"
    'SUMMARY:"""\n'
)

DEFAULT_SUMMARY_PROMPT = Prompt(
    DEFAULT_SUMMARY_PROMPT_TMPL, prompt_type=PromptType.SUMMARY
)

プロンプトは英語であるため、親ノードが持つ要約結果も英語となってしまっていることが多いです。

要約プロンプトを日本語に差し替えてみる

デフォルトのプロンプトを参考に日本語のプロンプトに変更してみます。

from llama_index.prompts.prompt_type import PromptType
from llama_index.prompts import Prompt

summpary_prompt_template = (
    "以下に関する要約を書きなさい。提供された情報のみを使用するようにしてください。"
    "その際、できるだけ多くの重要な詳細を含めるようにしてください。"
    "\n"
    "\n"
    "{context_str}\n"
    "\n"
    "\n"
    '要約文:"""\n'
)

summpary_prompt = Prompt(
    summpary_prompt_template, prompt_type=PromptType.SUMMARY
)

こちらをTreeIndexインスタンス化時に与えればOKです。

from llama_index import SimpleDirectoryReader
from llama_index import TreeIndex
from llama_index import ServiceContext
from llama_index.callbacks import CallbackManager, LlamaDebugHandler

documents = SimpleDirectoryReader(
    input_dir=INPUT_DIR,
    file_metadata=lambda x: {"file_name": str(x)}
).load_data()

llama_debug_handler = LlamaDebugHandler()
callback_manager = CallbackManager([llama_debug_handler])
service_context = ServiceContext.from_defaults(callback_manager=callback_manager)

index = TreeIndex.from_documents(documents,
    service_context=service_context,
    summary_template=summpary_prompt)

再度、親ノードに格納されているテキストを確認してみましょう。

for k,v in index.docstore.docs.items():
    node = v
    document_id = k

    # 子ノードのdoc_idを取得
    children_doc_ids = index.index_struct.get_children(node)

    # 末端(リーフノード)は飛ばす
    if len(children_doc_ids) == 0:
        continue

    print(f"parent_doc_id: {document_id}")

    print(f"\ttext: {node.text}")

以下のように日本語となっていることが確認できました。

parent_doc_id: 0d39b415-5104-445a-89b2-d8e8ab36fc2f
    text: クラスメソッドがラスベガスで開催されたDevelopers IEO in Las Vegas 2022の様子を報告している。参加者は104人で、CDKのアップデートやAmazonオープンサーチサービスのサーバーレス対応、Amazon ECSの新しいネットワーク機能などが注目されている。また、コンテナ間通信やマイクロサービス間通信の問題に対する解決策として、Amazon ECSサービスコネクトが紹介されている。さらに、Javaのコールドスタートの改善により、Javaのフレームワークも選択肢に入るようになった。また、インスペクターという脆弱性検証ツールのアップデートにより、ラムダのソースコードの脆弱性の確認が可能になった。AWSの新しいサービスや機能に期待が寄せられている。
parent_doc_id: 44566f96-639e-4e19-8cb4-5f1e4231e72f
    text: 開発者会議での話題は、VPC間の疎通やトランジットゲートウェイの使用、マルチアカウントでのVPC集約など、複雑なネットワーク構成に関するものが多かった。また、AWSによるコンテナのランタイムセキュリティの提供についても注目されており、セキュリティ商品の選択肢が限られている中でAWSが提供することに期待が寄せられている。さらに、アプリケーション開発に関連する更新サービスやML Governance、説明可能なAIなどについても期待が高まっている。
parent_doc_id: 07b84de4-862a-45ea-875b-85f6e6abfaf1
    text: このセッションでは、地理空間を使った機械学習や地理情報データを利用したモデル構築について話されました。AWS上の政治メーカーのスタジオで準備された衛星画像などの大きなデータを利用することも可能で、実際に機能とデモが行われました。また、自然言語処理のモデルやゲームモデルの説明可能性についても最先端の研究が紹介されました。さらに、空間シミュレーションのサービスや機械学習のトレンドについても話されました。AWSは実際のニーズに応えるためにリリースイベントで機能を提供しており、開発者はこれらの機能を活用して改善やプレイを進めることができると感じています。

...以降略...

TreeIndexのRetriever Modeとプロンプト

その１の記事に記載した通り、インデックスはその種類毎に使用可能なRetrieverModeで、ノードの探索方法を変更することが可能です。

以下がTreeIndexで使用可能なRetriever Modeとなります。

代表的なTreeRetrieverMode.SELECT_LEAF（デフォルト）の場合、ノードを選択に関する問い合わせをLLMを使って行います。

具体的にはquery_templateまたはquery_template_multipleに与えられるプロンプトを使ってLLMがノード選択を行います。

これらはデフォルトでは以下のように英語となっています。

# # single choice
DEFAULT_QUERY_PROMPT_TMPL = (
    "Some choices are given below. It is provided in a numbered list "
    "(1 to {num_chunks}),"
    "where each item in the list corresponds to a summary.\n"
    "---------------------\n"
    "{context_list}"
    "\n---------------------\n"
    "Using only the choices above and not prior knowledge, return "
    "the choice that is most relevant to the question: '{query_str}'\n"
    "Provide choice in the following format: 'ANSWER: <number>' and explain why "
    "this summary was selected in relation to the question.\n"
)
DEFAULT_QUERY_PROMPT = Prompt(
    DEFAULT_QUERY_PROMPT_TMPL, prompt_type=PromptType.TREE_SELECT
)

# multiple choice
DEFAULT_QUERY_PROMPT_MULTIPLE_TMPL = (
    "Some choices are given below. It is provided in a numbered "
    "list (1 to {num_chunks}), "
    "where each item in the list corresponds to a summary.\n"
    "---------------------\n"
    "{context_list}"
    "\n---------------------\n"
    "Using only the choices above and not prior knowledge, return the top choices "
    "(no more than {branching_factor}, ranked by most relevant to least) that "
    "are most relevant to the question: '{query_str}'\n"
    "Provide choices in the following format: 'ANSWER: <numbers>' and explain why "
    "these summaries were selected in relation to the question.\n"
)
DEFAULT_QUERY_PROMPT_MULTIPLE = Prompt(
    DEFAULT_QUERY_PROMPT_MULTIPLE_TMPL, prompt_type=PromptType.TREE_SELECT_MULTIPLE
)

この２つのプロンプトは、子ノードを何個ずつ選んで辿っていくかで代わります。

その設定を行うのがchild_branch_factorとなりデフォルトでは1となっています。

child_branch_factorは親ノードごとの個数ではなく、すべての親ノードに属する子ノードの中からchild_branch_factor個選択する動作となります。

そのため、その階層の深さに寄らず、最終的に選ばれるノード数はchild_branch_factor個となります。

今回は、child_branch_factorを3にしつつ、query_template_multipleを日本語にしてクエリしてみます。

クエリを試してみる

まずはDEFAULT_QUERY_PROMPT_MULTIPLE_TMPLを参考にして、日本語のプロンプトを以下のように準備します。

query_prompt_multiple_tmpl = (
    "以下にいくつかの選択肢を示す。"
    "これは番号付きリスト(1から{num_chunks})で提供され、リストの各項目は要約に対応する。"
    "---------------------"
    "{context_list}"
    "---------------------"
    "予備知識ではなく、上記の選択肢のみを使用して、以下の質問に最も関連する上位の選択肢を"
    "最も関連性の高いものから低いものへランク付けで返してください。"
    "質問 : {query_str}"
    "その際、選択肢の数は必ず{branching_factor}個以下にとなるようにしてください。"
    ""
    "また応答は以下のフォーマットで選択肢を提供してください。"
    "ANSWER： <numbers> と、これらの要約がなぜ質問に関連して選択されたかの説明"
)
default_query_prompt_multiple = Prompt(
    query_prompt_multiple_tmpl, prompt_type=PromptType.TREE_SELECT_MULTIPLE
)

次に以下のようにquery_engineをインスタンス化します。

from llama_index.indices.tree.base import TreeRetrieverMode

query_engine = index.as_query_engine(
    retriever_mode=TreeRetrieverMode.SELECT_LEAF,
    child_branch_factor=3,
    query_template_multiple=default_query_prompt_multiple)

ここまでのログを破棄して結果を確認します。

llama_debug_handler.flush_event_logs()

response = query_engine.query("機械学習に関するre:Invent 2022でのアップデートについて300字で教えてください。")

print(response.response)

得られた結果は以下です。

re:Invent 2022での機械学習に関するアップデートはいくつかあります。まず、Amazonオープンサーチサービスがサーバーレス対応となり、柔軟なスケーリングや管理の簡便さが向上しました。また、Amazon ECSの新しいネットワーク機能であるAmazon ECSサービスコネクトが発表され、サービス間の通信を効率的に行うことができるようになりました。さらに、MLプロジェクトの管理やモデルの説明を支援するモデルカードやAIサービスカードも新たに導入されました。これらのアップデートにより、機械学習の運用や開発がより効率的に行えるようになりました。

ノードは以下のように３つ選ばれていることが確認できました。

for node in response.source_nodes:
    pprint(node.node.id_)

'018bbc14-beee-4be7-8077-702db4bffce8'
'a6da34d0-c523-491b-b4de-d42383b51a72'
'203a45c7-2b90-448a-9526-000e445789bd'

LLMへの問い合わせログはLlamaDebugHandlerに３件登録されます。

このうち始めの２個がノードを探索する際の問い合わせとなり、最後の一つがユーザのクエリの問い合わせとなります。

以下で確認することができます。

from llama_index.callbacks import CBEventType

llama_debug_handler.get_event_pairs(CBEventType.LLM)

（出力は割愛いたします）

その他本記事で詳しく扱わなかったことについて

その他、TreeIndex関連で本記事で詳しく扱わなかったことは以下です。

• insert_prompt
  • ツリー構造のインデックスに新しいデータを挿入する際に使われるプロンプトです
  • summary_promptと同様にTreeIndexのインスタンス化時に与えます
• TreeRetrieverMode.SELECT_LEAF_EMBEDDING
  • こちらはListIndexと同様にノード情報にベクトルを含めてあげることでベクトルによるノード選択が可能です
  • ベクトルの追加方法はその４の記事で取り扱っていますのでそちらを参考にされてください。
• TreeRetrieverMode.ALL_LEAFとTreeRetrieverMode.ROOT
  • これらは自明かなと思いますのでここでは割愛します
• query時の応答合成に使用されるプロンプト
  • 説明したquery_template、query_template_multiple以外にもListIndexと同様、応答の合成に使用するプロンプトが別途準備されています
  • そちらもListIndexと同様に必要に応じてカスタマイズが可能です
• 子ノード数について
  • 子ノードを設定で増やしすぎると要約作成時にmax_tokenをオーバーする可能性があります
  • max_tokenをオーバーした場合は設定を見直しましょう

まとめ

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

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