このノートブックは、Vector Search で使用できるモデル サービス エンドポイントで、オープンソーステキスト埋め込みモデル e5-small-v2 を設定します。
- Hugging Face Hub からモデルをダウンロードします。
- MLflow モデル レジストリに登録します。
- モデルサービスエンドポイントを開始してモデルを提供します。
モデル e5-small-v2 は、 https://huggingface.co/intfloat/e5-small-v2で使用できます。
- MIT ライセンス
- バリアント：
Databricks Runtime に含まれるライブラリ バージョンの一覧については、Databricks Runtime バージョンの リリース ノート を参照してください。
Databricks Python SDK をインストールする
このノートブックでは、Python クライアントを使用してサービス エンドポイントを操作します。
%pip install -U databricks-sdk python-snappy
%pip install sentence-transformers
dbutils.library.restartPython()
モデルをダウンロードする
# Download model using the sentence_transformers library.
from sentence_transformers import SentenceTransformer
source_model_name = 'intfloat/e5-small-v2' # model name on Hugging Face Hub
model = SentenceTransformer(source_model_name)
# Test the model, just to show it works.
sentences = ["This is an example sentence", "Each sentence is converted"]
embeddings = model.encode(sentences)
print(embeddings)
モデルを MLflow に登録する
import mlflow
mlflow.set_registry_uri("databricks-uc")
# Specify the catalog and schema to use. You must have USE_CATALOG privilege on the catalog and USE_SCHEMA and CREATE_TABLE privileges on the schema.
# Change the catalog and schema here if necessary.
catalog = "main"
schema = "default"
model_name = "e5-small-v2"
# MLflow model name. The Model Registry uses this name for the model.
registered_model_name = f"{catalog}.{schema}.{model_name}"
# Compute input and output schema.
signature = mlflow.models.signature.infer_signature(sentences, embeddings)
print(signature)
model_info = mlflow.sentence_transformers.log_model(
model,
artifact_path="model",
signature=signature,
input_example=sentences,
registered_model_name=registered_model_name)
inference_test = ["I enjoy pies of both apple and cherry.", "I prefer cookies."]
# Load the custom model by providing the URI for where the model was logged.
loaded_model_pyfunc = mlflow.pyfunc.load_model(model_info.model_uri)
# Perform a quick test to ensure that the loaded model generates the correct output.
embeddings_test = loaded_model_pyfunc.predict(inference_test)
embeddings_test
# Extract the version of the model you just registered.
mlflow_client = mlflow.MlflowClient()
def get_latest_model_version(model_name):
client = mlflow_client
model_version_infos = client.search_model_versions("name = '%s'" % model_name)
return max([int(model_version_info.version) for model_version_info in model_version_infos])
model_version = get_latest_model_version(registered_model_name)
model_version
モデルサービングエンドポイントを作成する
詳細については、「エンドポイントを 提供する基盤モデルを作成する」を参照してください。
注: この例では、0 にスケールダウンする 小さな CPU エンドポイントを作成します。 これは、迅速で小規模なテスト用です。 より現実的なユース ケースの場合は、GPU エンドポイントを使用して埋め込み計算を高速化し、頻繁なクエリが必要な場合は 0 にスケールダウンしないことを検討してください。モデル サービス エンドポイントにはコールド スタートオーバーヘッドが発生するためです。
endpoint_name = "e5-small-v2" # Name of endpoint to create
from databricks.sdk import WorkspaceClient
from databricks.sdk.service.serving import EndpointCoreConfigInput
w = WorkspaceClient()
endpoint_config_dict = {
"served_entities": [
{
"name": f'{registered_model_name.replace(".", "_")}_{1}',
"entity_name": registered_model_name,
"entity_version": model_version,
"workload_type": "CPU",
"workload_size": "Small",
"scale_to_zero_enabled": True,
}
]
}
endpoint_config = EndpointCoreConfigInput.from_dict(endpoint_config_dict)
# The endpoint may take several minutes to get ready.
w.serving_endpoints.create_and_wait(name=endpoint_name, config=endpoint_config)
クエリ エンドポイント
上記の create_and_wait コマンドは、エンドポイントの準備ができるまで待機します。 Databricks UI でサービス エンドポイントの状態を確認することもできます。
詳細については、「 クエリ基盤モデル」を参照してください。
# Only run this command after the Model Serving endpoint is in the Ready state.
import time
start = time.time()
# If the endpoint is not yet ready, you might get a timeout error. If so, wait and then rerun the command.
endpoint_response = w.serving_endpoints.query(name=endpoint_name, dataframe_records=['Hello world', 'Good morning'])
end = time.time()
print(endpoint_response)
print(f'Time taken for querying endpoint in seconds: {end-start}')