見出し画像

wandb による LlamaIndex のトレースを試す

npaka

「wandb」による「LlamaIndex」のトレースを試したので、まとめました。


1. wandb (Weights & Biases)

「wandb」 (Weights & Biases)は、機械学習の実験の追跡および可視化プラットフォームです。「LlamaIndex」ではコールバックハンドラを1つ指定するだけで、「LlamaIndex」の内部状態をトレースすることができます。

「wandb」のサイトでアカウントを作成してください。個人利用は無料になります。詳しくは、「Pricing」を参照してください。

2. ドキュメントの準備

はじめに、チャットボットに教える専門知識を記述したドキュメントを用意します。

今回は、マンガペディアの「ぼっち・ざ・ろっく！」のあらすじのドキュメントを用意しました。

・bocchi.txt

3. トレースの実行

トレースの実行手順は、次のとおりです。

(1) パッケージのインストール。
「wandb」パッケージもインストールします。

# パッケージのインストール
!pip install llama-index
!pip install openai
!pip install wandb  # wandbも追加

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

import openai

# OpenAI APIキーの準備
openai.api_key = "<OpenAI_APIのAPIキー>"

(3) Wandbコールバックハンドラの準備。

from llama_index.callbacks import LlamaDebugHandler, WandbCallbackHandler, CallbackManager
from llama_index import ServiceContext

# Wandbコールバックハンドラの準備
llama_debug = LlamaDebugHandler(print_trace_on_end=True)
wandb_callback = WandbCallbackHandler(
    run_args=dict(
        project="llamaindex",  # プロジェクト名
    )
)
callback_manager = CallbackManager([llama_debug, wandb_callback])

「Google Colab」のログで、wandbのAPIキーの入力を促されたら、「https://wandb.ai/authorize」をクリックしてAPIキーをコピーし、テキストボックスに入力します。

(4) LlamaIndexの利用。
サービスコンテキストにwandbコールバックハンドラを設定する以外は、いつもと同じ使い方です。

from llama_index import SimpleDirectoryReader, GPTVectorStoreIndex

# サービスコンテキストの準備
service_context = ServiceContext.from_defaults(
    callback_manager=callback_manager,  # wandbコールバックハンドラ
    llm=llm
)

# ドキュメントの読み込み
docs = SimpleDirectoryReader("data").load_data()

# インデックスの準備
index = GPTVectorStoreIndex.from_documents(
    docs, 
    service_context=service_context
)

# クエリエンジンの準備
query_engine = index.as_query_engine()

# 質問応答
print(query_engine.query("ぼっちちゃんの得意な楽器は?"))
**********
Trace: index_construction
    |_node_parsing ->  0.170559 seconds
      |_chunking ->  0.166955 seconds
    |_embedding ->  0.591901 seconds
    |_embedding ->  0.505997 seconds
**********
wandb: Logged trace tree to W&B.
**********
Trace: query
    |_query ->  0.941976 seconds
      |_retrieve ->  0.150257 seconds
        |_embedding ->  0.141673 seconds
      |_synthesize ->  0.791522 seconds
        |_llm ->  0.787927 seconds
**********
wandb: Logged trace tree to W&B.

得意な楽器はギターです。

(5) wandbのダッシュボードでトレース内容を確認。
「Google Colab」のログのリンク (https://wandb.ai/<user>/uncategorized/runs/<id>) からも開くことができます。

「LlamaIndex」内でどのドキュメントを抽出してどのようにレスポンスを生成したかを確認できるため、レスポンス失敗時にどの処理が原因なのかがわかります。

参考



この記事が気に入ったらサポートをしてみませんか？