DatabricksのAI Playground上でSwallow-MS-7b-instruct-v0.1を試す

# 導入

以下で紹介されているDatabricksのAI Playground機能が、自分の環境下でも利用できるようになっていました。

https://qiita.com/taka_yayoi/items/3328b7eb2d841168addc

現状、日本リージョンだと自分で用意した「[プロビジョニング済みスループットのエンドポイント](https://docs.databricks.com/ja/machine-learning/foundation-models/deploy-prov-throughput-foundation-model-apis.html)」か「[外部モデルのサービング](https://docs.databricks.com/ja/generative-ai/external-models/index.html)」のみが利用可能なようなので、自前でプロビジョニングされたスループットの基盤モデルを用意して試してみます。

せっかくなので、先日ひっそり（？）とリリースされていた、TokyoTech-LLMのSwallow-MS-7b-instruct-v0.1を使ってます。

https://huggingface.co/tokyotech-llm/Swallow-MS-7b-instruct-v0.1

# Step1. 準備

Databricks上でノートブックを作成し、以下の処理を実行することで「プロビジョニングされたスループットの基盤モデル」を準備します。
(DBRは14.3ML、インスタンスタイプはg5.xlarge(AWS)で実行しました)

## パッケージインストール

```transformers```と```mlflow```を最新化します。
MLflowは2.12.1以降をインストールしてください。

```python
%pip install -U "mlflow>=2.12.1" transformers accelerate

dbutils.library.restartPython()
```

## MLflow登録

[こちらのモデル](https://huggingface.co/tokyotech-llm/Swallow-MS-7b-instruct-v0.1)を事前にダウンロードしておき、```transformers```で読み込みます。

```python
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_path = "/Volumes/training/llm/model_snapshots/models--tokyotech-llm--Swallow-MS-7b-v0.1/"

model = AutoModelForCausalLM.from_pretrained(
    model_path, 
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(
    model_path,
)
```

MLflow Model Registryにロギングします。
注意点として、MLflowのバージョンが2.12.1以降とその前で指定の仕方が異なります。
今回は2.12.1以降のやり方です。

詳しくは[公式Doc](https://docs.databricks.com/ja/machine-learning/foundation-models/deploy-prov-throughput-foundation-model-apis.html#log-fine-tuned-foundation-models)を確認ください。

```python
import mlflow
import numpy as np

# Models in UCを使っていない場合には下の行をコメントアウトします 
# そして、3レベル名前空間の名称ではなくシンプルにモデル名を指定します
mlflow.set_registry_uri('databricks-uc')
CATALOG = "training"
SCHEMA = "llm"
registered_model_name = f"{CATALOG}.{SCHEMA}.swallow-ms-7b-instruct"

# サンプル入力の定義
input_example = {
  "messages": [
    {
      "role": "system",
      "content": "You are a helpful assistant."
    },
    {
      "role": "user",
      "content": "Hello!"
    }
  ],
  "max_tokens": 10,
  "stream": false
}

# 新規MLflowランをスタート
with mlflow.start_run():
    components = {
        "model": model,
        "tokenizer": tokenizer,
    }
    mlflow.transformers.log_model(
        transformers_model=components,
        task="llm/v1/chat",
        artifact_path="model",
        registered_model_name=registered_model_name,
        input_example=input_example,
        example_no_conversion=True,
    )
```

## モデルサービングへの登録

まず、登録したモデルが「プロビジョニングされたスループットの基盤モデル」対応かどうかを確認。

```python
import requests
import json
from mlflow import MlflowClient


# 登録されたMLflowモデルの名前
CATALOG = "training"
SCHEMA = "llm"
registered_model_name = f"{CATALOG}.{SCHEMA}.swallow-ms-7b-instruct"

model_name = registered_model_name

# MLflowモデルの最新バージョンを取得
client=MlflowClient()
versions = [mv.version for mv in client.search_model_versions(f"name='{registered_model_name}'")]
model_version = versions[0]

# 現在のノートブックのコンテキストからAPIエンドポイントとトークンを取得
API_ROOT = dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiUrl().get() 
API_TOKEN = dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiToken().get()

headers = {"Context-Type": "text/json", "Authorization": f"Bearer {API_TOKEN}"}

response = requests.get(url=f"{API_ROOT}/api/2.0/serving-endpoints/get-model-optimization-info/{model_name}/{model_version}", headers=headers)

print(json.dumps(response.json(), indent=4))
```

```:出力
{
    "optimizable": true,
    "model_type": "mistral",
    "throughput_chunk_size": 970,
    "dbus": 24
}
```


optimizableがTrueのため、対象です。

では、こちらのモデルを使ってモデルサービングのエンドポイントを登録します。

```python
# MLflowエンドポイント名を設定
endpoint_name = "swallow-ms-7b-instruct-endpoint"

# プロビジョンされるスループットの最小値を指定 
min_provisioned_throughput = response.json()['throughput_chunk_size']

# プロビジョンされるスループットの最大値を指定
max_provisioned_throughput = response.json()['throughput_chunk_size']


data = {
    "name": endpoint_name,
    "config": {
        "served_entities": [
            {
                "entity_name": model_name,
                "entity_version": model_version,
                "min_provisioned_throughput": min_provisioned_throughput,
                "max_provisioned_throughput": min_provisioned_throughput,
            }
        ]
    },
}

headers = {"Context-Type": "text/json", "Authorization": f"Bearer {API_TOKEN}"}

response = requests.post(url=f"{API_ROOT}/api/2.0/serving-endpoints", json=data, headers=headers)

print(json.dumps(response.json(), indent=4))
```

エラーなく実行できれば、モデルサービングのエンドポイントを作成するビルド処理が開始されます。
使用可能になるまでは時間がかかるため、モデルサービングのUIを見ながら完了を待ちましょう。

# Step2. AI Playgroundで利用する

機械学習カテゴリからPlaygroundを実行します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/fb9d68e1-e350-50de-16ed-7bf86c899227.png)

左上のモデル選択ドロップダウンから、先ほど作成したエンドポイントを設定。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/8d8fad7e-c01c-488a-83c0-3151bf93be23.png)

全体の画面構成は以下のようになります。
また、日本語特化モデルなので、システムプロンプトに日本語で回答するように指定しました。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/e87ef585-7f25-5ceb-8403-4bbdd4d037fb.png)

