Aller au contenu principal
Dernière mise à jour le

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

Ouvrir dans Databricks

Ce Notebook configure le modèle d'intégration de texte open source e5-small-v2 dans un Endpoint Model Serving utilisable pour la recherche IA.

  • download le modèle à partir du Hub Hugging Face.
  • Enregistrez-le dans le MLflow Model Registry.
  • Start un Endpoint Model Serving pour diffuser le modèle.

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

Pour une liste des versions de bibliothèques incluses dans Databricks Runtime, consultez les notes de version de votre version de Databricks Runtime.

Installer Databricks Python SDK​

Ce Notebook utilise son client Python pour fonctionner avec des endpoints de service.

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

download model​

Python
# 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)
Python
# Test the model, just to show it works.
sentences = ["This is an example sentence", "Each sentence is converted"]
embeddings = model.encode(sentences)
print(embeddings)

Enregistrer le modèle dans MLflow​

Python
import mlflow
mlflow.set_registry_uri("databricks-uc")
Python

# 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"
Python
# MLflow model name. The Model Registry uses this name for the model.
registered_model_name = f"{catalog}.{schema}.{model_name}"
Python
# Compute input and output schema.
signature = mlflow.models.signature.infer_signature(sentences, embeddings)
print(signature)
Python
model_info = mlflow.sentence_transformers.log_model(
  model,
  artifact_path="model",
  signature=signature,
  input_example=sentences,
  registered_model_name=registered_model_name)
Python
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
Python
# 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 endpoint de service de modèle​

Pour plus de détails, consultez Créer des Endpoint de service de modèle de fondation.

Remarque : Cet exemple crée un petit endpoint de CPU qui se met à l'échelle jusqu'à 0. Ceci est destiné à des tests rapides et de petite taille. Pour des cas d’utilisation plus réalistes, envisagez d’utiliser des endpoints GPU pour un calcul d’intégration plus rapide et de ne pas réduire la mise à l’échelle à 0 si vous prévoyez des requêtes fréquentes, car les endpoints de Model Serving ont un certain temps de cold start.

Python
endpoint_name = "e5-small-v2"  # Name of endpoint to create
Python
from databricks.sdk import WorkspaceClient
from databricks.sdk.service.serving import EndpointCoreConfigInput

w = WorkspaceClient()
Python
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)

Endpoint de la query​

La commande create_and_wait ci-dessus attend que l'Endpoint soit prêt. Vous pouvez également vérifier le statut de l'endpoint de service dans l'interface utilisateur de Databricks.

Pour plus d'informations, voir query foundation models.

Python
# 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 Notebook​

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

Ouvrir le notebook dans un nouvel ongletOuvrir dans Databricks
Sur cette page