Semantic Search with Amazon OpenSearch Serverless and Titan

This two-part tutorial series will introduce the workflow and APIs needed to build a Q&A system; in this example, we’ll build our application based on the Academy Awards dataset from Kaggle. You will need an active Amazon Web Services (AWS) subscription to follow these steps.

We’ll be using Amazon Bedrock and Amazon Titan foundation models (FMs); Amazon has recently introduced the general availability of the embeddings models. With these two models, we can implement text generation, summarization, sentiment analysis, question answering and more.

This two-part tutorial series will walk you through the steps in implementing Retrieval Augmented Generation (RAG) based on Amazon Bedrock, Amazon Titan, and Amazon OpenSearch Serverless. RAG is a technique used for “grounding” a large language model (LLM) with information that is use-case specific and relevant.

You can find part two of the series here.

In the first part, we will set up the environment, convert the custom dataset into text embeddings, and ingest the embeddings into Amazon OpenSearch Serverless Vector DB, enabling us to implement RAG with Amazon Titan FMs.

Step 1: Configure the Environment

Let’s start by configuring the Python virtual environment with the required dependencies and modules.

mkdir qa-demo
cd qa-demo
python -m venv venv
source venv/bin/activate

Then, we will create a subdirectory to download and install the latest version of Boto3 and the AWS CLI updated for Bedrock.

mkdir sdk
cd sdk

curl -sS https://d2eo22ngex1n9g.cloudfront.net/Documentation/SDK/bedrock-python-sdk.zip > sdk.zip

pip install --no-build-isolation --force-reinstall \
 awscli-*-py3-none-any.whl \
 boto3-*-py3-none-any.whl \
 botocore-*-py3-none-any.whl

Finally, let’s install other dependencies related to OpenSearch, Jupyter, and Pandas.

pip install jupyter opensearch-py pandas

Step 2: Create the Amazon OpenSearch Serverless Collection and Index

In this step, we will provision the vector database for storing and searching the embeddings.

1. Open the Amazon OpenSearch Service console at https://console.aws.amazon.com/aos/home.
2. Choose Collections in the left navigation pane and choose Create collection.
3. Name the collection oscars-collection.
4. For collection type, choose Vector search.
5. Under Security, select Easy create to streamline your security configuration.
6. Choose Next.
7. Review your collection settings and choose Submit.

👁 Image

It will take several minutes for the collection to be ready. You can track the status on the AWS Management Console.

👁 Image

Once the collection is active, create an index programmatically through the Boto3 SDK. For this, we need the endpoint associated with the collection, which is available in the console.

👁 Image

Launch a new Jupyter Notebook to run the code shown in this tutorial.

Initialize the OpenSearch client and create the index. Replace HOST with the correct value shown in the console.

from opensearchpy import OpenSearch, RequestsHttpConnection, AWSV4SignerAuth, helpers
import boto3
from botocore.config import Config
import json

boto_config = Config(
 region_name = 'us-west-2',
 signature_version = 'v4',
 retries = {
 'max_attempts': 10,
 'mode': 'standard'
 }
)
client = boto3.client("opensearchserverless",config=boto_config)

host = "HOST.us-west-2.aoss.amazonaws.com"
region = "us-west-2"
service = "aoss"
credentials = boto3.Session().get_credentials()
auth = AWSV4SignerAuth(credentials, region, service)

client = OpenSearch(
 hosts = [{"host": host, "port": 443}],
 http_auth = auth,
 use_ssl = True,
 verify_certs = True,
 connection_class = RequestsHttpConnection,
 pool_maxsize = 20
)

index_name = "oscars-index"
index_body = {
 "mappings": {
 "properties": {
 "nominee_text": {"type": "text"},
 "nominee_vector": {
 "type": "knn_vector",
 "dimension": 4096,
 "method": {
 "engine": "nmslib",
 "space_type": "cosinesimil",
 "name": "hnsw",
 "parameters": {"ef_construction": 512, "m": 16},
 },
 },
 }
 },
 "settings": {
 "index": {
 "number_of_shards": 2,
 "knn.algo_param": {"ef_search": 512},
 "knn": True,
 }
 },
}

try:
 response = client.indices.create(index_name, body=index_body)
 print(json.dumps(response, indent=2))
except Exception as ex:
 print(ex)

