How to deploy NLP: Text Embeddings and Vector Search

As part of our natural language processing (NLP) blog series, we will walk through an example of using a text embedding model to generate vector representations of textual contents and demonstrating vector similarity search on generated vectors. We will deploy a publicly available model on Elasticsearch and use it in an ingest pipeline to generate embeddings from textual documents. We will then show how to use those embeddings in the vector similarity search to find semantically similar documents for a given query.

Vector similarity search or, as is commonly called semantic search, goes beyond the traditional keyword based search and allows users to find semantically similar documents that may not have any common keywords thus providing a wider range of results. Vector similarity search operates on dense vectors and uses k-nearest neighbour search to find similar vectors. For this, contents in the textual form first need to be converted to their numeric vector representations using a text embedding model.

We will use a public dataset from the MS MARCO Passage Ranking Task for demonstration. It consists of real questions from the Microsoft Bing search engine and human generated answers for them. This dataset is a perfect resource for testing vector similarity search, firstly, because question-answering is a one of the most common use cases for vector search, and secondly, the top papers in the MS MARCO leaderboard use vector search in some form.

In our example we will work with a sample of this dataset, use a model to produce text embeddings, and then run vector search on it. We hope to also do a quick verification of the quality of produced results from the vector search.

1. Deploy a text embedding model.

The first step is to install a text embedding model. For our model we use msmarco-MiniLM-L12-cos-v5 from Hugging Face. This is a sentence-transformer model that takes a sentence or a paragraph and maps it to a 384-dimensional dense vector. This model is optimized for semantic search and was specifically trained on the MS MARCO Passage dataset, making it suitable for our task. Besides this model, Elasticsearch supports a number of other models for text embedding. The full list can be found here.

We install the model with the Eland docker agent that we built in the NER example. Running a script below imports our model into our cluster and deploys it:

eland_import_hub_model 
--cloud-id <cloud-id> \ 
-u <username> -p <password> \ 
--hub-model-id sentence-transformers/msmarco-MiniLM-L12-cos-v5 \ 
--task-type text_embedding \
--start

This time, --task-type is set to text_embedding and the --start option is passed to the Eland script so the model will be deployed automatically without having to start it in the Model Management UI. To speed up inferences, you can increase the number of inference threads with inference_threads parameter.

We can test the successful deployment of the model by using this example in Kibana Console:

POST /_ml/trained_models/sentence-transformers__msmarco-minilm-l12-cos-v5/deployment/_infer
{
  "docs": {
    "text_field": "how is the weather in jamaica"
  }
}

We should see the predicted dense vector as the result:

{
  "predicted_value" : [
    0.051237598061561584,
    -0.04680659621953964,
    0.03971194103360176
    …
  ]
}

2. Loading initial data

As mentioned in the introduction, we use the MS MARCO Passage Ranking dataset. The dataset is quite big, consisting of over 8 million passages. For our example, we use a subset of it that was used in the testing stage of the 2019 TREC Deep Learning Track. The dataset msmarco-passagetest2019-top1000.tsv used for the re-ranking task contains 200 queries and for each query a list of relevant text passages extracted by a simple IR system. From that dataset, we’ve extracted all unique passages with their ids, and put them into a separate tsv file, totaling 182469 passages. We use this file as our dataset.

We use Kibana's file upload feature to upload this dataset. Kibana file upload allows us to provide custom names for fields, let’s call them id with type long for passages’ ids, and text with type text for passages’ contents. The index name is collection. After the upload, we can see an index named collection with 182469 documents.

3. Creating pipeline

We want to process the initial data with an Inference processor that will add an embedding for each passage. For this, we create a text embedding ingest pipeline and then reindex our initial data with this pipeline.

In the Kibana Console we create an ingest pipeline (as we did in the previous blog post), this time for text embeddings, and call it text-embeddings. The passages are in a field named text. As we did before, we’ll define a field_map to map text to the field text_field that the model expects. Similarly on_failure handler is set to index failures into a different index:

PUT _ingest/pipeline/text-embeddings
{
  "description": "Text embedding pipeline",
  "processors": [
    {
      "inference": {
        "model_id": "sentence-transformers__msmarco-minilm-l12-cos-v5",
        "target_field": "text_embedding",
        "field_map": {
          "text": "text_field"
        }
      }
    }
  ],
  "on_failure": [
    {
      "set": {
        "description": "Index document to 'failed-<index>'",
        "field": "_index",
        "value": "failed-{{{_index}}}"
      }
    },
    {
      "set": {
        "description": "Set error message",
        "field": "ingest.failure",
        "value": "{{_ingest.on_failure_message}}"
      }
    }
  ]
}

4. Reindex data

We want to reindex documents from the collection index into the new collection-with-embeddings index by pushing documents through text-embeddings pipeline, so that documents in the collection-with-embeddings index have an additional field for passages’ embeddings. From Elasticsearch v 8.11, it is not necessary anymore to define index mapping for dense_vector field - long float arrays are automatically mapped as dense_vector field with correct number of dimensions set. But if we want to have more control on index mapping, we can define it as following:

PUT collection-with-embeddings
{
  "mappings": {
    "properties": {
      "text_embedding.predicted_value": {
        "type": "dense_vector"
      },
      "text": {
        "type": "text"
      }
    }
  }
}

Note: from Elasticsearch v 8.11, it is optional to provide dims, index and similarity parameters in the mapping of dense_vector.

