MLFlow 2.11で追加されたChatModelを試してみる on Databricks

# 導入

2024年3月にMLFlowのバージョン2.11がリリースされました。

https://mlflow.org/releases/2.11.0

リリースノートを見ていると、[「ChatModel interface for a unified Chat experience with pyfunc models」](https://mlflow.org/releases/2.11.0#chatmodel-interface-for-a-unified-chat-experience-with-pyfunc-models
)という内容があり、気になったので確認してみました。


# どういう機能？

おそらく以下のチュートリアルがこの機能に対応する内容だと思われます。

https://mlflow.org/docs/latest/llms/transformers/tutorials/conversational/pyfunc-chat-model.html?highlight=chatmodel

一部抜粋&邦訳。

> MLflow 2.11 以降では、[ChatModel](https://mlflow.org/docs/latest/python_api/mlflow.types.html#) クラスが追加され、OpenAI API 仕様に準拠した提供されるモデルを簡単に作成できるようになりました。これにより、チャットアプリのバッキングLLMをシームレスに交換したり、クライアント側のコードを編集することなくさまざまなモデルを簡単に評価したりできます。

内容を見るに、OpenAI API仕様に準拠したモデルとして、MLFlowにモデルを登録・利用できる模様。

挙動をシンプルに確認するために、ダミーの[カスタムChatModelクラスを作成](https://mlflow.org/docs/latest/llms/transformers/tutorials/conversational/pyfunc-chat-model.html?highlight=chatmodel#Customizing-the-model)し、どのような挙動となるか見てみます。

確認はDatabricks on AWS上で実施しました。

:::note warn
[ChatModel](https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html?highlight=chatmodel#mlflow.pyfunc.ChatModel)はDocumentを見るにまだ実験クラスのようで、将来的に警告なしに変更・削除される可能性があります。
:::


# Step1. カスタムChatModelクラスを作成

カスタムChatModelクラスを定義します。
カスタムChatModelクラスは```mlflow.pyfunc.ChatModel```を継承するクラスを定義することで実装します。

[従来のカスタムpyfuncクラス](https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html#pyfunc-create-custom)と同様、```predict```メソッドをオーバーライドする必要があるのですが、引数の内容が少し異なり、```messages```という引数を取るようになります。

今回は挙動を確認するためにこの```messages```引数の内容と```params```引数の内容を出力し、ダミーの推論結果を返すように実装します。

```python
from typing import List

import mlflow
from mlflow.types.llm import ChatResponse

class MyChatModel(mlflow.pyfunc.ChatModel):
    def load_context(self, context):
        pass

    def predict(
        self,
        context,
        messages: List[mlflow.types.llm.ChatMessage],
        params: mlflow.types.llm.ChatParams,
    ):
        
        print(messages)
        print(params)

        id = "0"
        text = "Sample Text"
        usage = {
            "prompt_tokens": 10,
            "completion_tokens": 20,
            "total_tokens": 30,
        }

        response = {
            "id": id,
            "model": "MyChatModel",
            "choices": [
                {
                    "index": 0,
                    "message": {"role": "assistant", "content": text},
                    "finish_reason": "stop",
                }
            ],
            "usage": usage,
        }

        return ChatResponse(**response)
```

# Step2. mlflowにカスタムモデルを登録

最新のMLFlowをインストールし、Step1で作成したモデルを登録してみます。

```python
%pip install -U -qq "mlflow-skinny[databricks]>=2.11.1"

dbutils.library.restartPython()
```

```python
import mlflow

# Unity Catalogのモデルとして登録
mlflow.set_registry_uri("databricks-uc")

with mlflow.start_run() as run:
    _ = mlflow.pyfunc.log_model(
        artifact_path="model",
        python_model=MyChatModel(),
        registered_model_name="training.llm.test_chatmodel",  # 登録モデル名 in Unity Catalog
    )
```

無事に登録されました。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/310ba7c5-de05-00f5-263a-04f2cbe9b5ef.png)

今更ですが、mlflowはモデル登録の際に```predict```を一度試験実行するのですね。
```predict```メソッド内のprint文の実行結果が表示されています。

登録した結果、モデルスキーマは以下のように登録されました。
画像上は全て表示できていませんが、log_model時の```signature```引数を指定していなくても規定の入力/出力が登録されるようです。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/0ac1c9a7-64c9-64d7-1b01-caed7d1e8945.png)

逆に、```signature```引数を指定すると以下のエラーが発生しました。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/304ae682-ea24-8f47-b9fb-75846f76c3d7.png)


# Step3. モデルのロードと出力確認

では、登録したモデルをロードして、挙動を確認してみます。

```python
# 最新モデルバージョンの特定
model_uri = f"models:/training.llm.test_chatmodel/1"
model = mlflow.pyfunc.load_model(model_uri)
```

まずは入力/出力/パラメータのスキーマを確認。

```python
model.metadata.signature
```

```:出力
inputs: 
  ['messages': Array({content: string (required), name: string (optional), role: string (required)}) (required), 'temperature': double (optional), 'max_tokens': long (optional), 'stop': Array(string) (optional), 'n': long (optional), 'stream': boolean (optional)]
outputs: 
  ['id': string (required), 'object': string (required), 'created': long (required), 'model': string (required), 'choices': Array({finish_reason: string (required), index: long (required), message: {content: string (required), name: string (optional), role: string (required)} (required)}) (required), 'usage': {completion_tokens: long (required), prompt_tokens: long (required), total_tokens: long (required)} (required)]
params: 
  None
```

[OpenAI API Chat仕様](https://platform.openai.com/docs/api-reference/chat/create)の主要パラメータをカバーしているのがわかります。

では、ダミーの推論処理を実行してみます。

```python
from pprint import pprint

messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "Write me a hello world program in python"},
]

result = model.predict({"messages": messages, "temperature": 0.1})

print("-----")

print(type(result))
pprint(result)
```

```:出力
[ChatMessage(role='system', content='You are a helpful assistant.', name=None), ChatMessage(role='user', content='Write me a hello world program in python', name=None)]
ChatParams(temperature=0.1, max_tokens=None, stop=None, n=1, stream=False)
-----
<class 'dict'>
{'choices': [{'finish_reason': 'stop',
              'index': 0,
              'message': {'content': 'Sample Text', 'role': 'assistant'}}],
 'created': 1709979715,
 'id': '0',
 'model': 'MyChatModel',
 'object': 'chat.completion',
 'usage': {'completion_tokens': 20, 'prompt_tokens': 10, 'total_tokens': 30}}
```

出力結果から、辞書型形式で```predict```メソッドに渡した```messages```等のパラメータは、```predict```メソッド内ではChatMessagesとChatParamsのそれぞれに分かれて格納されて利用可能になっています。
```ChatResponse```型で返した結果は、いくつかの情報を付与されて辞書型のインスタンスとして返されているのもわかります。

# まとめ

MLFlow 2.11で実装されたpyfunc ChatModelを試用し、挙動を確認してみました。

最近はOpenAI API互換仕様で推論サーバを提供するパッケージも多いため、簡単に同種のモデルを作ることが出来るのは便利ですね。
チュートリアルでは実際に```mlflow models serve```機能でOpenAI API互換のAPIサーバを実行していましたし、手軽にサーバを立てることも出来そうです。
また、FastAPI等を使ってPython上カスタムしたAPIサーバを手軽に建てることもできそうです。

今回は挙動確認だけでしたが、実際にLLMによる推論処理含めたカスタムモデルの作成もやってみようと思います。