End to End pipeline for building a semantic search engine

Image generated using Stable Diffusion 2

Searching for a piece of a sentence in a vast corpus of sentences can be quickly done by fuzzy matching or pattern matching. Suppose you have a query that says “captain of the Indian cricket team” You will get all the articles that contain this exact query, but this is not what we always want because, in some cases, we might search for “Rohit Sharma “, but still we would want to see results related to the Indian cricket team. Simple pattern matching cannot do this, so "Semantic Search" comes to the rescue for this use case.

So, what exactly is semantic search? How to implement it? How to make it work for a vast corpus of data? Be patient, and all of these questions will be answered in this blog.

Semantic Search

Firstly, a semantic search is a data-searching procedure that denotes a search with meaning. It is different from lexical search, where the search engine looks for keywords of the query or variants of them without understanding the overall meaning of the query. In semantic search, we try to understand the intent and contextual purpose of the words. Semantic search improves search accuracy by understanding the content of the search query, in contrast to conventional search engines, which only find records based on lexical matches.

The NLP behind search engine

While working with textual data, one of the most significant steps is to convert words or sentences into vectors. There are multiple techniques like “Bag of words” and “Tf-IDF”, but these techniques do not give much importance to the meaning of words and also, with the increase in the size of vocabulary, they might face the curse of dimensionality. So, in our use case to convert sentences into embeddings, we will use state-of-the-art sentence transformers.

Sentence transformers are the models used to generate dense vector embeddings for words, sentences, or even paragraphs. Multiple transformer models are available (like BERT, RoBERTa, DistilBERT etc.). Out of which, we will use DistilBERT

Training Pipeline

The first step in any NLP project usually is preprocessing of the text, this step is use-case specific, and thus I will leave it. Some common preprocessing steps include removing stop-words, removing numbers, removing abbreviations, and converting to lowercase. Etc.

Once we have preprocessed text, we will pass it to a transformer network to convert these sentences into embeddings.

We are using a pre-trained “distilbert-base-nli-mean-tokens” transformer from the hugging face “sentence_transformer” library it converts given sentences into a 768-dimensional vector. Once the entire corpus of sentences has been converted to vector embeddings, we will save it in any database.

Inference Pipeline

All the training steps have been completed. Now, whenever a new query comes, it will go through similar preprocessing steps, and then it will pass through the same transformer model, which will spit out a query vector of 768 Dimension.

The Search Pipeline

To find the vector from the saved corpus that is most similar to the query vector, we will calculate the cosine similarity between all the vectors in the corpus and the query vector. In data analysis, cosine similarity is a measure of similarity between two vectors; its values lie between [-1,1]. Two exactly similar vectors have a similarity of 1, and two exactly opposite vectors have a cosine similarity of -1.

With the above-described linear search method, we can easily find the most similar vector for the query vector. But is it efficient?

When dealing with a huge corpus of data (~10 M records), linear search is inefficient for exact matching by scanning the whole repository. If we are trying to get output in real-time, linear search cannot fulfill our demands due to the curse of dimensionality.

These problems make linear search almost impractical, and approximate nearest neighbour (ANN) comes into the picture.

What is ANN?

ANNs are a set of algorithmic techniques that aims to return the approximate nearest neighbour as fast as possible.

The neighbour returned by ANN might not be the exact nearest neighbour, but it is the optimal neighbour given the time constraint. The difference between KNN and ANN is that instead of searching all the training data points in the prediction phase, we only search for a small subset of candidate points.

How does ANN work?

With the increase in internet usage these days, a lot of information is consumed by users daily, so finding efficient ways of searching a query has always been a hot topic for search engines like Google and e-commerce giants like Amazon. Query searching plays an important role. Therefore, finding the nearest neighbour has become a hot research topic increasing the chance of users finding the information they seek in a reasonable time.

There are multiple techniques available for finding ANN. But the core of all the methods is they speed up the search by preprocessing the data into an efficient data structure, then dividing the entire vector space into multiple regions and then searching only in some of the areas instead of the entire vector space.

Some of the common preprocessing steps for ANN include.

· Vector transformation — Before indexing vector goes through different transformation techniques, e.g., dimensionality reduction and vector rotation.