Finally, we are ready to reindex. Given that reindex will take some time to process all documents and infer on them, we do reindex in the background by invoking the API with the wait_for_completion=false flag.

POST _reindex?wait_for_completion=false
{
  "source": {
    "index": "collection"
  },
  "dest": {
    "index": "collection-with-embeddings",
    "pipeline": "text-embeddings"
  }
}

The above returns a task id. We can monitor progress of the task with:

GET _tasks/<task_id>

Alternatively, track progress by watching Inference count increase in the model stats API or model stats UI.

The reindexed documents now contain the inference results – vector embeddings. As an example one of the documents looks like this:

{
    "id": "G7PPtn8BjSkJO8zzChzT",
    "text": "This is the definition of RNA along with examples of types of RNA molecules. This is the definition of RNA along with examples of types of RNA molecules. RNA Definition",
    "text_embedding":
    {
     "predicted_value":
        [
            0.057356324046850204,
            0.1602816879749298,
            -0.18122544884681702,
            0.022277727723121643,
            ....
        ],
        "model_id": "sentence-transformers__msmarco-minilm-l12-cos-v5"
    }
}

5. Vector Similarity Search

From Elasticsearch v 8.7, we support implicit generation of embeddings from query terms during a search request using query_vector_builder parameter of knn search. For this you simply need to provide your model_id (in our case it will be "sentence-transformers__msmarco-minilm-l12-cos-v5"), and model_text – the query string from which the model will generate the dense vector representation.

GET collection-with-embeddings/_search
{
  "knn": {
    "field": "text_embedding.predicted_value",
    "query_vector_builder": {
      "text_embedding": { 
        "model_id": "sentence-transformers__msmarco-minilm-l12-cos-v5", 
        "model_text": "how is the weather in jamaica"
      }
    },
    "k": 10,
    "num_candidates": 100
  },
  "fields": [
    "id",
    "text"
  ],
  "_source": false
}

Note: from Elasticsearch v 8.13, it is optional to provide k and num_candidates parameters for knn search.

As a result, we get top 10 closest to the query documents sorted by their proximity to the query:

"hits" : [
      {
        "_index" : "collection-with-embeddings",
        "_id" : "47TPtn8BjSkJO8zzKq_o",
        "_score" : 0.94591534,
        "fields" : {
          "id": [
            434125
          ],
          "text": [
            "The climate in Jamaica is tropical and humid with warm to hot temperatures all year round. The average temperature in Jamaica is between 80 and 90 degrees Fahrenheit. Jamaican nights are considerably cooler than the days, and the mountain areas are cooler than the lower land throughout the year. Continue Reading."
          ]
        }
      },
      {
        "_index" : "collection-with-embeddings",
        "_id" : "3LTPtn8BjSkJO8zzKJO1",
        "_score" : 0.94536424,
        "_source" : {
          "id": [
            4498474
          ],
          "text": [
            "The climate in Jamaica is tropical and humid with warm to hot temperatures all year round. The average temperature in Jamaica is between 80 and 90 degrees Fahrenheit. Jamaican nights are considerably cooler than the days, and the mountain areas are cooler than the lower land throughout the year"
          ]
        }
      },
      {
        "_index" : "collection-with-embeddings",
        "_id" : "KrXPtn8BjSkJO8zzPbDW",
        "_score" :  0.9432083,
        "_source" : {
          "id": [
            190804
          ],
          "text": [
            "Quick Answer. The climate in Jamaica is tropical and humid with warm to hot temperatures all year round. The average temperature in Jamaica is between 80 and 90 degrees Fahrenheit. Jamaican nights are considerably cooler than the days, and the mountain areas are cooler than the lower land throughout the year. Continue Reading"
          ]
        }
      },
...

6. Quick Verification

As we used only a subset of MS MARCO dataset, we can not do a full evaluation. What we can do instead is a simple verification on a few queries just to get a sense that we indeed are getting relevant results, and not some random ones. From the TREC 2019 Deep Learning Track judgements for Passage Ranking Task, we take the 3 last queries, submit them to our vector similarity search, get top 10 results and consult the TREC judgments to see how relevant are the results that we have received. For the Passage Ranking task, passages are judged on a four-point scale of Irrelevant (0), Related (the passage is on-topic but does not answer the question) (1), Highly Relevant (2), and Perfectly Relevant (3).

Please note that our verification is not a rigorous evaluation, it is used only for our quick demo. Since we only indexed passages that are known to be related to the queries, it is a much easier task than the original passage retrieval task. In the future we intend to do a rigorous evaluation on the MS MARCO dataset.

Query #1124210 “tracheids are part of _____” submitted to our vector search returns the following results:

| | |

Query #1129237 “hydrogen is a liquid below what temperature” returns the following results:

| | |

Query #1133167 “how is the weather in jamaica” returns the following results:

| | |

As we can see for all 3 queries Elasticsearch returned mostly relevant results, and the top results for all queries were mostly either highly or perfectly relevant.

Trying it out

NLP is a powerful feature in the Elastic Stack with an exciting roadmap. Discover new features and keep up with the latest developments by building your cluster in Elastic Cloud. Sign up for a free 14-day trial today and try the examples in this blog.

If you want more NLP reads:

• How to deploy NLP named entity recognition NER
• How to deploy NLP sentiment analysis
• How to deploy natural language processing: Getting started