MLflow LlamaIndexフレーバーを試す

# 導入

2024/7月末にMLflow 2.15.0がリリースされました。（現時点の最新は2.15.1）

https://mlflow.org/releases/2.15.0

このバージョンでリリースされた主要機能の一つにLlamaIndex Flavorの追加があります。

> - **LlamaIndex フレーバー🦙** - MLflowは、カスタムデータを中心としたGenAIアプリを構築するための最も人気のあるライブラリの1つであるLlamaIndexとのネイティブ統合を提供します。この統合により、MLflow内でLlamaIndexインデックスをログに記録し、異なるエンジンタイプでの推論タスクのためにインデックスされたデータをロードおよびデプロイすることができます。MLflowはまた、LlamaIndex操作の包括的なトレースサポートを提供し、複雑なクエリに対する前例のない透明性を提供します。[MLflow LlamaIndexドキュメント](https://mlflow.org/docs/latest/llms/llama-index/index.html)をチェックして始めましょう！

普段LlamaIndexをあまり使ったことが無かったのですが、最近のLlamaIndexアップデートで気になる機能などもあり（いつか記事書きたい）、勉強の意味でMLflow LlamaIndexのQuick StartをDatabricks上でウォークスルーしてみます。
※ 一部、OpenAIのサービスではなく、Databricksのモデルサービングを利用するように変更しています。

https://mlflow.org/docs/latest/llms/llama-index/notebooks/llama_index_quickstart.html

ちなみに、DatabricksでのLlamaIndex利用は以下のような素晴らしい記事が既にありますので参照ください。

https://qiita.com/taka_yayoi/items/e56ee77ad6779088b790

https://qiita.com/taka_yayoi/items/3bbf70806b31a7ca148c

検証環境はDatabricks on AWS、ノートブックのクラスタは"サーバレス"です（便利！）。

各ステップの解説はQuick Startの内容を邦訳する形で記載しています。

# Step0. 学習内容

> このチュートリアルの終わりまでに、以下のことができるようになります：
> 
> - LlamaIndexでMVP VectorStoreIndexを作成する。
> - そのインデックスをMLflowトラッキングサーバーにログする。
> - そのインデックスをMLflowモデルレジストリに登録する。
> - モデルをロードして推論を行う。
> - ログされたアーティファクトについてMLflow UIを探索する。
> 
> これらの基本を通じて、MLflowにおけるLlamaIndexのユーザージャーニーに慣れ親しむことができます。

# Step1. セットアップ

> まず、必要な依存関係と環境変数を確認する必要があります。デフォルトでは、LlamaIndexはLLMと埋め込みモデルのソースとしてOpenAIを使用するため、ここでも同様に行います。必要なライブラリをインストールし、OpenAI APIキーを提供することから始めましょう。

**註:** 今回はOpenAIのAPIを利用しないため、OpenAI APIキーは設定していません。

まず、パッケージをインストール。

```python
%pip install "mlflow-skinny[databricks]>=2.15.1" llama-index>=0.10.44 llama-index-llms-databricks llama-index-embeddings-databricks

dbutils.library.restartPython()
```

モジュールのインポートと、環境変数にDatabricksモデルサービングへアクセスするためのAPIキーやエンドポイントURLを設定します。
下記記載のURLはダミーです。実際のURLはDatabricksのサービング画面から確認してください。

```python

from llama_index.core import Document, VectorStoreIndex, Settings
from llama_index.core.llms import ChatMessage
from llama_index.embeddings.databricks import DatabricksEmbedding
from llama_index.llms.databricks import Databricks

import os
import mlflow

# 試験なので、ノートブック実行環境からAPIキーを取得
API_TOKEN = (
    dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiToken().get()
)

os.environ["DATABRICKS_TOKEN"] = API_TOKEN
os.environ["DATABRICKS_SERVING_ENDPOINT"] = "https://xxxxxx.com/serving-endpoints/"
```

LlamaIndex上で標準的に利用する埋め込みモデルとLLMを指定します。
今回は、以下の記事の際に作成したエンドポイントを利用しました。

https://qiita.com/isanakamishiro2/items/66eaedf438fd62913804

https://qiita.com/isanakamishiro2/items/7ce94c22a0ad5d2e59dd

```python
embed_model = DatabricksEmbedding(model="embedding_bge_m3_endpoint")
Settings.embed_model = embed_model

llm = Databricks(model="EZO-Common-T2-2B-gemma-2-it-sglang-endpoint")
Settings.llm = llm
```

# Step2. インデックスの作成

> ベクトルストアインデックスはLlamaIndexの主要なコンポーネントの一つです。これらは、取り込まれたドキュメントチャンクの埋め込みベクトル（時にはドキュメントチャンク自体も）を含んでいます。これらのベクトルは、クエリエンジン、チャットエンジン、リトリーバーなど、LlamaIndexで異なる目的に応じた推論を可能にします。
> 
> 1. **クエリエンジン:**
> 
>     - **使用法:** ユーザーの質問に基づいて関連情報を取得するための簡単なクエリを実行します。
>     - **シナリオ:** 検索エンジンのように、特定のクエリに一致する簡潔な回答やドキュメントを取得するのに理想的です。
> 
> 2. **チャットエンジン:**
> 
>     - **使用法:** 複数のインタラクションにわたってコンテキストと履歴を維持する必要がある会話型AIタスクに従事します。
>     - **シナリオ:** カスタマーサポートボットやバーチャルアシスタントのような対話型アプリケーションに適しており、会話のコンテキストが重要です。
> 
> 3. **リトリーバー:**
> 
>     - **使用法:** 与えられた入力に対して意味的に類似したドキュメントやテキストセグメントを取得します。
>     - **シナリオ:** 要約や質問応答のタスクで生成された回答の質を向上させるために、関連するコンテキストや背景情報を取得するRAGシステムに役立ちます。
> 
> これらの異なるタイプの推論を活用することで、LlamaIndexはユーザーと大規模言語モデルの間のインタラクションを強化し、さまざまなユースケースに合わせた堅牢なAIアプリケーションを構築することができます。

というわけでインデックスを作成し、クエリやチャット、Retrieverを実行してみます。

```python
print("------------- Example Document used to Enrich LLM Context -------------")
llama_index_example_document = Document.example()
print(llama_index_example_document)

index = VectorStoreIndex.from_documents([llama_index_example_document])

print("\n------------- Example Query Engine -------------")
query_response = index.as_query_engine().query("What is llama_index?")
print(query_response)

print("\n------------- Example Chat Engine  -------------")
chat_response = index.as_chat_engine().chat(
    "What is llama_index?",
    chat_history=[ChatMessage(role="system", content="You are an expert on RAG!")],
)
print(chat_response)


print("\n------------- Example Retriever   -------------")
retriever_response = index.as_retriever().retrieve("What is llama_index?")
print(retriever_response)
```

```:出力
------------- Example Document used to Enrich LLM Context -------------
Doc ID: 1a0e2277-45db-4830-91ba-a7969d3a2c11
Text: Context LLMs are a phenomenal piece of technology for knowledge
generation and reasoning. They are pre-trained on large amounts of
publicly available data. How do we best augment LLMs with our own
private data? We need a comprehensive toolkit to help perform this
data augmentation for LLMs.  Proposed Solution That's where LlamaIndex
comes in. Ll...

------------- Example Query Engine -------------
LlamaIndex is a data framework designed to augment Large Language Models (LLMs) with private data. It offers tools to ingest data from various sources (APIs, PDFs, documents, etc.), structure it into indices and graphs for LLM usage, and

------------- Example Chat Engine  -------------
LlamaIndex is a data framework that allows users to incorporate private data into Large Language Models (

------------- Example Retriever   -------------
[NodeWithScore(node=TextNode(id_='6d5539ac-e03e-445f-ac23-956b0d452f8d', embedding=None, metadata={'filename': 'README.md', 'category': 'codebase'}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={<NodeRelationship.SOURCE: '1'>: RelatedNodeInfo(node_id='1a0e2277-45db-4830-91ba-a7969d3a2c11', node_type=<ObjectType.DOCUMENT: '4'>, metadata={'filename': 'README.md', 'category': 'codebase'}, hash='3183371414f6a23e9a61e11b45ec45f808b148f9973166cfed62226e3505eb05')}, text='Context\nLLMs are a phenomenal piece of technology for knowledge generation and reasoning.\nThey are pre-trained on large amounts of publicly available data.\nHow do we best augment LLMs with our own private data?\nWe need a comprehensive toolkit to help perform this data augmentation for LLMs.\n\nProposed Solution\nThat\'s where LlamaIndex comes in. LlamaIndex is a "data framework" to help\nyou build LLM  apps. It provides the following tools:\n\nOffers data connectors to ingest your existing data sources and data formats\n(APIs, PDFs, docs, SQL, etc.)\nProvides ways to structure your data (indices, graphs) so that this data can be\neasily used with LLMs.\nProvides an advanced retrieval/query interface over your data:\nFeed in any LLM input prompt, get back retrieved context and knowledge-augmented output.\nAllows easy integrations with your outer application framework\n(e.g. with LangChain, Flask, Docker, ChatGPT, anything else).\nLlamaIndex provides tools for both beginner users and advanced users.\nOur high-level API allows beginner users to use LlamaIndex to ingest and\nquery their data in 5 lines of code. Our lower-level APIs allow advanced users to\ncustomize and extend any module (data connectors, indices, retrievers, query engines,\nreranking modules), to fit their needs.', mimetype='text/plain', start_char_idx=1, end_char_idx=1279, text_template='{metadata_str}\n\n{content}', metadata_template='{key}: {value}', metadata_seperator='\n'), score=0.5558356896342923)]
```

# Step3. MLflowでインデックスをロギング

> 以下のコードは、MLflowを使用してLlamaIndexモデルをロギングし、異なる環境間での永続化と管理を可能にします。MLflowを使用することで、モデルの追跡、バージョン管理、および再現が確実に行えます。このスクリプトはパラメータ、入力例をログし、特定の名前でモデルを登録します。model_uriは、後でモデルを取得するための一意の識別子を提供します。この永続化は、開発、テスト、および本番環境での一貫性と再現性を確保するために不可欠です。MLflowでモデルを管理することで、ロード、デプロイ、および共有が簡素化され、整理されたワークフローが維持されます。
> 
> 主要なパラメータ
> 
> - **engine_type:** pyfuncおよびspark_udf推論タイプを定義します
> - **input_example:** 入力シグネチャを定義し、予測を通じて出力シグネチャを推測します
> - **registered_model_name:** MLflowモデルレジストリ内のモデル名を定義します

今回のメインですね。
MFflow LlamaIndex flavorを使ってロギングします。
```registered_model_name```は環境に合わせて変更しました。

```python
mlflow.llama_index.autolog()  # This is for enabling tracing

with mlflow.start_run() as run:
    mlflow.llama_index.log_model(
        index,
        artifact_path="llama_index",
        engine_type="query",  # Defines the pyfunc and spark_udf inference type
        input_example="hi",  # Infers signature
        registered_model_name="training.llm.my_llama_index_vector_store",  # Stores an instance in the model registry
    )

    run_id = run.info.run_id
    model_uri = f"runs:/{run_id}/llama_index"
    print(f"Unique identifier for the model location for loading: {model_uri}")
```

# Step4. インデックスのロードと推論の実行

> 以下のコードは、ロードされたモデルで実行できる3つの主要な推論タイプを示しています。
> 
> 1. **LlamaIndexを使用したロードと推論の実行:** この方法は、mlflow.llama_index.load_modelを使用してモデルをロードし、直接クエリ、チャット、または検索を実行します。基盤となるLlamaIndexオブジェクトの完全な機能を活用したい場合に理想的です。
> 
> 1. **MLflow PyFuncを使用したロードと推論の実行:** この方法は、mlflow.pyfunc.load_modelを使用してモデルをロードし、ログ時に指定されたエンジンタイプで汎用PyFunc形式でモデル予測を実行します。mlflow.evaluateでモデルを評価したり、サービングのためにモデルをデプロイしたりするのに便利です。
> 
> 1. **MLflow Spark UDFを使用したロードと推論の実行:** この方法は、mlflow.pyfunc.spark_udfを使用してモデルをSpark UDFとしてロードし、Spark DataFrame内の大規模データセットに対して分散推論を実行します。大規模データ処理に最適で、PyFunc推論と同様に、ログ時に定義されたエンジンタイプのみをサポートします。

LlamaIndex FlavorとPyFuncを使用したロードの2種を実行してみます。

```python
print("\n------------- Inference via Llama Index   -------------")
index = mlflow.llama_index.load_model(model_uri)
query_response = index.as_query_engine().query("hi")
print(query_response)

print("\n------------- Inference via MLflow PyFunc -------------")
index = mlflow.pyfunc.load_model(model_uri)
query_response = index.predict("hi")
print(query_response)
```

LlamaIndex FlavorだとLlamaIndexのIndexをそのまま取得できるようですね。

また、```autolog```を有効化しているため、実行結果は以下のようにMLflow Tracingの結果も同時に表示されます。便利。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/489b05e2-747d-8d02-57e0-ad9ad8d32477.png)

Quick Startはこの後UI上でのモデル内容確認がありますが、こちらは割愛します。
実際に実行して結果を見てみてください！

# まとめ

MLflowのLlamaIndex flavorを試してみました。
autologも対応しているので、LangChain同様かなり管理しやすくなったのではないかと思います。

実はIndex部分を[Databricks Mosaic AI Vector Search](https://docs.databricks.com/ja/generative-ai/vector-search.html)に変更して同様にやってみたのですが、こちらはロギング後にモデルをロードする際にエラーが発生してうまくいかず。上手くやる方法募集中です。

私個人はLangChainを使うことが多い（最近はLangGraphが多い）のですが、RAGシステムを構築する際はLlamaIndexを使うのも良いと思います。Databricks用モジュールも出てますし。
これを機にもう少し勉強したいと思います。