Tutorial: semantic search with the inference API

edit

Tutorial: semantic search with the inference API

edit

The instructions in this tutorial shows you how to use the inference API workflow with various services to perform semantic search on your data.

For the easiest way to perform semantic search in the Elastic Stack, refer to the semantic_text end-to-end tutorial.

The following examples use the:

You can use any Cohere and OpenAI models, they are all supported by the inference API. For a list of recommended models available on HuggingFace, refer to the supported model list.

Click the name of the service you want to use on any of the widgets below to review the corresponding instructions.

Requirements

edit

A Cohere account is required to use the inference API with the Cohere service.

Create an inference endpoint

edit

Create an inference endpoint by using the Create inference API:

resp = client.inference.put(
    task_type="text_embedding",
    inference_id="cohere_embeddings",
    inference_config={
        "service": "cohere",
        "service_settings": {
            "api_key": "<api_key>",
            "model_id": "embed-english-v3.0",
            "embedding_type": "byte"
        }
    },
)
print(resp)
const response = await client.inference.put({
  task_type: "text_embedding",
  inference_id: "cohere_embeddings",
  inference_config: {
    service: "cohere",
    service_settings: {
      api_key: "<api_key>",
      model_id: "embed-english-v3.0",
      embedding_type: "byte",
    },
  },
});
console.log(response);
PUT _inference/text_embedding/cohere_embeddings 
{
    "service": "cohere",
    "service_settings": {
        "api_key": "<api_key>", 
        "model_id": "embed-english-v3.0", 
        "embedding_type": "byte"
    }
}

The task type is text_embedding in the path and the inference_id which is the unique identifier of the inference endpoint is cohere_embeddings.

The API key of your Cohere account. You can find your API keys in your Cohere dashboard under the API keys section. You need to provide your API key only once. The Get inference API does not return your API key.

The name of the embedding model to use. You can find the list of Cohere embedding models here.

When using this model the recommended similarity measure to use in the dense_vector field mapping is dot_product. In the case of Cohere models, the embeddings are normalized to unit length in which case the dot_product and the cosine measures are equivalent.

Create the index mapping

edit

The mapping of the destination index - the index that contains the embeddings that the model will create based on your input text - must be created. The destination index must have a field with the dense_vector field type for most models and the sparse_vector field type for the sparse vector models like in the case of the elasticsearch service to index the output of the used model.

resp = client.indices.create(
    index="cohere-embeddings",
    mappings={
        "properties": {
            "content_embedding": {
                "type": "dense_vector",
                "dims": 1024,
                "element_type": "byte"
            },
            "content": {
                "type": "text"
            }
        }
    },
)
print(resp)
response = client.indices.create(
  index: 'cohere-embeddings',
  body: {
    mappings: {
      properties: {
        content_embedding: {
          type: 'dense_vector',
          dims: 1024,
          element_type: 'byte'
        },
        content: {
          type: 'text'
        }
      }
    }
  }
)
puts response
const response = await client.indices.create({
  index: "cohere-embeddings",
  mappings: {
    properties: {
      content_embedding: {
        type: "dense_vector",
        dims: 1024,
        element_type: "byte",
      },
      content: {
        type: "text",
      },
    },
  },
});
console.log(response);
PUT cohere-embeddings
{
  "mappings": {
    "properties": {
      "content_embedding": { 
        "type": "dense_vector", 
        "dims": 1024, 
        "element_type": "byte"
      },
      "content": { 
        "type": "text" 
      }
    }
  }
}

The name of the field to contain the generated tokens. It must be refrenced in the inference pipeline configuration in the next step.

The field to contain the tokens is a dense_vector field.

The output dimensions of the model. Find this value in the Cohere documentation of the model you use.

The name of the field from which to create the dense vector representation. In this example, the name of the field is content. It must be referenced in the inference pipeline configuration in the next step.

The field type which is text in this example.

Create an ingest pipeline with an inference processor

edit

Create an ingest pipeline with an inference processor and use the model you created above to infer against the data that is being ingested in the pipeline.