では、いくつか動かしてみます。

## デフォルトプロンプトの「要約」

Playgroundのサンプルとして組み込まれている「要約」を実行。

![llm_chat3a.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/4ebb9b86-ad50-703f-4714-9730d7a359e9.gif)

要約というより翻訳っぽくなってしまいましたが、チャットボットのように回答を出してくれています。
速度は30 tokens/secでした。

## オリジナルの問い合わせ

大阪の観光名所は？

![llm_chat3b.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/6bb75486-309c-36c1-b42c-8ceddded59eb.gif)

## おまけ

界隈でのデファクトとなっているチェック。

![llm_chat3c.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/26305012-d075-30d7-5e7c-69e9aa0a8c6d.gif)

<br>

というわけで、少しだけですが試してみました。

画像をつけていませんが、temperatureなどの設定や複数モデルでの応答比較など、ベーシックな機能は提供されています。

# まとめ

AI Playgroundを試してみました。
モデルをチャット形式で試せる仕組みは既に多くのものが出ていますが、Production Readyでサービングされたモデルを手軽に動かせるものというのはまだ少ないのではないかと思います。

「プロビジョニングされたスループットの基盤モデル」に対応した基盤モデルはそれほどバリエーションが多いわけではないのですが、LlamaやMistral/Mixtral、DBRXなど主要なものはカバーされています。
これらのカスタムモデルなどを手軽に確認するのには良い環境だと思います。

Swallow-MS-7B-Instructの性能は今回さわったぐらいではわかりませんので、これちらも継続的に確認したいですね。