Demo Notebook for deploying CLIPTextModel to OpenSearch

Download notebook

Related Docs: * OpenSearch ML Framework * Huggingface - CLIP * Huggingface - export to torchscript

This notebook provides a walkthrough for users to trace, register, and deploy a CLIPTextModel from a local file. CLIPTextModel can be used with the Neural Search plugin to generate embeddings of documents and ingest time and of user queries at search time.

Step 0: Import packages and set up client

Step 1: Trace CLIPTextModel and export to TorchScript

Step 2: Prep files for registration

Step 3: Register model to OpenSearch

Step 4: Deploy model

Step 0: Import packages and set up client

from transformers import CLIPProcessor, CLIPTextModel
import torch

import opensearch_py_ml as oml
from opensearch_py_ml.ml_commons import MLCommonClient
from opensearchpy import OpenSearch

import warnings
warnings.filterwarnings("ignore", message="Unverified HTTPS request")

# Connect to OpenSearch cluster
host = 'localhost'
port = 9200
auth = ('admin', '< admin password >') # For testing only. Don't store credentials in code.

def get_os_client(host = host, port = port, auth = auth):
    '''
    Get OpenSearch client
    :param cluster_url: cluster URL like https://ml-te-netwo-1s12ba42br23v-ff1736fa7db98ff2.elb.us-west-2.amazonaws.com:443
    :return: OpenSearch client
    '''
    client = OpenSearch(
        hosts = [{'host': host, 'port': port}],
        http_compress = True, # enables gzip compression for request bodies
        http_auth = auth,
        use_ssl = False,
        verify_certs = False,
        ssl_assert_hostname = False,
        ssl_show_warn = False,
    )
    return client

client = get_os_client()
ml_client = MLCommonClient(client)

Step 1: Trace CLIPTextModel and export to TorchScript

To use a model in OpenSearch, you’ll need to export the model into a portable format. As of Version 2.5, OpenSearch only supports the TorchScript and ONNX formats.

Exporting a model to TorchScript requires two things:

• model instantiation with the torchscript flag

• a forward pass with dummy inputs

The dummy inputs are used for a model’s forward pass. While the inputs’ values are propagated through the layers, PyTorch keeps track of the different operations executed on each tensor. These recorded operations are then used to create the trace of the model.

As of OpenSearch 2.6, the ML Framework supports text-embedding models only. CLIP is multi-modal, but we will use CLIPTextModel only here.

model_name = "openai/clip-vit-base-patch32" #See https://huggingface.co/models for other options
text_to_encode = "example search query" #See https://huggingface.co/docs/transformers/torchscript for more info on dummy inputs

# Instantiate CLIPTextModel and CLIPProcessor with pretrained weights
model = CLIPTextModel.from_pretrained(model_name, torchscript=True, return_dict=False)
processor = CLIPProcessor.from_pretrained(model_name)

# Use processor to generate tensors and create dummy input
text_inputs =processor(text=text_to_encode, return_tensors="pt",max_length=77, padding="max_length", truncation=True)
dummy_input = [text_inputs['input_ids'], text_inputs['attention_mask']]

# Trace model and convert to torchscript object
traced_model = torch.jit.trace(model, dummy_input)

# Save model in portable format
torch.jit.save(traced_model, "traced_model_example.pt")

Step 2: Prep files for registration

OpenSearch requires two files zipped together for registration: * Model in TorchScript format * tokenizor.json file

The tokenizor for the model used in this example can be found here

Additionally, a config.json file with the following details must be passed with the .zip. More info on model config

# config.json sample contents
"""
{
    "name": "clip-vit-base-patch32",
    "version": '1.0.0',
    "model_format": "TORCH_SCRIPT",
    "model_config": {
        "model_type": "clip",
        "embedding_dimension": 512,
        "framework_type": "huggingface_transformers"
    }
}
""";

Step 3: Register model to OpenSearch

• Model name in config.json should match .pt torchscript file name

• Record the model ID from the output of the next cell

model_path = "<your_path>/traced_model_example.zip"
model_config_path = "<your_path>/config.json"

model_id_file_system = ml_client.register_model(model_path, model_config_path, isVerbose=True, deploy_model = False)

Total number of chunks 19
Sha1 value of the model file:  62f4786ef2d546180dbbaf8fe6b5be218243c8b806e6623840b1fe9d11bcad4a
Model meta data was created successfully. Model Id:  -uy7rooBhmcN7ynH0lgK
uploading chunk 1 of 19
Model id: {'status': 'Uploaded'}
uploading chunk 2 of 19
Model id: {'status': 'Uploaded'}
uploading chunk 3 of 19
Model id: {'status': 'Uploaded'}
uploading chunk 4 of 19
Model id: {'status': 'Uploaded'}
uploading chunk 5 of 19
Model id: {'status': 'Uploaded'}
uploading chunk 6 of 19
Model id: {'status': 'Uploaded'}
uploading chunk 7 of 19
Model id: {'status': 'Uploaded'}
uploading chunk 8 of 19
Model id: {'status': 'Uploaded'}
uploading chunk 9 of 19
Model id: {'status': 'Uploaded'}
uploading chunk 10 of 19
Model id: {'status': 'Uploaded'}
uploading chunk 11 of 19
Model id: {'status': 'Uploaded'}
uploading chunk 12 of 19
Model id: {'status': 'Uploaded'}
uploading chunk 13 of 19
Model id: {'status': 'Uploaded'}
uploading chunk 14 of 19
Model id: {'status': 'Uploaded'}
uploading chunk 15 of 19
Model id: {'status': 'Uploaded'}
uploading chunk 16 of 19
Model id: {'status': 'Uploaded'}
uploading chunk 17 of 19
Model id: {'status': 'Uploaded'}
uploading chunk 18 of 19
Model id: {'status': 'Uploaded'}
uploading chunk 19 of 19
Model id: {'status': 'Uploaded'}
Model registered successfully

Step 4: Deploy model

model_id = '-uy7rooBhmcN7ynH0lgK' #your model ID from previous step
ml_client.deploy_model(model_id)

task_id: --y8rooBhmcN7ynHYFg6
Model deployed successfully

{'model_id': '-uy7rooBhmcN7ynH0lgK',
 'task_type': 'DEPLOY_MODEL',
 'function_name': 'TEXT_EMBEDDING',
 'state': 'COMPLETED',
 'worker_node': ['4K6CeIPPTkKiwZMplvJ6CQ'],
 'create_time': 1695148695608,
 'last_update_time': 1695148703362,
 'is_async': True}

# Check model status
ml_client.get_model_info(model_id)

{'name': 'traced_model_example3',
 'model_group_id': '-ey7rooBhmcN7ynH0Vji',
 'algorithm': 'TEXT_EMBEDDING',
 'model_version': '1',
 'model_format': 'TORCH_SCRIPT',
 'model_state': 'DEPLOYED',
 'model_content_size_in_bytes': 186945250,
 'model_content_hash_value': '62f4786ef2d546180dbbaf8fe6b5be218243c8b806e6623840b1fe9d11bcad4a',
 'model_config': {'model_type': 'clip',
  'embedding_dimension': 512,
  'framework_type': 'HUGGINGFACE_TRANSFORMERS'},
 'created_time': 1695148659209,
 'last_updated_time': 1695148703362,
 'last_deployed_time': 1695148703362,
 'total_chunks': 19,
 'planning_worker_node_count': 1,
 'current_worker_node_count': 1,
 'planning_worker_nodes': ['4K6CeIPPTkKiwZMplvJ6CQ'],
 'deploy_to_all_nodes': True}