Mendaftar dan melayani model penyematan OSS

Notebook ini menyiapkan model penyematan teks sumber terbuka e5-small-v2 dalam endpoint Servis Model yang bisa digunakan untuk Pencarian Vektor.

  • Unduh model dari Hugging Face Hub.
  • Daftarkan ke Registri Model MLflow.
  • Mulai endpoint Penyajian Model untuk menyajikan model.

Model e5-small-v2 ini tersedia di https://huggingface.co/intfloat/e5-small-v2.

Untuk daftar versi pustaka yang disertakan dalam Databricks Runtime, lihat catatan rilis untuk versi Databricks Runtime Anda.

Menginstal Databricks Python SDK

Notebook ini menggunakan klien Python untuk bekerja dengan melayani titik akhir.

%pip install -U databricks-sdk python-snappy
%pip install sentence-transformers
dbutils.library.restartPython()

Mengunduh model

# 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)

Mendaftarkan model ke 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

Membuat titik akhir penyajian model

Untuk detail selengkapnya, lihat Membuat model fondasi yang melayani titik akhir.

Catatan: Contoh ini membuat titik akhir CPU kecil yang menurunkan skala ke 0. Ini untuk tes cepat dan kecil. Untuk kasus penggunaan yang lebih realistis, pertimbangkan untuk menggunakan titik akhir GPU untuk penghitungan embedding yang lebih cepat dan tidak mengurangi skala hingga 0 jika Anda mengharapkan kueri yang sering, karena titik akhir Model Serving memiliki beberapa overhead cold start.

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)

Titik akhir kueri

Perintah di atas create_and_wait menunggu hingga titik akhir siap. Anda juga dapat memeriksa status titik akhir penyajian di antarmuka pengguna Databricks.

Untuk informasi selengkapnya, lihat Model fondasi kueri.

# 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}')

Contoh buku catatan

Mendaftarkan dan menerapkan model penyematan OSS

Dapatkan buku catatan