Adit Shrimal

Varun Hande

Neil Chitre

Manav Middha

In the rapidly expanding online retail industry, data is being generated at an unprecedented scale. To provide a seamless shopping experience to customers, it’s crucial to build an efficient, scalable product search architecture for storing and processing data. In this article, we will compare four popular natural language processing (NLP) techniques — namely BM25, TF-IDF, Word2Vec, and BERT — to empirically find the most optimal solution for retrieving the most relevant results for a search query from a large corpus of products.


## Where’s the code?

The full code for this project can be found [here](https://github.com/mongodb-developer/nlp-product-search).


## Data processing<a id="data-processing"></a>

Our data is sourced from Asos’s (an online retailer for clothing, footwear, accessories, etc.) real-time product data APIs. We utilized Apache Spark, Google Cloud Storage (GCS), and MongoDB to store and process data orchestrated by Apache Airflow.

Our data processing pipeline consists of three main parts:

### 1. Fetch data from APIs and write to GCS<a id="fetch-data-to-gcs"></a>

We first fetch the product categories. Then, we get the products in each category and store this data to Google Cloud Storage. The Python snippet below illustrates how we fetched the data using RapidAPI and stored the data to GCS:

```
import json
import requests

url = "https://asos2.p.rapidapi.com/categories/list"
headers = {
   "x-rapidapi-key": "YOUR_RAPIDAPI_KEY",
   "x-rapidapi-host": "asos2.p.rapidapi.com",
}
response = requests.request("GET", url, headers=headers)
categories = json.loads(response.text)
for category in categories:a
   url = f"https://asos2.p.rapidapi.com/products/v2/list/{category['id']}"
   response = requests.request("GET", url, headers=headers)
   products = json.loads(response.text)
   # Store products in GCS
   bucket = os.environ.get('GS_BUCKET_NAME')
   product_filename = f"products_{datetime.now().strftime('%H%m%S')}.json"
   write(bucket, str(date.today()) + "/" + product_filename, products)
```

### 2. Pre-processing<a id="pre-processing"></a>

The next step is preprocessing. Since product names are often not clean, we tokenized the data and removed stop words. For this, we used the `RegexTokenizer` and `StopWordsRemover` from SparkML.

```
from pyspark.ml.feature import RegexTokenizer, StopWordsRemover

# Instantiate and set input and output column parameters
regexTokenizer = RegexTokenizer(
   inputCol="product_name", outputCol="words", pattern="\\W"
)
remover = StopWordsRemover(inputCol="words", outputCol="filtered")

# Transform the data
regexTokenized = regexTokenizer.transform(df)
removedStopWords = remover.transform(regexTokenized)
```

### 3. Write transformed data to MongoDB<a id="write-data-to-mongodb"></a>

The last step is to store the transformed documents in a MongoDB collection. MongoDB’s document model, ability to efficiently handle large data volumes, built-in full-text search functionality, and easy integrations with big data tools such as PySpark make it a fitting choice for our data processing pipeline. The following code snippet shows how to write data to MongoDB:

```
import json

from pymongo import MongoClient

client = MongoClient(
   "mongodb://<username>:<password>@cluster0.mongodb.net/test?retryWrites=true&w=majority"
)
db = client.get_database("product_db")
products = db.products

# Fetch data from GCS
for product in gcs_data:
   products.insert_one(product)
```


## Experimenting with different NLP techniques<a id="experimenting-with-nlp-techniques"></a>

Our primary objective was to find the most similar products for a given search query. We experimented with several methodologies:

### 1. TF-IDF + cosine similarity<a id="tfidf-cosine"></a>

[TF-IDF](https://en.wikipedia.org/wiki/Tf%E2%80%93idf), short for Term Frequency–Inverse Document Frequency, is a measure of importance of a word to a document in a corpus. It is a widely used method to encode text. Given a user query, we first encode it using TF-IDF and then use cosine similarity to find the most similar products from our catalog.

```
from pyspark.ml.feature import IDF, HashingTF

hashingTF = HashingTF(inputCol="filtered", outputCol="rawFeatures", numFeatures=20)
featurizedData = hashingTF.transform(removedStopWords)

idf = IDF(inputCol="rawFeatures", outputCol="features")
idfModel = idf.fit(featurizedData)
rescaledData = idfModel.transform(featurizedData)
```

### 2. Word2Vec + cosine similarity<a id="word2vec-cosine"></a>

[Word2Vec](https://en.wikipedia.org/wiki/Word2vec) is a family of model architectures to learn word embeddings from large datasets. Here, we use the Word2Vec implementation in SparkML to embed the user query and then use cosine similarity to find the most similar products from our catalog.

```
from pyspark.ml.feature import Word2Vec

word2Vec = Word2Vec(
   vectorSize=300, minCount=0, inputCol="filtered", outputCol="features"
)
model = word2Vec.fit(removedStopWords)
result = model.transform(removedStopWords)
```

### 3. BM25<a id="bm25"></a>

BM25, short for “Best Match 25,” is a ranking function used in information retrieval systems to rank documents based on their relevance to a given search query. It is a probabilistic ranking models and is known for its ability to balance term frequency (TF) and inverse document frequency (IDF), thus making it a more sophisticated version of TF-IDF.

In our project, we implemented BM25 from scratch in PySpark due to the lack of a direct implementation in the library. Here is a more detailed look at our approach and code:

Firstly, we calculated the IDF values for each term, which is a measure of how much information the word provides, i.e., if it’s common or rare across all documents.

Secondly, we computed the term frequency in each document. Term frequency is simply the number of times a word appears in a document.

With these two metrics, we could then compute the BM25 score for each document relative to a query. Below is a simplified pseudocode for the process:

```
# Compute IDF values for each term
# Compute term frequency in each document for each term
# For each query term, do:
# For each document containing the term do:
# Compute BM25 Score
# Sort documents by score
```

### Building an inverted index<a id="building-an-inverted-index"></a>

In order to enhance the time efficiency of calculating TF and IDF every time for a product, we created an inverted index. This inverted index saves these important values beforehand to improve the retrieval time when someone searches for a product.

The first part of preprocessing involves stemming, which is the process of reducing inflected (or sometimes derived) words to their word stem, base, or root form. It’s done to get to the basic part of a word that carries the meaning. For example, the word stem of “running” is “run.” In this work, we used the [Porter Stemming algorithm](https://www.nltk.org/howto/stem.html), a popular and robust stemming algorithm.

```
def stemming(words):
   ps = PorterStemmer()
   result = []

   for w in words:
       root_word = ps.stem(w)
       result.append(root_word)

   return result

stemming_udf = udf(lambda x: stemming(x), ArrayType(elementType=StringType()))
preprocessed_data = preprocessed_data.withColumn(
"tokens", stemming_udf("filtered_words")
)
```

Next, we flatten the list of tokens into individual words and count the total number of unique words and total words.

```
words = preprocessed_data.select(explode("tokens").alias("word"))
total_unique_words = words.select("word").distinct().count()
total_words = words.count()
```

Once we have a list of words, we start building the inverted index. For each word in our dataset, we create an entry in our index, where each entry points to a list of documents that contain the word, along with the word’s frequency in each document.

```
id_str = concat_ws("", col("_id").cast(StringType()))

word_count_df = (
   preprocessed_data.select("_id", explode("tokens").alias("word"))
   .groupBy("word", "_id")
   .agg(size(collect_list("word")).alias("occurrence"))
   .groupBy("word")
   .agg(
       map_from_entries(collect_list(struct(id_str.alias("id"), "occurrence"))).alias(
           "documents"
       )
   )
   .withColumnRenamed("word", "_id")
).rdd.map(lambda row: (row["_id"], row["documents"]))

inverted_index = word_count_df.collectAsMap()
```

Finally, we also compute the length of all documents and the unique word count for each document. These values are used later in the BM25 formula.

```
doc_length_df = (
   preprocessed_data.select("_id", size("tokens").alias("length"))
   .groupBy("_id")
   .agg(sum(col("length")).alias("doc_length"))
   .withColumn("_id", concat_ws("", col("_id").cast(StringType())))
   .rdd.map(lambda row: (row["_id"], row["doc_length"]))
)
doc_length = doc_length_df.collectAsMap()

doc_length["word_counter"] = total_unique_words
doc_length["doc_counter"] = total_words

unique_word_count_df = preprocessed_data.select(
   "_id",
   size(array_distinct("tokens")).alias("num_unique_words"),
).withColumn("_id", concat_ws("", col("_id").cast(StringType()))).rdd.map(lambda row: (row["_id"], row["num_unique_words"]))
unique_word = unique_word_count_df.collectAsMap()
```

This concludes the preprocessing and inverted index-building step. The inverted index is a critical component of the BM25 ranking function, as it enables fast retrieval of relevant documents for a given query.

Now, let’s dive into the implementation of the BM25 algorithm.

### Implementing BM25<a id="implementing-bm25"></a>

The mathematical formula for BM25 is as follows:

![BM25 mathematical formula][1]

Where:
* f(q_i, D): Frequency of query term q_i in document D
* |D|: Word count of document D
* avgdl: Average word count across all documents
* k_1 and b: Tuning parameters, usually set as k_1 in [1.2, 2.0] and b=0.75
* IDF(q_i): Weight of query term q_i based on its rarity across all documents

[Source](https://en.wikipedia.org/wiki/Okapi_BM25)

Here’s how we’ve implemented the above algorithm in Python:

```
def bm25(doc_length,
  inverted_index, 
   unique_word, 
   word_counter, 
   doc_counter, 
   query):
   k1 = 1.2
   b = 0.75
   k2 = 100
   N = len(doc_length) - 2
   avgdl = doc_length["doc_counter"] / N

   score = defaultdict(int)
   for word in query:
       if word in inverted_index:
           n = len(inverted_index[word])
           idf = math.log((N - n + 0.5) / (n + 0.5))
           for doc, freq in inverted_index[word].items():
               f = freq
               qf = query.count(word)
               dl = doc_length[doc]
               K = k1 * ((1 - b) + b * (float(dl) / float(avgdl)))
               R1 = idf * ((f * (k1 + 1)) / (f + K))
               R2 = (qf * (k2 + 1)) / (qf + k2)
               R = R1 * R2
               score[doc] += R
       else:
           Continue

   return score
```

This function takes the document length dictionary, the inverted index, the unique word count dictionary, the total word counter, the document counter, and the search query. It calculates the BM25 score for each document in relation to the query and returns a dictionary where the keys are the document IDs, and the values are the corresponding BM25 scores. The documents with the highest BM25 scores are the ones most relevant to the query. In this study, we chose values for k1, b, and k2 based on other existing studies.

### 4. BERT<a id="bert"></a>

[BERT](https://en.wikipedia.org/wiki/BERT_(language_model)), short for Bidirectional Encoder Representations from Transformers, was one of the early examples of large language models (LLMs) and uses a transformer encoder architecture to represent text as vectors. We used a pre-trained BERT model to map the user query to a high-dimensional vector space and then use cosine similarity to find the most similar products from our catalog.

```
from sentence_transformers import SentenceTransformer

model = SentenceTransformer('bert-base-nli-mean-tokens')
sentence_embeddings = model.encode(data['product_name'])
```


## Results

Each of the models had their strengths and weaknesses. While TF-IDF was the easiest to interpret, it struggled to capture the semantic meaning of the data. Word2Vec was limited to handling queries present in the training corpus. BM25 was very quick, but it lacked semantic understanding. Among the algorithms we tested, the best model was BERT, as it could incorporate the semantic meaning of the data and handle queries in multiple languages.

For our project, we used a GPU-accelerated compute instance on AWS. Below are the details of the specifications we used:
* Databricks Runtime version: 12.1 ML (includes [Apache Spark](https://www.whizlabs.com/blog/databricks-certified-associate-developer-apache-spark/) 3.3.1, GPU, Scala 2.12)
* Worker/Driver types: g4dn.xlarge
* Number of workers: minimum 2, maximum 5
* Resources per worker: 32–80 GB Memory, 8–20 Cores
* Driver resources: 16 GB Memory, 4 Cores

We compared the performance of the four search algorithms (BM25, TF-IDF, Word2Vec, and BERT) on three distinct queries: "sneakers", "t-shirt", and "chappals" (which means sandals/slippers in Hindi). The results varied in terms of speed and relevance for each query and algorithm combination. A summary of the results is as follows:

* **Sneakers query**:  BM25 was the fastest in returning relevant results. Both TF-IDF and Word2Vec were slower compared to BM25, while BERT offered a balance between speed and relevance.

* **T-shirt query**:  While BM25 was the fastest to return results, it provided irrelevant results due to the limited vocabulary of the corpus. Word2Vec performed poorly in terms of both speed and relevance, similar to TF-IDF. BERT was moderately slow but provided highly relevant results.

* **Chappals query**: Chappal is a Hindi word for leather sandals. Again, BM25 excelled in terms of speed but delivered irrelevant results due to language limitations in the corpus. Word2Vec and TF-IDF had similar performance metrics as with the "t-shirt" query. BERT again demonstrated its effectiveness by delivering highly relevant results even with the complexity of handling multiple languages.

![TD-IDF results on “chappals” query][2]

TF-IDF's results for the “chappals” query show the algorithm's struggle with multi-language queries.

![ Word2Vec results on “chappals” query][3]

Word2Vec results demonstrate its limitations in handling diverse languages.

![BM25 results on “chappals” query][4]

BM25 responds quickly but produces irrelevant results due to language limitations in the corpus.

![BERT results on “chappals” query][5]

BERT successfully retrieves relevant results, showcasing its strength in handling multi-language queries.


## Conclusion

In this article, we compared four different natural language processing (NLP) techniques — namely BM25, TF-IDF, Word2Vec, and BERT — to empirically find the most optimal solution for retrieving the most relevant results for a search query from a large corpus of products. 

Our preliminary experiments show that the BM25 algorithm is the most time-efficient, while BERT provides the best trade-off between speed and relevance. This observation is in-line with what we see in practice in the real world, with BM25 being a popular algorithm for lexical search, and vector search being the preferred choice for semantic search. Hybrid search, a technique to combine the strengths of keyword and semantic search, is becoming increasingly common. 

MongoDB Atlas supports lexical, vector, as well as hybrid search. Here are some resources to get you started:
* [Lexical search](https://www.mongodb.com/docs/atlas/atlas-search/tutorial/?tck=as_web_get_started)
* [Vector search](https://www.mongodb.com/docs/atlas/atlas-vector-search/tutorials/vector-search-quick-start/?tck=ai_as_web)
* [Hybrid search](https://www.mongodb.com/docs/atlas/atlas-vector-search/tutorials/reciprocal-rank-fusion/)

For further questions, reach out to us in our [community forums](https://www.mongodb.com/community/forums/)!


  [1]: https://images.contentstack.io/v3/assets/blt39790b633ee0d5a7/blt4aac5c0f6ca6b699/66c66a74e712efbfe231261f/image5.png
  [2]: https://images.contentstack.io/v3/assets/blt39790b633ee0d5a7/blte07874680978fce1/66c66c4aab1b6925953cbc99/image3.png
  [3]: https://images.contentstack.io/v3/assets/blt39790b633ee0d5a7/blt8c0675665aa68c07/66c66c74038881647b1f04a0/image1.png
  [4]: https://images.contentstack.io/v3/assets/blt39790b633ee0d5a7/blt721d9a8b57c4d82f/66c66c9703888150a01f04a5/image4.png
  [5]: https://images.contentstack.io/v3/assets/blt39790b633ee0d5a7/blt7f9b1c834e0f5d1f/66c66cdce712ef296d312639/image2.png

In this article, we will compare four popular natural language processing (NLP) techniques to find the most optimal solution for retrieving the most relevant results for a search query from a large corpus of products. 

Varun Hande

Content by Varun Hande

Comparing NLP Techniques for Scalable Product Search