Registro e implementación de un modelo de incrustación de OSS

Abrir la versión 'notebook' de esta página

Este cuaderno configura el modelo de inserción de texto código abierto e5-small-v2 en un punto de conexión de servicio de modelos que se puede usar para la búsqueda de vectores.

  • Descargue el modelo del Hub de Hugging Face.
  • Regístrelo en el Registro de modelos de MLflow.
  • Inicie un endpoint de servicio de modelos para servir el modelo.

El modelo e5-small-v2 está disponible en https://huggingface.co/intfloat/e5-small-v2.

Para ver una lista de las versiones de las bibliotecas incluidas en Databricks Runtime, consulte las notas de la release para su versión de Databricks Runtime.

Instalación del SDK de Python de Databricks

Este cuaderno usa su cliente Python para trabajar con puntos de servicio.

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

Descarga del modelo

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

Registro del modelo en 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

Creación de un punto de conexión de servicio de modelos

Para más información, consulte Creación de un modelo de base que atiende puntos de conexión.

Nota: En este ejemplo se crea un punto de conexión de CPU pequeño que se reduce verticalmente a 0. Esto es para pruebas rápidas y pequeñas. Para casos de uso más realistas, considere usar endpoints de GPU para un cálculo de embedding más rápido y no escalarla a 0 si espera consultas frecuentes, ya que los endpoints de servición de modelos tienen cierta sobrecarga de arranque en frío.

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)

Punto de conexión de consulta

El comando anterior create_and_wait espera hasta que el punto de conexión esté listo. También puede verificar el estado del punto final de servicio en la interfaz de usuario de Databricks.

Para obtener más información, consulte Modelos de base de consultas.

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

Cuaderno de ejemplo

Registro y puesta en servicio de un modelo de incrustación de software de código abierto

Obtener el portátil

  • Last updated on