We created an index based on the KNN and cosine similarity search algorithms. It has two properties – nominee_text and nominee_vector to store the text and the corresponding embeddings.

Step 3: Pre-processing the Dataset

Download the Oscar Award dataset from Kaggle and move the CSV file to a subdirectory named data. The dataset has all the categories, nominations, and winners of the Academy Awards from 1927 to 2023. I renamed the CSV file to oscars.csv Start by importing the Pandas library and loading the dataset:

import pandas as pd
df=pd.read_csv('./data/oscars.csv')
df.head()

The dataset is well-structured, with column headers and rows that represent the details of each category, including the name of the actor/technician, the film, and whether the nomination was won or lost.

Since we are most interested in awards related to 2023, let’s filter them and create a new Pandas dataframe. At the same time, we will also convert the category to lowercase while dropping the rows where the value of a film is blank.

df=df.loc[df['year_ceremony'] == 2023]
df=df.dropna(subset=['film'])
df['category'] = df['category'].str.lower()
df.head()

With the filtered and cleansed dataset, let’s add a new column to the data frame that has an entire sentence representing a nomination. This complete sentence will be used to generate the text embeddings later.

df['text'] = df['name'] + ' got nominated under the category, ' + df['category'] + ', for the film ' + df['film'] + ' to win the award'
df.loc[df['winner'] == False, 'text'] = df['name'] + ' got nominated under the category, ' + df['category'] + ', for the film ' + df['film'] + ' but did not win'
df.head()['text']

Notice how we concatenate the values to generate a complete sentence. For example, the column “text” in the first two rows of the data frame has the below values:

Austin Butler got nominated under the category, actor in a leading role, for the film Elvis but did not win

Colin Farrell got nominated under the category, actor in a leading role, for the film The Banshees of Inisherin but did not win

Step 4: Generate Embeddings with Titan

Now that we have the text constructed from the dataset let’s convert it into word embeddings. This is a crucial step, as the tokens generated by the embedding model will help us perform a semantic search to retrieve the sentences from the dataset with similar meanings.

Let’s define a function to convert the input text into an embedding.

def text_embedding(text):
 body=json.dumps({"inputText": text})
 response = bedrock.invoke_model(body=body, modelId='amazon.titan-e1t-medium', accept='application/json', contentType='application/json')
 response_body = json.loads(response.get('body').read())
 embedding = response_body.get('embedding')
 return embedding

We can invoke this per row of the data frame and create a new column to store the embeddings.

df=df.assign(embedding=(df["text"].apply(lambda x : text_embedding(x))))

👁 Image

We are ready to ingest the text and the corresponding embeddings into Amazon OpenSearch Serverless.

Step 5: Insert the Embeddings and Performing a Similarity Search

Let’s create a function to insert each row of the data frame into the vector database.

def add_document(vector,text):
 document = {
 "nominee_vector": vector,
 "nominee_text": text
 }
 
 response = client.index(
 index = 'oscars-index',
 body = document
 )
 print('\nAdding document:')
 print(response) 

Invoking this function directly by calling the apply method on the data frame is easy.

df.apply(lambda row: add_document(row['embedding'], row['text']), axis=1)

Once the data is ingested, we are ready to perform the search. Let’s create a function that accepts a vector and returns text that’s similar to the input.

def search_index(vector):
 document = {
 "size": 15,
 "_source": {"excludes": ["nominee_vector"]},
 "query": {
 "knn": {
 "nominee_vector": {
 "vector": vector,
 "k":15
 }
 }
 }
 }
 response = client.search(
 body = document,
 index = "oscars-index"
 )
 return response

The line "_source": {"excludes": ["nominee_vector"]}, excludes the vector in the result, which saves bandwidth.

Let’s now run a query on the vector database. Before that, we need to convert our query to an embedding based on Titan.

query='who won the award for best music?'
vector=text_embedding(query)

Now, let’s pass the embedding to the search function and see the results.

response=search_index(vector)
response['hits']['hits']

👁 Image

As you can see, we got all the nominations related to sound and music from the vector database, along with the similarity score. By extracting the contents of the nominee_text element, we can construct the context that will be used to augment the prompt sent to the LLM.

In the next part of this tutorial, we will explore how to leverage the output to increase the accuracy of the Titan model’s response. Stay tuned.