Building a Search Engine with BERT and TensorFlow | by Denis Antyukhov | Towards Data Science

Feature extractors based on deep Neural Probabilistic Language Models such as BERT may extract features that are relevant for a wide array of downstream NLP tasks. For that reason, they are sometimes referred to as Natural Language Understanding (NLU) modules.
These features may also be utilized for computing the similarity between text samples, which is useful for instance-based learning algorithms (e.g. K-NN). To illustrate this, we will build a nearest neighbour search engine for text, using BERT for feature extraction.
The plan for this experiment is:
getting the pre-trained BERT model checkpointextracting a sub-graph optimized for inferencecreating a feature extractor with tf.Estimatorexploring vector space with T-SNE and Embedding Projectorimplementing a nearest neighbour search engineaccelerating search queries with mathexample: building a movie recommendation system
getting the pre-trained BERT model checkpoint
extracting a sub-graph optimized for inference
creating a feature extractor with tf.Estimator
exploring vector space with T-SNE and Embedding Projector
implementing a nearest neighbour search engine
accelerating search queries with math
example: building a movie recommendation system
This guide contains implementations of two things: a BERT text feature extractor and a nearest neighbour search engine.
This guide should be useful for researchers interested in using BERT for natural language understanding tasks. It may also serve as a worked example of interfacing with tf.Estimator API.
For a reader familiar with TensorFlow it should take around 30 minutes to finish this guide.
The code for this experiment is available in Colab here. Also, check out the repository I set up for my BERT experiments: it contains bonus stuff!
Now, let’s start.
We start with a pre-trained BERT checkpoint. For demonstration purposes, I will be using the uncased English model pre-trained by Google engineers. To train a model for a different language, check out my other guide.
For configuring and optimizing the graph for inference we will make use of the awesome bert-as-a-service repository. This repository allows for serving BERT models for remote clients over TCP.
Having a remote BERT-server is beneficial in multi-host environments. However, in this part of the experiment we will focus on creating a local (in-process) feature extractor. This is useful if one wishes to avoid additional latency and potential failure modes introduced by a client-server architecture.
Now, let us download the model and install the package.
!wget https://storage.googleapis.com/bert_models/2019_05_30/wwm_uncased_L-24_H-1024_A-16.zip!unzip wwm_uncased_L-24_H-1024_A-16.zip!pip install bert-serving-server --no-deps
Normally, to modify the model graph we would have to do some low-level TensorFlow programming. However, thanks to bert-as-a-service, we can configure the inference graph using a simple CLI interface.
There are a couple of parameters there to look out for.
For each text sample, BERT encoding layers output a tensor of shape [sequence_len, encoder_dim] with one vector per token. We need to apply some sort of pooling if we are to obtain a fixed representation.
POOL_STRAT parameter defines the pooling strategy applied to the encoder layer number POOL_LAYER. The default value ‘REDUCE_MEAN’ averages the vectors for all tokens in a sequence. This strategy works best for most sentence-level tasks when the model is not fine-tuned. Another option is NONE, in which case no pooling is applied at all. This is useful for word-level tasks such as Named Entity Recognition or POS tagging. For a detailed discussion of these options check out the Han Xiao’s blog post.
SEQ_LEN affects the maximum length of sequences processed by the model. Smaller values will increase the model inference speed almost linearly.
Running the above command will put the model graph and weights into a GraphDef object which will be serialized to a pbtxt file at GRAPH_OUT. The file will often be smaller than the pre-trained model because the nodes and variables required for training will be removed. This results in a very portable solution: for example, the english model only takes 380 MB after serializing.
Now, we will use the serialized graph to build a feature extractor using the tf.Estimator API. We will need to define two things: input_fn and model_fn
input_fn manages getting the data into the model. That includes executing the whole text preprocessing pipeline and preparing a feed_dict for BERT.
First, each text sample is converted into a tf.Example instance containing the necessary features listed in INPUT_NAMES. The bert_tokenizer object contains the WordPiece vocabulary and performs the text preprocessing. After that, the examples are re-grouped by feature name in a feed_dict.
tf.Estimators have a fun feature which makes them rebuild and reinitialize the whole computational graph at each call to the predict function. So, in order to avoid the overhead, we will pass a generator to the predict function, and the generator will yield the features to the model in a never-ending loop. Ha-ha.
model_fn contains the specification of the model. In our case, it is loaded from the pbtxt file we saved in the previous step. The features are mapped explicitly to the corresponding input nodes via input_map.
Now we have almost everything we need to perform inference. Let’s do this!
Due to our use of generators, consecutive calls to bert_vectorizer will not trigger the model rebuild.
>>> bert_vectorizer = build_vectorizer(estimator, build_input_fn)>>> bert_vectorizer(64*['sample text']).shape(64, 768)
A standalone version of the feature extractor described above can be found in the repository.
Now it’s time for a demonstration!
Using the vectorizer, we will generate embeddings for articles from the Reuters-21578 benchmark corpus.To visualize and explore the embedding vector space in 3D, we will use a dimensionality reduction technique called T-SNE.
Let’s get the article embeddings first.
The interactive visualization of generated embeddings is available on the Embedding Projector.
From the link you can run T-SNE yourself or load a checkpoint using the bookmark in lower-right corner (loading works only on Chrome).
Building a supervised model with the generated features is straightforward:
Now, let’s say we have a knowledge base of 50k text samples and we need to answer queries based on this data, fast. How do we retrieve the sample most similar to a query from a text database? The answer is nearest neighbour search.
Formally, the search problem we will be solving is defined as follows:given a set of points S in vector space M, and a query point Q ∈ M, find the closest point in S to Q. There are multiple ways to define ‘closest’ in vector space, we will use Euclidean distance.
So, to build an Search Engine for text we will follow these steps:
Vectorize all samples from the knowledge base — that gives SVectorize the query — that gives QCompute euclidean distances D between Q and SSort D in ascending order — providing indices of the most similar samplesRetrieve labels for said samples from the knowledge base
Vectorize all samples from the knowledge base — that gives S
Vectorize the query — that gives Q
Compute euclidean distances D between Q and S
Sort D in ascending order — providing indices of the most similar samples
Retrieve labels for said samples from the knowledge base
To make the simple matter of implementing this a bit more exciting, we will do it in pure TensorFlow.
First, we create the placeholders for Q and S
Define euclidean distance computation
Finally, get the most similar sample indices
Now that we have a basic retrieval algorithm set up, the question is: can we make it run faster? With a tiny bit of math, we can.
For a pair of vectors p and q, the euclidean distance is defined as follows:
Which is exactly how we did compute it in Step 4.
However, since p and q are vectors, we can expand and rewrite this:
where ⟨...⟩ denotes inner product.
In TensorFlow this can be written as follows:
By the way, in the formula above PP and QQ are actually squared L2 norms of the respective vectors. If both vectors are L2 normalized, then PP = QQ = 1. That gives an interesting relation between inner product and euclidean distance:
However, doing L2 normalization discards the information about vector magnitude which in many cases is undesirable .
Instead, we may notice that as long as the knowledge base does not change, PP, its squared vector norm, also remains constant. So, instead of recomputing it every time we can just do it once and then use the precomputed result further accelerating the distance computation.
Now let us put it all together.
Naturally, you could use this implementation with any vectorizer model, not just BERT. It is quite effective at nearest neighbour retrieval, being able to process dozens of requests per second on a 2-core Colab CPU.
For this example we will use a dataset of movie summaries from IMDB. Using the NLU and Retriever modules, we will build a movie recommendation system that suggests movies with similar plot features.
First, let’s download and prepare the IMDB dataset.
Vectorize movie plots with the BERT NLU module:
Finally, using the L2Retriever, find movies with plot vectors most similar to the query movie, and return it to user.
Let’s check it out!
>>> recommend = buildMovieRecommender(names, X_vect)>>> recommend("The Matrix")Impostor Immortel Saturn 3 Terminator Salvation The Terminator Logan's Run Genesis II Tron: Legacy Blade Runner
In this guide we have built a general-purpose BERT feature extractor. Models built with the features extracted from BERT perform adequately on classification and retrieval tasks. While their performance can be further improved by fine-tuning, the described approach to text feature extraction provides a solid unsupervised baseline for downstream NLP solutions.
Pre-training BERT from scratch with cloud TPUBuilding a Search Engine with BERT and Tensorflow [you are here]Fine-tuning BERT with Keras and tf.ModuleImproving sentence embeddings with BERT and Representation Learning
Pre-training BERT from scratch with cloud TPU
Building a Search Engine with BERT and Tensorflow [you are here]
Fine-tuning BERT with Keras and tf.Module
Improving sentence embeddings with BERT and Representation Learning