Elasticsearch meets BERT: Building Search Engine with Elasticsearch and BERT | by Hironsan | Towards Data Science

In this post, we use a pre-trained BERT model and Elasticsearch to build a search engine. Elasticsearch has recently released text similarity search with vector fields. On the other hand, you can convert text into a fixed-length vector using BERT. So once we convert documents into vectors by BERT and store them into Elasticsearch, we can search similar documents with Elasticsearch and BERT.
This post implements a search engine with Elasticsearch and BERT by the following architecture. Here, we use Docker to divide the whole system into three parts: application, BERT, and Elasticsearch. The purpose is to make it easier to scale each service.
I show you only important part in this post, but the entire system is written in docker-compose.yamlin the following GitHub repository. Please take a look at the repository:
https://github.com/Hironsan/bertsearch
The plan for this post is:
Download a pre-trained BERT modelSetting environment variablesLaunching Docker containerCreating Elasticsearch indexCreating documentsIndexing document
Download a pre-trained BERT model
Setting environment variables
Launching Docker container
Creating Elasticsearch index
Creating documents
Indexing document
First, download a pre-trained BERT model. The following command is an example of downloading an English model:
$ wget https://storage.googleapis.com/bert_models/2018_10_18/cased_L-12_H-768_A-12.zip$ unzip cased_L-12_H-768_A-12.zip
You need to set a pre-trained BERT model and Elasticsearch’s index name as environment variables. These variables are used in Docker containers. Below is an example that specifies jobsearch as index name and ./cased_L-12_H-768_A-12 as model’s path:
$ export PATH_MODEL=./cased_L-12_H-768_A-12$ export INDEX_NAME=jobsearch
Now, let’s launch Docker containers by using Docker compose. There are three containers to launch here: application container, BERT container, and Elasticsearch container. Check docker-compose.yaml for more details.
$ docker-compose up
Note that I recommend you that you should assign more memory(over 8GB) to Docker. Because BERT container needs large memory.
You can use the create index API to add a new index to an Elasticsearch cluster. When creating an index, you can specify the following:
Settings for the index
Mappings for fields in the index
Index aliases
For example, if you want to create jobsearch index with title, text and text_vector fields, you can create the index by the following command:
$ python example/create_index.py --index_file=example/index.json --index_name=jobsearch# index.json{  "settings": {    "number_of_shards": 2,    "number_of_replicas": 1  },  "mappings": {    "dynamic": "true",    "_source": {      "enabled": "true"    },    "properties": {      "title": {        "type": "text"      },      "text": {        "type": "text"      },      "text_vector": {        "type": "dense_vector",        "dims": 768      }    }  }}
CAUTION: The dims value of text_vector must need to match the dims of a pre-trained BERT model.
Once you created an index, you’re ready to index some document. The point here is to convert your document into a vector using BERT. The resulting vector is stored in the text_vector field. Let’s convert your data into a JSON document:
$ python example/create_documents.py --data=example/example.csv --index_name=jobsearch# example/example.csv"Title","Description""Saleswoman","lorem ipsum""Software Developer","lorem ipsum""Chief Financial Officer","lorem ipsum""General Manager","lorem ipsum""Network Administrator","lorem ipsum"
After finishing the script, you can get a JSON document like follows:
# documents.jsonl{"_op_type": "index", "_index": "jobsearch", "text": "lorem ipsum", "title": "Saleswoman", "text_vector": [...]}{"_op_type": "index", "_index": "jobsearch", "text": "lorem ipsum", "title": "Software Developer", "text_vector": [...]}{"_op_type": "index", "_index": "jobsearch", "text": "lorem ipsum", "title": "Chief Financial Officer", "text_vector": [...]}...
After converting your data into a JSON, you can add a JSON document to the specified index and makes it searchable.
$ python example/index_documents.py
Go to http://127.0.0.1:5000. Below is an example of a job search. You can see that the “I’m looking for a lawyer” query returns jobs like “Legal assistant” and “Lawyer”.
In this post, we implemented the search engine with Elasticsearch and BERT. Although there is a problem with the execution speed of BERT, it is easy to scale by treating BERT as an independent container like this architecture, so I think that the problem can be solved. I hope this article will be useful to you.