resp = client.ingest.put_pipeline(
    id="cohere_embeddings_pipeline",
    processors=[
        {
            "inference": {
                "model_id": "cohere_embeddings",
                "input_output": {
                    "input_field": "content",
                    "output_field": "content_embedding"
                }
            }
        }
    ],
)
print(resp)
const response = await client.ingest.putPipeline({
  id: "cohere_embeddings_pipeline",
  processors: [
    {
      inference: {
        model_id: "cohere_embeddings",
        input_output: {
          input_field: "content",
          output_field: "content_embedding",
        },
      },
    },
  ],
});
console.log(response);
PUT _ingest/pipeline/cohere_embeddings_pipeline
{
  "processors": [
    {
      "inference": {
        "model_id": "cohere_embeddings", 
        "input_output": { 
          "input_field": "content",
          "output_field": "content_embedding"
        }
      }
    }
  ]
}

The name of the inference endpoint you created by using the Create inference API, it’s referred to as inference_id in that step.

Configuration object that defines the input_field for the inference process and the output_field that will contain the inference results.

Load data

edit

In this step, you load the data that you later use in the inference ingest pipeline to create embeddings from it.

Use the msmarco-passagetest2019-top1000 data set, which is a subset of the MS MARCO Passage Ranking data set. It consists of 200 queries, each accompanied by a list of relevant text passages. All unique passages, along with their IDs, have been extracted from that data set and compiled into a tsv file.

Download the file and upload it to your cluster using the Data Visualizer in the Machine Learning UI. After your data is analyzed, click Override settings. Under Edit field names, assign id to the first column and content to the second. Click Apply, then Import. Name the index test-data, and click Import. After the upload is complete, you will see an index named test-data with 182,469 documents.

Ingest the data through the inference ingest pipeline

edit

Create embeddings from the text by reindexing the data through the inference pipeline that uses your chosen model. This step uses the reindex API to simulate data ingestion through a pipeline.

resp = client.reindex(
    wait_for_completion=False,
    source={
        "index": "test-data",
        "size": 50
    },
    dest={
        "index": "cohere-embeddings",
        "pipeline": "cohere_embeddings_pipeline"
    },
)
print(resp)
const response = await client.reindex({
  wait_for_completion: "false",
  source: {
    index: "test-data",
    size: 50,
  },
  dest: {
    index: "cohere-embeddings",
    pipeline: "cohere_embeddings_pipeline",
  },
});
console.log(response);
POST _reindex?wait_for_completion=false
{
  "source": {
    "index": "test-data",
    "size": 50 
  },
  "dest": {
    "index": "cohere-embeddings",
    "pipeline": "cohere_embeddings_pipeline"
  }
}

The default batch size for reindexing is 1000. Reducing size to a smaller number makes the update of the reindexing process quicker which enables you to follow the progress closely and detect errors early.

The rate limit of your Cohere account may affect the throughput of the reindexing process.

The call returns a task ID to monitor the progress:

resp = client.tasks.get(
    task_id="<task_id>",
)
print(resp)
const response = await client.tasks.get({
  task_id: "<task_id>",
});
console.log(response);
GET _tasks/<task_id>

Reindexing large datasets can take a long time. You can test this workflow using only a subset of the dataset. Do this by cancelling the reindexing process, and only generating embeddings for the subset that was reindexed. The following API request will cancel the reindexing task:

resp = client.tasks.cancel(
    task_id="<task_id>",
)
print(resp)
const response = await client.tasks.cancel({
  task_id: "<task_id>",
});
console.log(response);
POST _tasks/<task_id>/_cancel

Semantic search

edit

After the data set has been enriched with the embeddings, you can query the data using semantic search. In case of dense vector models, pass a query_vector_builder to the k-nearest neighbor (kNN) vector search API, and provide the query text and the model you have used to create the embeddings. In case of a sparse vector model like ELSER, use a sparse_vector query, and provide the query text with the model you have used to create the embeddings.

