Inscrire et servir un modèle d’incorporation OSS

Ce notebook configure le modèle d’incorporation de texte open source e5-small-v2 dans un point de terminaison Model Serving utilisable pour la recherche vectorielle.

  • Téléchargez le modèle à partir du Hub Hugging Face.
  • Inscrivez-le dans le Registre du modèle MLflow.
  • Démarrez un point de terminaison de service de modèle afin de servir le modèle.

Le modèle e5-small-v2 est disponible à l’adresse https://huggingface.co/intfloat/e5-small-v2.

Pour obtenir la liste des versions de bibliothèque incluses dans Databricks Runtime, consultez les notes de publication de votre version de Databricks Runtime.

Installer le Kit de développement logiciel (SDK) Databricks Python

Ce notebook utilise son client Python pour travailler avec les points de terminaison de service.

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

Télécharger un modèle

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

Inscrire un modèle dans 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

Créer un point de terminaison pour le service de modèle

Pour plus d’informations, consultez Créer un modèle de base servant des points de terminaison.

Remarque : Cet exemple crée un petit point de terminaison d’UC qui est réduit à 0. C’est pour des tests rapides et petits. Pour des cas d’usage plus réalistes, envisagez d’utiliser des points de terminaison GPU pour accélérer le calcul d’incorporation et ne pas effectuer un scale-down à 0 si vous attendez des requêtes fréquentes, car les points de terminaison Model Service ont une surcharge de démarrage à froid.

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)

Point de terminaison de requête

La commande ci-dessus create_and_wait attend que le point de terminaison soit prêt. Vous pouvez également vérifier l’état du point de terminaison de service dans l’interface utilisateur Databricks.

Pour plus d’informations, consultez Modèles de base de requête.

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

Exemple de bloc-notes

Enregistrer et déployer un modèle d'intégration OSS

Obtenir un ordinateur portable