· Vector Encoding — This step converts the vector into the actual index that can be searched. These techniques include data structure-based techniques like Trees, LSH and quantisation-based techniques to encode vectors to a much more compact form

· Non-Exhaustive Search — All the ANN technique uses this step in one way or another while searching index files. Instead of searching all the data points, we only searched in the probable region where there is the possibility of finding the nearest neighbour.

ANN reduces the search space by only comparing with a small subset of candidate points. This can be done by using Tree-Based algorithms (Like Annoy, ScaNN ) or LSH-based algorithms (like FAISS).

ANNOY is one of the most famous ANNs in the market, and it was developed by a team of researchers at Spotify. ANNOY is the algorithm that provides us with music recommendations based on songs we have previously played. You can read more about ANNOY **here.**

Source: https://erikbern.com/2015/10/01/nearest-neighbors-and-vector-models-part-2-how-to-search-in-high-dimensional-spaces.html

Source: https://erikbern.com/2015/10/01/nearest-neighbors-and-vector-models-part-2-how-to-search-in-high-dimensional-spaces.html

Basically, ANNOY works by creating many trees (forests) by picking two points at random and splitting the space into two parts by using hyperplane; we recursively keep splitting into subspaces until the number of data points in the node is small enough.

FAISS is another ANN technique developed by the Facebook AI team. The core of FAISS is Location-sensitive hashing.

Source: https://brc7.github.io/2019/09/19/Visual-LSH.html

In LSH, to construct the index, we apply multiple hash functions to map data points into buckets so that data points near each other are in the same buckets with high probability. In contrast, data points far from each other will likely fall into different buckets.

How to select which ANN algorithm to use?

The answer to this question varies from use case to use case. These points should be considered while selecting which ANN algorithm to use.

· Algorithm should give results as close as possible to the actual neighbour

· There might be use cases where the addition of new indexes and removing old indexes are important.

· Some algorithms don’t support user-defined IDs, which might be important for some use cases.

· There are multiple ways to calculate similarity scores, like Manhattan distance, L2 distance, cosine distance etc. Based on the use case, this might be important.

· GPU-based implementations are much faster than Only CPU-based ANN. Thus, depending on the availability, you can select GPU-based algorithms.

Performance Benchmarking

For Benchmarking purposes, we used a dataset of 1M records. Based on our use case, we will use FAISS-GPU for our semantic search engine.

Performance benchmarking of different ANN algorithms

Final Search Architecture

The final architecture of our search engine is going to look like this.

The architecture of a basic search engine

As everything is final, let's code it.

We will start with importing and installing the required libraries

!pip install faiss-cpu
!pip install -U sentence-transformers
import numpy as np
import torch
import os
import pandas as pd
import faiss
import time
from sentence_transformers import SentenceTransformer

Now import any text data, and apply some preprocessing on it.

df=pd.read_csv("any_data.csv")
df=preprocessing(df)

As per the architecture, convert the textual data into vector embeddings using a sentence transformer.

model = SentenceTransformer('distilbert-base-nli-mean-tokens')
encoded_data = model.encode(data)

Once embeddings have been generated, we will convert them into the FAISS index file and save them for future reference.

## dimension of embeddings
vector_dimension=768

## number of clusters in which data will be partitined
ncluster=5

quantizer = faiss.IndexFlatL2(vector_dimension)

## creating partition index
index = faiss.IndexIVFFlat(quantizer,vector_dimension,ncluster)

## always search in nearest 2 clusters
index.nprobe=2

index.train(encoded_data)
index.add(encoded_data)

## saving the index in local for future use
faiss.write_index(index,'index_file_name')

Now our training pipeline is done; whenever a new query comes, it will pass through similar steps to give us the most similar results.

## reading the saved index file
index = faiss.read_index('index_file_name')

## converting query into vector
search_query= preprocessing(search_query)
query_vector= model.encode([search_query])

## searching for nearest 5 neighbours in the index file
k= 5

## indices will contain index of nearest 5 neighbour
## distances will contain the similarity score of correspoding index
distances, indices = index.search(query_vector, k)

Voila!! Our Semantic Search engine is ready.

References

1-https://www.elastic.co/what-is/semantic-search

2-https://github.com/facebookresearch/faiss

3-https://github.com/spotify/annoy

4-https://huggingface.co/thenlper/gte-large