If you cancelled the reindexing process, you run the query only a part of the data which affects the quality of your results.

resp = client.search(
    index="cohere-embeddings",
    knn={
        "field": "content_embedding",
        "query_vector_builder": {
            "text_embedding": {
                "model_id": "cohere_embeddings",
                "model_text": "Muscles in human body"
            }
        },
        "k": 10,
        "num_candidates": 100
    },
    source=[
        "id",
        "content"
    ],
)
print(resp)
response = client.search(
  index: 'cohere-embeddings',
  body: {
    knn: {
      field: 'content_embedding',
      query_vector_builder: {
        text_embedding: {
          model_id: 'cohere_embeddings',
          model_text: 'Muscles in human body'
        }
      },
      k: 10,
      num_candidates: 100
    },
    _source: [
      'id',
      'content'
    ]
  }
)
puts response
const response = await client.search({
  index: "cohere-embeddings",
  knn: {
    field: "content_embedding",
    query_vector_builder: {
      text_embedding: {
        model_id: "cohere_embeddings",
        model_text: "Muscles in human body",
      },
    },
    k: 10,
    num_candidates: 100,
  },
  _source: ["id", "content"],
});
console.log(response);
GET cohere-embeddings/_search
{
  "knn": {
    "field": "content_embedding",
    "query_vector_builder": {
      "text_embedding": {
        "model_id": "cohere_embeddings",
        "model_text": "Muscles in human body"
      }
    },
    "k": 10,
    "num_candidates": 100
  },
  "_source": [
    "id",
    "content"
  ]
}

As a result, you receive the top 10 documents that are closest in meaning to the query from the cohere-embeddings index sorted by their proximity to the query:

"hits": [
      {
        "_index": "cohere-embeddings",
        "_id": "-eFWCY4BECzWLnMZuI78",
        "_score": 0.737484,
        "_source": {
          "id": 1690948,
          "content": "Oxygen is supplied to the muscles via red blood cells. Red blood cells carry hemoglobin which oxygen bonds with as the hemoglobin rich blood cells pass through the blood vessels of the lungs.The now oxygen rich blood cells carry that oxygen to the cells that are demanding it, in this case skeletal muscle cells.ther ways in which muscles are supplied with oxygen include: 1  Blood flow from the heart is increased. 2  Blood flow to your muscles in increased. 3  Blood flow from nonessential organs is transported to working muscles."
        }
      },
      {
        "_index": "cohere-embeddings",
        "_id": "HuFWCY4BECzWLnMZuI_8",
        "_score": 0.7176013,
        "_source": {
          "id": 1692482,
          "content": "The thoracic cavity is separated from the abdominal cavity by the  diaphragm. This is a broad flat muscle.    (muscular) diaphragm The diaphragm is a muscle that separat…e the thoracic from the abdominal cavity. The pelvis is the lowest part of the abdominal cavity and it has no physical separation from it    Diaphragm."
        }
      },
      {
        "_index": "cohere-embeddings",
        "_id": "IOFWCY4BECzWLnMZuI_8",
        "_score": 0.7154432,
        "_source": {
          "id": 1692489,
          "content": "Muscular Wall Separating the Abdominal and Thoracic Cavities; Thoracic Cavity of a Fetal Pig; In Mammals the Diaphragm Separates the Abdominal Cavity from the"
        }
      },
      {
        "_index": "cohere-embeddings",
        "_id": "C-FWCY4BECzWLnMZuI_8",
        "_score": 0.695313,
        "_source": {
          "id": 1691493,
          "content": "Burning, aching, tenderness and stiffness are just some descriptors of the discomfort you may feel in the muscles you exercised one to two days ago.For the most part, these sensations you experience after exercise are collectively known as delayed onset muscle soreness.urning, aching, tenderness and stiffness are just some descriptors of the discomfort you may feel in the muscles you exercised one to two days ago."
        }
      },
      (...)
    ]

Interactive tutorials

edit

You can also find tutorials in an interactive Colab notebook format using the Elasticsearch Python client: