Registrieren und Bedienen eines OSS-Einbettungsmodells

Notizbuchversion dieser Seite öffnen

In diesem Notizbuch wird das Open Source Texteinbettungsmodell e5-small-v2 in einem Modellbereitstellungsendpunkt eingerichtet, der für die Vektorsuche verwendet werden kann.

  • Laden Sie das Modell aus dem Hugging Face Hub herunter.
  • Registrieren Sie sie bei der MLflow-Modellregistrierung.
  • Starten Sie einen Modellbereitstellungsendpunkt, um das Modell zu bedienen.

Das Modell e5-small-v2 ist unter https://huggingface.co/intfloat/e5-small-v2.

Eine Liste der Bibliotheksversionen, die in Databricks Runtime enthalten sind, finden Sie in den Versionshinweisen für Ihre Databricks-Runtime-Version.

Installieren von Databricks Python SDK

Dieses Notizbuch verwendet seinen Python-Client, um mit Dienstendpunkten zu arbeiten.

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

Herunterladen des Modells

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

Registrieren des Modells für 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

Erstellen eines Modellbereitstellungsendpunkts

Weitere Informationen finden Sie unter Endpunkte für Servicemodelle erstellen.

Hinweis: In diesem Beispiel wird ein kleiner CPU-Endpunkt erstellt, der auf 0 skaliert wird. Dies ist für schnelle, kleine Tests vorgesehen. Für realistischere Anwendungsfälle sollten Sie die Verwendung von GPU-Endpunkten für eine schnellere Einbettung der Berechnung in Betracht ziehen und nicht auf 0 skalieren, wenn Sie häufige Abfragen erwarten, da Modellbereitstellungsendpunkte einen Kaltstartaufwand haben.

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)

Abfrageendpunkt

Der obige create_and_wait Befehl wartet, bis der Endpunkt bereit ist. Sie können auch den Status des dienenden Endpunkts in der Databricks-Benutzeroberfläche überprüfen.

Weitere Informationen finden Sie unter Query Foundation-Modelle.

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

Beispiel-Notebook

Registrieren und Bedienen eines OSS-Einbettungsmodells

Notebook abrufen

  • Last updated on