Perform Semantic Search with Elasticsearch and OpenAI

Semantic search using Elasticsearch and OpenAI

This notebook demonstrates how to:

• Index the OpenAI Wikipedia vector dataset into Elasticsearch
• Embed a question with the OpenAI embeddings endpoint
• Perform semantic search on the Elasticsearch index using the encoded question

Install packages and import modules

# install packages

! python3 -m pip install -qU openai pandas wget elasticsearch

# import modules

from getpass import getpass
from elasticsearch import Elasticsearch, helpers
import wget
import zipfile
import pandas as pd
import json
from openai import OpenAI

Connect to Elasticsearch

ℹ️ We're using an Elastic Cloud deployment of Elasticsearch for this notebook.
If you don't already have an Elastic deployment, you can sign up for a free Elastic Cloud trial.

To connect to Elasticsearch, you need to create a client instance with the Cloud ID and password for your deployment.

Find the Cloud ID for your deployment by going to https://cloud.elastic.co/deployments and selecting your deployment.

CLOUD_ID = getpass("Elastic deployment Cloud ID")
CLOUD_PASSWORD = getpass("Elastic deployment Password")
client = Elasticsearch(
  cloud_id = CLOUD_ID,
  basic_auth=("elastic", CLOUD_PASSWORD) # Alternatively use `api_key` instead of `basic_auth`
)

# Test connection to Elasticsearch
print(client.info())

Download the dataset

In this step we download the OpenAI Wikipedia embeddings dataset, and extract the zip file.

embeddings_url = 'https://cdn.openai.com/API/examples/data/vector_database_wikipedia_articles_embedded.zip'
wget.download(embeddings_url)

with zipfile.ZipFile("vector_database_wikipedia_articles_embedded.zip",
"r") as zip_ref:
    zip_ref.extractall("data")

Read CSV file into a Pandas DataFrame

Next we use the Pandas library to read the unzipped CSV file into a DataFrame. This step makes it easier to index the data into Elasticsearch in bulk.

wikipedia_dataframe = pd.read_csv("data/vector_database_wikipedia_articles_embedded.csv")

Create index with mapping

Now we need to create an Elasticsearch index with the necessary mappings. This will enable us to index the data into Elasticsearch.

We use the dense_vector field type for the title_vector and content_vector fields. This is a special field type that allows us to store dense vectors in Elasticsearch.

Later, we'll need to target the dense_vector field for kNN search.

index_mapping= {
    "properties": {
      "title_vector": {
          "type": "dense_vector",
          "dims": 1536,
          "index": "true",
          "similarity": "cosine"
      },
      "content_vector": {
          "type": "dense_vector",
          "dims": 1536,
          "index": "true",
          "similarity": "cosine"
      },
      "text": {"type": "text"},
      "title": {"type": "text"},
      "url": { "type": "keyword"},
      "vector_id": {"type": "long"}
      
    }
}

client.indices.create(index="wikipedia_vector_index", mappings=index_mapping)

Index data into Elasticsearch

The following function generates the required bulk actions that can be passed to Elasticsearch's Bulk API, so we can index multiple documents efficiently in a single request.

For each row in the DataFrame, the function yields a dictionary representing a single document to be indexed.

def dataframe_to_bulk_actions(df):
    for index, row in df.iterrows():
        yield {
            "_index": 'wikipedia_vector_index',
            "_id": row['id'],
            "_source": {
                'url' : row["url"],
                'title' : row["title"],
                'text' : row["text"],
                'title_vector' : json.loads(row["title_vector"]),
                'content_vector' : json.loads(row["content_vector"]),
                'vector_id' : row["vector_id"]
            }
        }

As the dataframe is large, we will index data in batches of 100. We index the data into Elasticsearch using the Python client's helpers for the bulk API.

start = 0
end = len(wikipedia_dataframe)
batch_size = 100
for batch_start in range(start, end, batch_size):
    batch_end = min(batch_start + batch_size, end)
    batch_dataframe = wikipedia_dataframe.iloc[batch_start:batch_end]
    actions = dataframe_to_bulk_actions(batch_dataframe)
    helpers.bulk(client, actions)

Let's test the index with a simple match query.

print(client.search(index="wikipedia_vector_index", body={
    "_source": {
        "excludes": ["title_vector", "content_vector"]
    },
    "query": {
        "match": {
            "text": {
                "query": "Hummingbird"
            }
        }
    }
}))

Encode a question with OpenAI embedding model

To perform semantic search, we need to encode queries with the same embedding model used to encode the documents at index time.
In this example, we need to use the text-embedding-3-small model.

You'll need your OpenAI API key to generate the embeddings.

# Create OpenAI client
openai_client = OpenAI()

# Define question
question = 'Is the Atlantic the biggest ocean in the world?'

question_embedding = openai_client.embeddings.create(
    input=question,
    model="text-embedding-3-small"
)

Run semantic search queries

Now we're ready to run queries against our Elasticsearch index using our encoded question. We'll be doing a k-nearest neighbors search, using the Elasticsearch kNN query option.

First, we define a small function to pretty print the results.

# Function to pretty print Elasticsearch results

def pretty_response(response):
    for hit in response['hits']['hits']:
        id = hit['_id']
        score = hit['_score']
        title = hit['_source']['title']
        text = hit['_source']['text']
        pretty_output = (f"\nID: {id}\nTitle: {title}\nSummary: {text}\nScore: {score}")
        print(pretty_output)

Now let's run our kNN query.

response = client.search(
  index = "wikipedia_vector_index",
  knn={
      "field": "content_vector",
      "query_vector":  question_embedding.data[0].embedding,
      "k": 10,
      "num_candidates": 100
    }
)
pretty_response(response)

Next steps

Success! Now you know how to use Elasticsearch as a vector database to store embeddings, encode queries by calling the OpenAI embeddings endpoint, and run semantic search.

Play around with different queries, and if you want to try with your own data, you can experiment with different embedding models.

ℹ️ Check out our other notebook Retrieval augmented generation using Elasticsearch and OpenAI. That notebook builds on this example to demonstrate how to use Elasticsearch together with the OpenAI chat completions API for retrieval augmented generation (RAG).