導入
Databricks上でなるべく簡単に日本語利用可能かつ本番利用も見据えたEmbedding用のAPIサーバを用意したいなと思っていたら、MarketplaceにBGE-M3モデルが公開されていることに気づきました。
BGE-M3については、以前記事を書いているので参考まで。
BGE-M3自体はHuggingFace上でモデルが公開されており、利用自体はそんなに難しくないのですが、より手軽かつ本番利用想定での準備となるとなかなか難しいです。
MarketplaceではModel Servingにこのモデルの公開をするところまでのサンプルノートブックも付いており、こちらを利用してEmbedding用サーバを準備・利用してみます。
ちなみに、DatabricksではFoundation Model API(pay-per-token)といういくつかのLLM基盤モデルをAPIとして公開しており、その中にEmbeddingモデルもあるのですが、公開されているモデルがBGE-Large(EN)という英語用モデルであること、またそもそもFoundation Model APIは2024/3月時点で日本リージョンにはまだ公開されていないため、現状は自分でAPIサーバを立てるのがよさそうです。
Step1. Marketplaceからモデルを取得
Databricks Marketplace上にBGE-M3 modelという名前で公開されていますので、こちらのアクセス権を取得してください。
デフォルトでは、databricks_bge_m3_modelという名前の共有カタログが作られ、modelsスキーマにモデルが格納されています。
また、Marketplace上にこのモデルを利用するためのノートブックも共有公開されています。
以後は主にそのノートブックの内容をなぞる形でModel Servingに登録します。
Step2. モデルをModel Servingに登録
取得したモデルをModel Serving機能を使ってAPIサーバとして公開します。
UI上からでも実行できるのですが、今回はMarketpalceの指示に基づいて、ノートブック上に登録処理を記述・実行します。
まずはdatabricsk-sdkパッケージをアップグレード。
# Upgrade to use the newest Databricks SDK
%pip install --upgrade databricks-sdk
dbutils.library.restartPython()
Model Serving機能で構築するエンドポイントの設定情報を変数に格納。
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インスタンスの場合だと、数分程度で完了します。
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))
無事に完了すると、サービングメニューから該当のエンドポイントが作成できていることを確認できます。
Step3. テスト
テスト用に、作成したエンドポイントを埋め込み処理を実行させてみます。
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機能で構築されたエンドポイントと連携するモジュールが公開されており、そちらの機能を利用して埋め込み結果を取得してみます。
まずはパッケージをインストール。
%pip install -U langchain
dbutils.library.restartPython()
LangChainのDatabricksモジュールを利用して埋め込み処理を実行。
同一ワークスペースのエンドポイントを利用する場合、(権限があれば)エンドポイント名だけの指定で利用できる模様。
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としてエンドポイントを構築することはできませんでした。BGE-Largeはできるようなのですが、M3はアーキテクチャが異なるからNGなのかな。
LLM等の生成AI利用における本番化はどんどん進んでいくと思いますが、セキュリティポリシーなどで外部APIを手軽に使えない企業も多くあるのではないかと考えています。
こういった際に、Databricks上で自らエンドポイントを構築・利用できる環境ニーズは何となく増えるのではないかと思います。