Exploring hybrid search and custom models in OpenSearch on TrueFoundry (Part II)

Introduction

In the previous part of this article, we deployed a small end-to-end application that utilizes one of the examples of the new OpenSearch hybrid search feature. This time, we'll replicate the same example with an important difference: we are hosting a copy of the same model from Hugging Face on a TrueFoundry cluster.

The idea is to set up the scenario to train a custom model using TrueFoundry training capabilities while keeping it plugged-in to OpenSearch. This whole platform will eventually allow to use private data to fine tune hybrid search on premise without having to send the data outside an organization's boundary.

Part II: a second iteration

For the second iteration, we are going to repeat the installation of a single OpenSearch node just as we did in the first part. The approach will be basically the same. However, before jumping to the Jupyter notebook, we will register a Hugging Face model on TrueFoundry.

Hosting a single node of OpenSearch

Create a new workspace (or reuse the previous one). We'll create a different one in this case named "hybrid-search-2". Be aware that this change may affect the URLs shown below.

Workspace creation wizard.

Go ahead an push "+New Deployment" as we did in the previous entry. Create a node for the Docker image "opensearchproject/opensearch:2.15.0". Remember to remove the command override, change the port, and adjust the resources, just as we did before. Finally, add the environment variables "OPENSEARCH_INITIAL_ADMIN_PASSWORD" and “discovery.type=single-node” exactly in the same way.

NOTE: by the time we wrote this article OpenSearch version 2.15.0 wasn't available yet. This release candidate can be used instead: opensearchstaging/opensearch:2.15.0.9964 Check out the OpenSearch release schedule for further information.

Registering a Hugging Face model on TrueFoundry

Go to Deployments dashboard and click “+New Deployment”, but this time, let’s click on “Model” on the left, instead of “Service”.

The wizard is different, to ease deployment of existing models. Select your workspace and enter “https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2" in the “Enter HuggingFace URI of a Model to deploy”.

Deployment dashboard to deploy a HuggingFace model.

Hit “Next”, and choose “CPU (TEI)” so we don incur in GPU costs in the TrueFoundry platform. This can be changed later by editing the deployment.

GPU/CPU selection for your model.

And you will be presented with the wizard of the Docker image corresponding to the model. Everything should be correct and automatically filled in. In our case we will uncheck the “Expose” checkbox to expose the model, as we only will access it through OpenSearch.

Click “Submit” and the deployment should start. Copy the endpoint as we’re gonna need it for our next move.

Testing the model from a notebook

Deploy a base image of a Jupyter notebook kernel if you haven't done so already, as we did in the previous blogpost if you haven't done it yet. Then, create a cell with this content:

!curl -X POST http://all-minilm-l6-v2.hybrid-search-2.svc.cluster.local:8000/ \
 -d '["check this embedding"]' \
 -H 'Content-Type: application/json'

If everything went okay, you should see the embedding, a long JSON array of floating-point numbers, similar to this one:

[[0.0074966755,-0.05388556,-0.016110366,-0.06457912,0.0745039,0.0070791887,
 /* ... skipped ... */ 0.022871481 ]]

Great! Let's move to the next phase.

Registering the model in OpenSearch

Let’s copy over the first two cells of the notebook from the previous blogpost entry as we’re going to use them again:

import getpass
opensearch_password = getpass.getpass("opensearch_password:")

Run it using the same password we were using for OPENSEARCH_INITIAL_ADMIN_PASSWORD.

# Disable HTTPS verification warnings (not for production).
import urllib3
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

###

import requests
import json

from requests.auth import HTTPBasicAuth

endpoint = "https://opensearch.hybrid-search-2.svc.cluster.local:9200"
auth = HTTPBasicAuth("admin", opensearch_password)

# A function that calls an HTTP request with an optional json body,
# prints the json output and returns it.
def request(method, uri, body=None):
    r = requests.request(method=method, url=f"{endpoint}{uri}", auth=auth, verify=False, json=body)
    r.raise_for_status()
    result = r.json()
    print(json.dumps(result, indent=2))
    return result

request("GET", "/");

Beware of the URL, it could have changed if you changed the name of the workspace.

Now. In order to be able to use the external model, we need to tell OpenSearch that we need to allow external HTTP calls. The way to do this is by modifying the cluster settings providing a more permissive value for "trust_connector_endpoint_regex" than the default. In our next cell, put:

request("PUT", "/_cluster/settings", {
  "persistent": {
    "plugins": {
      "ml_commons": {
        "only_run_on_ml_node": "false",
        "model_access_control_enabled": "true",
        "connector_access_control_enabled": "true",
        "native_memory_threshold": "99",
        # You might want to provide a narrower URL regex here 
        # in a prod environment.
        "trusted_connector_endpoints_regex": [
          "^.*$"
        ],
        "connector.private_ip_enabled": "true"
      }
    }
  }
});

