Part #3: Semantically Search Your Data With MongoDB Atlas Vector Search
.jpg&w=3840&q=75)
Dominic Frei6 min read • Published Feb 09, 2024 • Updated Feb 14, 2024
Rate this tutorial
This final part of the series will show you how to use the Amazon SageMaker endpoint created in the previous part and perform a semantic search on your data using MongoDB Atlas Vector Search. The two parts shown in this tutorial will be:
- Creating and updating embeddings/vectors for your data.
- Creating vectors for a search query and sending them via Atlas Vector Search.
If you haven’t done so, create a new cluster in your MongoDB Atlas account. Make sure to check
[Add sample dataset](https://www.mongodb.com/developer/products/atlas/atlas-sample-datasets/) to get the sample data we will be working with right away into your cluster.
If you are using a cluster that has already been set up, click on the three dots in your cluster and then
Load Sample Dataset.
There is one more step we need to take in Atlas, which is creating a search index, specifically for Vector Search.
In your database overview, click on
Create Index.
The Search page will be shown. Click on
Create Search Index here.
Then choose
Atlas Vector Search -> JSON Editor.
Open the
sample_mflix database and choose the embedded_movies collection (not movies). Name your index (here, VectorSearchIndex).
The configuration for the index needs to state the number of dimensions. That depends on the model and in our case, it’s 384. The path tells the index which field will be used to store the vectors. We’ll call it
embedding here. The similarity for text is usually done with the cosine function.
Hit
Create and you’re good to go.Are you ready for the final part?
Let’s have a look at the code (here, in Python)!
In the following section, we will look at the three relevant files that show you how you can implement a server app that uses the Amazon SageMaker endpoint.
The
sagemaker.py module is the wrapper around the Lambda/Gateway endpoint that we created in the previous example.Make sure to create a
.env file with the URL saved in EMBDDING_SERVICE.It should look like this:
The following function will then attach the query that we want to search for to the URL and execute it.
As a result, we expect to find the vector in a JSON field called
embedding.The module
atlas.py is the wrapper around everything MongoDB Atlas.Similar to
sagemaker.py, we first grab the MONGODB_CONNECTION_STRING that you can retrieve in Atlas by clicking on Connect in your cluster. It’s the authenticated URL to your cluster. We need to save MONGODB_CONNECTION_STRING to our .env file too.We then go ahead and define a bunch of variables that we’ve set in earlier parts, like
VectorSearchIndex and embedding, along with the automatically created sample_mflix demo data.Using the Atlas driver for Python (called PyMongo), we then create a
MongoClient which holds the connection to the Atlas cluster.The first step will be to actually prepare the already existing data with embeddings.
This is the sole purpose of the
add_missing_embeddings function.We’ll create a filter for the documents with missing embeddings and retrieve those from the database, only showing their plot, which is the only field we’re interested in for now.
Assuming we will only find a couple every time, we can then go through them and call the
create_embedding endpoint for each, creating an embedding for the plot of the movie.We’ll then add those new embeddings to the
movies_to_update array so that we eventually only need one bulk_write to the database, which makes the call more efficient.Note that for huge datasets with many embeddings to create, you might want to adjust the lambda function to take an array of queries instead of just a single query. For this simple example, it will do.
Now that the data is prepared, we add two more functions that we need to offer a nice REST service for our client application.
First, we want to be able to update the plot, which inherently means we need to update the embeddings again.
The
update_plot is similar to the initial add_missing_embeddings function but a bit simpler since we only need to update one document.The other function we need to offer is the actual vector search. This can be done using the MongoDB Atlas aggregation pipeline that can be accessed via the Atlas driver.
The
$vectorSearch stage needs to include the index name we want to use, the path to the embedding, and the information about how many results we want to get. This time, we only want to retrieve the title, so we add a $project stage to the pipeline. Make sure to use list to turn the cursor that the search returns into a python list.Now, we can put it all together. Let’s use Flask to expose a REST service for our client application.
One route we want to expose is
/movies/<title> that can be executed with a PUT operation to update the plot of a movie given the title. The title will be a query parameter while the plot is passed in via the body. This function is using the update_plot that we created before in atlas.py and returns the movie with its new plot on success.The other endpoint, finally, is the vector search:
/movies/search.A
query is POST’ed to this endpoint which will then use create_embedding first to create a vector from this query. Note that we need to also create vectors for the query because that’s what the vector search needs to compare it to the actual data (or rather, its embeddings).We then call
execute_vector_search with this embedding to retrieve the results, which will be returned on success.And that’s about all you have to do. Easy, wasn’t it?
Go ahead and run the Flask app (main.py) and when ready, send a cURL to see Atlas Vector Search in action. Here is an example when running it locally:
This should lead to the following result:
War of the Worlds — a movie about Earth, Mars, and an invasion. And what a great one, right?
Of course, this is just a quick and short overview of how to use Amazon SageMaker to create vectors and then search via Vector Search.
We do have a full workshop for you to learn about all those parts in detail. Please visit the Search Lab GitHub page to learn more.
✅ Already have an AWS account? Atlas supports paying for usage via the AWS Marketplace (AWS MP) without any upfront commitment — simply sign up for MongoDB Atlas via AWS Marketplace.