Building a k-NN Similarity Search Engine using Amazon Elasticsearch and SageMaker | by Yi Ai | Towards Data Science

Amazon Elasticsearch Service recently added support for k-nearest neighbor search. It enables you to run high scale and low latency k-NN search across thousands of dimensions with the same ease as running any regular Elasticsearch query.
k-NN similarity search is powered by Open Distro for Elasticsearch, an Apache 2.0-licensed distribution of Elasticsearch.
In this post, I’ll show you how to build a scalable similarity questions search api using Amazon Sagemaker, Amazon Elasticsearch, Amazon Elastic File System (EFS) and Amazon ECS.
Deploy and run a Sagemaker notebook instance in VPC.
Mount EFS to notebook instance.
Download Quora Question Pairs dataset, then map variable-length questions from dataset to fixed-length vectors using DistilBERT model.
Create downstream task to reduce embedding dimensions and save sentence embedder to EFS.
Transform questions text to vectors, and index all vectors to Elasticsearch.
Deploy a containerized Flask rest api to ECS.
The following diagram shows the architecture of above steps:
First, Let’s create a Sagemaker notebook instance which connects to Elasticsearch and make sure then are in the same VPC.
To configure VPC options in Sagemaker console, within Network section of Create notebook instance page, set VPC network configuration details such as VPC subnet IDs and security group IDs:
We will do all the necessary sentence transforming steps in a SageMaker notebook (code found here).
Now, mounting EFS to model directory, for more details about EFS, please check AWS official document.
NOTE:
fs-xxxxx.efs.ap-southeast-2.amazonaws.com is DNS name of EFS.
EFS Mount targets and Sagemaker are in same VPC.
To run nearest neighbor search, we have to get sentence and tokens embeddings. We can use sentence-transformers which is a sentence embeddings using BERT / RoBERTa / DistilBERT / ALBERT / XLNet with PyTorch. It allows us to map sentences into fixed-length representations in just a few lines of code.
We will use light weight DistilBERT model to generate sentence embeddings in this example, please note that the number of hidden units of DistilBERT is 768. This dimension seems too big to Elasticsearch index, we can reduce the dimension to 256 by adding a dense layer after the pooling:
Next, Save sentence embedder to EFS mounted directory:
transformer.save("model/transformer-v1/")
We need to make sure dataset has been downloaded, the data set for this example is quora question paris datas.
Next, extract full text of each question to dataframe:
To start with, create a kNN index,
then transform and index question vectors to Elasticsearch.
Questions in Elasticsearch have the following structure:
{'question_vector': [-0.06435434520244598, ... ,0.0726890116930008],'question': 'How hard is it to learn to play piano as an adult?'}
We have embedded questions into fixed-length vectors and indexed all of vectors to Elasticsearch. Let’s create a rest api connects to Elasticsearch and test it out!
We will use sample cloud formation template to create ECS Cluster and service in VPC (templates and bash scripts found here).
We will use EFS volumes with ECS, search flow is 1)Flask application loads saved sentence embedder in EFS volume, 2)transform input parameter sentence to vectors, 3)then query K-Nearest neighbors in Elasticsearch.
We now have flask api running in ECS container, let’s use the basic search function to find similar questions when it comes to the query: “What is best way to make money online?”:
$curl --data 'question=What is best way to make money online?' --data 'size=5' --data 'min_score=0.3'  -X POST http://knn-s-publi-xxxx-207238135.ap-southeast-2.elb.amazonaws.com/search
Check out the result:
As you can see, the result was pretty amazing, you can also fine-tune your own sentence embedding methods, so that you get task-specific sentence embeddings for k-NN search.
Great! We have what we need! I hope you have found this article useful.
The complete scripts can be found in my GitHub repo.