Next, create a connector for the model:

body = request(
    "POST",
    "/_plugins/_ml/connectors/_create",
    {
        "name": "Custom local model",
        "description": "The connector to a locally hosted model",
        "version": "1",
        "protocol": "http",
        "parameters": {"model": "just a placeholder", "input_type": "text"},
        "actions": [
            {
                "action_type": "predict",
                "method": "POST",
                "url": "http://all-minilm-l6-v2.hybrid-search-2.svc.cluster.local:8000",
                "request_body": '["${parameters.inputText}"]',
                "pre_process_function": '\n    StringBuilder builder = new StringBuilder();\n    builder.append("\\"");\n    String first = params.text_docs[0];\n    builder.append(first);\n    builder.append("\\"");\n    def parameters = "{" +"\\"inputText\\":" + builder + "}";\n    return  "{" +"\\"parameters\\":" + parameters + "}";',
                "post_process_function": "connector.post_process.default.embedding",
            }
        ],
    },
)
connector_id = body['connector_id']
connector_id

Register the connector id in the model group as with any new model. Wait for the task to be completed and also deploy the model just as before. But first, let's create a model group.

model_group_name = "my_model_group"

body = request("GET", "/_plugins/_ml/model_groups/_search", {
  "query": {
    "bool": {
      "must": [
        {
          "terms": {
            "name": [model_group_name]
          }
        }
      ]
    }
  }
})

model_group_id = None
model_group_docs = body.get('hits', {}).get('hits', {})
if model_group_docs:
    model_group_id = model_group_docs[0]['_id']

if not model_group_id:
    body = request("POST", "/_plugins/_ml/model_groups/_register", {
      "name": model_group_name,
      "description": f"A model group named {model_group_name}",
      "access_mode": "public"
    })
    model_group_id = body['model_group_id']

model_group_id

Then register the model with the new connector:

body = request("POST", "/_plugins/_ml/models/_register", {
  "name": "Our very own all-MiniLM-L6-v2",
  "function_name": "remote",
  "model_group_id": model_group_id,
  "connector_id": connector_id
})

task_id = body['task_id']

# Write this as a function, because we're going to use it a few times.
def await_for_task(task_id):
    import time

    print(f"Awaiting for the task {task_id} to finish", end='')

    model_id = None
    while True:
        r = requests.get(f"{endpoint}/_plugins/_ml/tasks/{task_id}", auth=auth, verify=False)
        r.raise_for_status()

        body = r.json()
        state = body['state']
        if state == "COMPLETED":
            print("Done!")
            return body['model_id']

        print('.', end='')
        time.sleep(1)

model_id = await_for_task(task_id)
model_id

Deploy the model:

body = request("POST", f"/_plugins/_ml/models/{model_id}/_deploy")
task_id = body['task_id']

model_id = await_for_task(task_id)

Test the model:

request("POST", f"/_plugins/_ml/_predict/text_embedding/{model_id}", {
  "text_docs": [ "best embedding ever"],
  "return_number": False,
  "target_response": ["text_embedding"]
});

If you got a valid JSON response like this one, we are golden!

{
  "inference_results": [
    {
      "output": [
        {
          "name": "sentence_embedding",
          "data_type": "FLOAT32",
          "shape": [
            384
          ],
          "data": [
            -0.048571117,
            -0.06612981,
            0.035486978,
            -0.07163426,
            // ... skipped ...
            0.020217668
          ]
        }
      ],
      "status_code": 200
    }
  ]
}

So now everything is just like before. Like for any other internal model. So let's quickly go over the rest.

Create or update the ingest pipeline so it points to the new model.

ingest_pipeline = "my-ingest-pipeline"

request("PUT", f"/_ingest/pipeline/{ingest_pipeline}", {
  "processors": [
    {
      "text_embedding": {
        "model_id": model_id,
        "field_map": {
          "text": "passage_embedding"
        }
      }
    }
  ]
});

Reindex everything, recreating the index first.

index_name = "my-index-name"

# Delete the index if already exists, so the notebook can be replayed.
r = requests.get(f"{endpoint}/{index_name}", auth=auth, verify=False)
if r.status_code != 404:
    r = requests.delete(f"{endpoint}/{index_name}", auth=auth, verify=False)
    r.raise_for_status()

request("PUT", f"/{index_name}", {
  "settings": {
    "index.knn": True,
    "default_pipeline": ingest_pipeline
  },
  "mappings": {
    "properties": {
      "id": {
        "type": "text"
      },
      "passage_embedding": {
        "type": "knn_vector",
        "dimension": 384,
        "method": {
          "engine": "lucene",
          "space_type": "l2",
          "name": "hnsw",
          "parameters": {}
        }
      },
      "text": {
        "type": "text"
      }
    }
  }
});

