Databricks Model ServingでBGE-M3モデルを利用する

# 導入

Databricks上でなるべく簡単に日本語利用可能かつ本番利用も見据えたEmbedding用のAPIサーバを用意したいなと思っていたら、Marketplaceに[BGE-M3](https://huggingface.co/BAAI/bge-m3)モデルが公開されていることに気づきました。

BGE-M3については、以前記事を書いているので参考まで。

https://qiita.com/isanakamishiro2/items/e4f67586b4cb5f171ea9

BGE-M3自体はHuggingFace上でモデルが公開されており、利用自体はそんなに難しくないのですが、より手軽かつ本番利用想定での準備となるとなかなか難しいです。

Marketplaceでは[Model Serving](https://docs.databricks.com/en/machine-learning/model-serving/index.html)にこのモデルの公開をするところまでのサンプルノートブックも付いており、こちらを利用してEmbedding用サーバを準備・利用してみます。

:::note info
ちなみに、Databricksでは[Foundation Model API(pay-per-token)](https://docs.databricks.com/en/machine-learning/foundation-models/index.html)といういくつかのLLM基盤モデルをAPIとして公開しており、その中にEmbeddingモデルもあるのですが、公開されているモデルがBGE-Large(EN)という英語用モデルであること、またそもそもFoundation Model APIは[2024/3月時点で日本リージョンにはまだ公開されていない](https://docs.databricks.com/en/machine-learning/model-serving/model-serving-limits.html#region-availability)ため、現状は自分でAPIサーバを立てるのがよさそうです。
:::

# Step1. Marketplaceからモデルを取得

Databricks Marketplace上にBGE-M3 modelという名前で公開されていますので、こちらのアクセス権を取得してください。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/3c24a564-564f-9215-26d1-1bb45de8a179.png)

デフォルトでは、```databricks_bge_m3_model```という名前の共有カタログが作られ、```models```スキーマにモデルが格納されています。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/921b9f4e-089a-645f-030e-c40d56115cd3.png)

また、Marketplace上にこのモデルを利用するためのノートブックも共有公開されています。
以後は主にそのノートブックの内容をなぞる形でModel Servingに登録します。

# Step2. モデルをModel Servingに登録

取得したモデルをModel Serving機能を使ってAPIサーバとして公開します。
UI上からでも実行できるのですが、今回はMarketpalceの指示に基づいて、ノートブック上に登録処理を記述・実行します。

まずは```databricsk-sdk```パッケージをアップグレード。

```python
# Upgrade to use the newest Databricks SDK
%pip install --upgrade databricks-sdk
dbutils.library.restartPython()
```

Model Serving機能で構築するエンドポイントの設定情報を変数に格納。

```python
catalog_name = "databricks_bge_m3_model"
version = "2"
model_name = "bge_m3"
model_uc_path = f"{catalog_name}.models.{model_name}"
endpoint_name = f'{model_name}_endpoint'

# エンドポイントのコンピュートタイプ。CPUかGPU各種から選択。今回はテスト的にCPUを選択。
workload_type = "CPU"
```

上記で設定した内容を基に、エンドポイントの作成を実行。
CPUインスタンスの場合だと、数分程度で完了します。

```python
import datetime

from databricks.sdk import WorkspaceClient
from databricks.sdk.service.serving import EndpointCoreConfigInput
w = WorkspaceClient()

config = EndpointCoreConfigInput.from_dict({
    "served_models": [
        {
            "name": endpoint_name,
            "model_name": model_uc_path,
            "model_version": version,
            "workload_type": workload_type,
            "workload_size": "Small",
            "scale_to_zero_enabled": "True",
        }
    ]
})
model_details = w.serving_endpoints.create(name=endpoint_name, config=config)
model_details.result(timeout=datetime.timedelta(minutes=40))
```

無事に完了すると、サービングメニューから該当のエンドポイントが作成できていることを確認できます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/d5dedaef-2390-71b5-980f-e2ed8815cd83.png)


# Step3. テスト

テスト用に、作成したエンドポイントを埋め込み処理を実行させてみます。

```python
from databricks.sdk import WorkspaceClient
from pprint import pprint

dataframe_records = [
    {
        "input": "Databricksは、クラウド上の統合分析プラットフォームです。そのコンセプトは「データとAIの民主化」であり、データ統合、データ分析、AI活用をすべて行えます"
    }
]

w = WorkspaceClient()
result = w.serving_endpoints.query(
    name=endpoint_name,
    dataframe_records=dataframe_records,
)

pprint(result)
```

```:出力
QueryEndpointResponse(choices=[],
                      created=None,
                      data=[],
                      id=None,
                      model=None,
                      object=None,
                      predictions={'data': [{'embedding': [0.0043464601039886475,
                                                           -0.022760512307286263,
                                                           -0.039256516844034195,
                                                           -0.021614814177155495,
（省略）
                                                           0.017395777627825737,
                                                           -0.026914486661553383],
                                             'index': 0,
                                             'object': 'embedding'}],
                                   'object': 'list',
                                   'usage': {'prompt_tokens': 46,
                                             'total_tokens': 46}},
                      served_model_name='bge_m3_endpoint',
                      usage=None)
```

正しく埋め込み結果が得られているように見えます。
このモデル、トークン数も出力するようになっているのが良いですね。

# Step4. LangChainとの連携

実際にはLangChainやLlamaIndexなど、他のパッケージから利用することが多いと思います。
LangChainではDatabricks Model Serving機能で構築されたエンドポイントと連携するモジュールが公開されており、そちらの機能を利用して埋め込み結果を取得してみます。

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

```python
%pip install -U langchain
dbutils.library.restartPython()
```

LangChainの[Databricksモジュール](https://python.langchain.com/docs/integrations/providers/databricks)を利用して埋め込み処理を実行。
同一ワークスペースのエンドポイントを利用する場合、(権限があれば)エンドポイント名だけの指定で利用できる模様。

```python
from langchain_community.embeddings import DatabricksEmbeddings

model_name = "bge_m3"
endpoint_name = f"{model_name}_endpoint"

embeddings = DatabricksEmbeddings(endpoint=endpoint_name)
embeddings.embed_query(
    "Databricksは、クラウド上の統合分析プラットフォームです。そのコンセプトは「データとAIの民主化」であり、データ統合、データ分析、AI活用をすべて行えます"
)
```

```:出力
[0.0043464601039886475,
 -0.022760512307286263,
 -0.039256516844034195,
 -0.021614814177155495,
（省略）
 -0.05037514865398407,
 0.03156855329871178,
 ...]
```

無事に埋め込み結果を取得できました。

# まとめ

Databricks MarketplaceとModel Servingの利用によって簡単に日本語対応のEmbeddingモデルAPIサーバ(エンドポイント)を作成・公開できました。手軽ですし、いわゆるProduction Readyの状態で構築できるのはありがたいです。

ちなみに、BGE-M3を[プロビジョニングされたスループット モデルAPI](https://docs.databricks.com/ja/machine-learning/foundation-models/deploy-prov-throughput-foundation-model-apis.html)としてエンドポイントを構築することはできませんでした。BGE-Largeはできるようなのですが、M3はアーキテクチャが異なるからNGなのかな。

LLM等の生成AI利用における本番化はどんどん進んでいくと思いますが、セキュリティポリシーなどで外部APIを手軽に使えない企業も多くあるのではないかと考えています。
こういった際に、Databricks上で自らエンドポイントを構築・利用できる環境ニーズは何となく増えるのではないかと思います。