Then the documents. It may take longer this time, because of access to a different host:

docs = [{
  "text": "A West Virginia university women 's basketball team , officials , and a small gathering of fans are in a West Virginia arena .",
  "id": "4319130149.jpg"
}, {
  "text": "A wild animal races across an uncut field with a minimal amount of trees .",
  "id": "1775029934.jpg"
}, {
  "text": "People line the stands which advertise Freemont 's orthopedics , a cowboy rides a light brown bucking bronco .",
  "id": "2664027527.jpg"
}, {
  "text": "A man who is riding a wild horse in the rodeo is very near to falling off .",
  "id": "4427058951.jpg"
}, {
  "text": "A rodeo cowboy , wearing a cowboy hat , is being thrown off of a wild white horse .",
  "id": "2691147709.jpg"
}]

for doc in docs:
    request("PUT", f"/{index_name}/_doc/{doc['id']}", doc)

Create or update the search template to use the new model_id.

search_template = "my-search-template"

request("POST", f"/_scripts/{search_template}", {
  "script": {
    "lang": "mustache",
    "source": {
      "from": "{{from}}{{^from}}0{{/from}}",
      "size": "{{size}}{{^size}}10{{/size}}",
      "_source": {
        "exclude": [
          "passage_embedding"
        ]
      },
      "query": {
        "hybrid": {
          "queries": [
            {
              "match": {
                "passage_text": {
                  "query": "{{query}}"
                }
              }
            },
            {
              "neural": {
                "passage_embedding": {
                  "query_text": "{{query}}",
                  "model_id": "{{model_id}}",
                  "k": 5
                }
              }
            }
          ]
        }
      },
      "search_pipeline": {
        "phase_results_processors": [
          {
            "normalization-processor": {
              "normalization": {
                "technique": "min_max"
              },
              "combination": {
                "technique": "arithmetic_mean",
                "parameters": {
                  "weights": [
                    0.3,
                    0.7
                  ]
                }
              }
            }
          }
        ]
      }
    },
    "params": {
      "query": ""
    }
  }
});

If everything went well, you should be able to search like this:

request("GET", f"/{index_name}/_search/template", {
  "id": search_template,
  "params": {
    "query": "sport",
    "model_id": model_id,
    "size": 2
  }
});

And the response should look something similar to:

{
  "took": 305,
  "timed_out": false,
  "_shards": {
    "total": 1,
    "successful": 1,
    "skipped": 0,
    "failed": 0
  },
  "hits": {
    "total": {
      "value": 2,
      "relation": "eq"
    },
    "max_score": 0.7,
    "hits": [
      {
        "_index": "my-index-name",
        "_id": "4319130149.jpg",
        "_score": 0.7,
        "_source": {
          "text": "A West Virginia university women 's basketball team , officials , and a small gathering of fans are in a West Virginia arena .",
          "id": "4319130149.jpg"
        }
      },
      {
        "_index": "my-index-name",
        "_id": "2691147709.jpg",
        "_score": 0.00070000003,
        "_source": {
          "text": "A rodeo cowboy , wearing a cowboy hat , is being thrown off of a wild white horse .",
          "id": "2691147709.jpg"
        }
      }
    ]
  }
}

We are all set. Congratulations!

The rest of the application should be similar to what we have done in the previous blog entry. So we aren't copying it over here.

Notice that the time taken to search will be significantly higher, as the model is being accessed remotely to generate an embedding for the query itself. This is a disadvantage of using an external model. However, we are here for the advantages.

Also, there are a few countermeasure's we can take to reduce the impact of the performance overhead. We are to discuss that on the next section.

Conclusions and Further Work

The advantages of having an external model are plenty. Not only we are now enabled to fine tune the model, but we could use another model entirely, with different weights, with different architecture, etc.

Check out one of the models for feature extraction in the TrueFoundry model catalogue. They're ready to deploy.

Because of the impact of accessing the an external service is there, there are several countermeasures we could take to improve the model inference performance: move the model to GPU and incrementally scale out the GPU as needed, horizontally as well as vertically.

At each step, it will be useful to rely on benchmarking the performance of the system in order to characterize the impact of every experiment. There are tools to benchmark the performance on TrueFoundry.

Additionally, you could consider adding a cache layer between the model and the consumer, the OpenSearch cluster. This would be useful for performance, but also useful to reduce GPU-associated costs.

Whilst a hello world application like these ones don’t look like much, they help establishing a skeleton of the most significant interactions of a system’s components, so everybody can handle them intellectually and iterating over becomes a coordinated and understandable task all the time.

The problem of generating embeddings for a collection of documents has very often noticeable complexity depending on the externally services used. The OpenSearch proposed architecture for the problem makes this problem trivial.

As stated in the previous blog entry, the ability to fine tune a model could be a game changer for an organization.

Fine tune in part III?