# How to Use Cohere Embeddings and Rerank Modules with MongoDB Atlas

The daunting task that developers currently face while developing solutions powered by the retrieval augmented generation (RAG) framework is the choice of retrieval mechanism. Augmenting the large language model (LLM) prompt with relevant and exhaustive information creates better responses from such systems.. One is tasked with choosing the most appropriate embedding model in the case of semantic similarity search. Alternatively, in the case of full-text search implementation, you have to be thorough about your implementation to achieve a precise recall and high accuracy in your results. Sometimes, the solutions require a combined implementation that benefits from both retrieval mechanisms.

If your current full-text search scoring workflow is leaving things to be desired, or if you find yourself spending too much time writing numerous lines of code to get semantic search functionality working within your applications, then Cohere and MongoDB can help. To prevent these issues from holding you back from leveraging powerful AI search functionality or machine learning within your application, Cohere and MongoDB offer easy-to-use and fully managed solutions.

Cohere is an AI company specializing in large language models.

1. With a powerful tool for embedding natural language in their projects, it can help you represent more accurate, relevant, and engaging content as embeddings. The Cohere language model also offers a simple and intuitive API that allows you to easily integrate it with your existing workflows and platforms. 
2. The Cohere Rerank module is a component of the Cohere natural language processing system that helps to select the best output from a set of candidates. The module uses a neural network to score each candidate based on its relevance, semantic similarity, theme, and style. The module then ranks the candidates according to their scores and returns the top N as the final output.

MongoDB Atlas is a fully managed developer data platform service that provides scalable, secure, and reliable data storage and access for your applications. One of the key features of MongoDB Atlas is the ability to perform vector search and full-text search on your data, which can enhance the capabilities of your AI/ML-driven applications. MongoDB Atlas can help you build powerful and flexible AI/ML-powered applications that can leverage both structured and unstructured data. You can easily create and manage search indexes, perform queries, and analyze results using MongoDB Atlas's intuitive interface, APIs, and drivers. MongoDB Atlas Vector Search provides a unique feature — pre-filtering and post-filtering on vector search queries — that helps users control the behavior of their vector search results, thereby improving the accuracy and retrieval performance, and saving money at the same time.

Therefore, with Cohere and MongoDB Atlas, we can demonstrate techniques where we can easily power a semantic search capability on your private dataset with very few lines of code. Additionally, you can enhance the existing ranking of your full-text search retrieval systems using the Cohere Rerank module. Both techniques are highly beneficial for building more complex GenAI applications, such as RAG- or LLM-powered summarization or data augmentation.

## What will we do in this tutorial?

### Store embeddings and prepare the index

1. Use the Cohere Embed Jobs to generate vector embeddings for the first time on large datasets in an asynchronous and scheduled manner.
2. Add vector embeddings into MongoDB Atlas, which can store and index these vector embeddings alongside your other operational/metadata. 
3. Finally, prepare the indexes for both vector embeddings and full-text search on our private dataset.

### Search with vector embeddings

1. Write a simple Python function to accept search terms/phrases and pass it through the Cohere embed API again to get a query vector.
2. Take these resultant query vector embeddings and perform a vector search query using the $vectorsearch operator in the MongoDB Aggregation Pipeline.
3. Pre-filter documents using meta information to narrow the search across your dataset, thereby speeding up the performance of vector search results while retaining accuracy.
4. The retrieved semantically similar documents can be post-filtered (relevancy score) to demonstrate a higher degree of control over the semantic search behaviour.

### Search with text and Rerank with Cohere

1. Write a simple Python function to accept search terms/phrases and prepare a query using the $search operator and MongoDB Aggregation Pipeline.
2. Take these resultant documents and perform a reranking operation of the retrieved documents to achieve higher accuracy with full-text search results using the Cohere rerank module.

- Cohere CLI tool

Also, if you have not created a MongoDB Atlas instance for yourself, you can follow the tutorial to create one. This will provide you with your `MONGODB_CONNECTION_STR`. 

Run the following lines of code in Jupyter Notebook to initialize the Cohere secret or API key and MongoDB Atlas connection string.

```python
import os
import getpass
# cohere api key
try:
    cohere_api_key = os.environ"COHERE_API_KEY"]
except KeyError:
    cohere_api_key = getpass.getpass("Please enter your COHERE API KEY (hit enter): ")

# MongoDB connection string
try:
    MONGO_CONN_STR = os.environ["MONGODB_CONNECTION_STR"]
except KeyError:
    MONGO_CONN = getpass.getpass("Please enter your MongoDB Atlas Connection String (hit enter): ")
```

### Load dataset from the S3 bucket

Run the following lines of code in Jupyter Notebook to read data from an AWS S3 bucket directly to a pandas dataframe.

```python
import pandas as pd
import s3fs
df = pd.read_json("s3://ashwin-partner-bucket/cohere/movies_sample_dataset.jsonl", orient="records", lines=True)
df.to_json("./movies_sample_dataset.jsonl", orient="records", lines=True)
df[:3]
```

![Loaded AWS S3 Dataset][2]

### Initialize and schedule the Cohere embeddings job to embed the "sample_movies" dataset

Here we will create a movies dataset in Cohere by uploading our sample movies dataset that we fetched from the S3 bucket and have stored locally. Once we have created a dataset, we can use the Cohere embed jobs API to schedule a batch job to embed all the entire dataset.

You can run the following lines of code in your Jupyter Notebook to upload your dataset to Cohere and schedule an embedding job.

```python
import cohere  
co_client = cohere.Client(cohere_api_key, client_name='mongodb')
# create a dataset in Cohere Platform
dataset = co_client.create_dataset(name='movies',
                                   data=open("./movies_sample_dataset.jsonl",'r'),
                                   keep_fields=["overview","title","year"],
                                   dataset_type="embed-input").wait()
dataset.wait()
dataset

dataset.wait()
# Schedule an Embedding job to run on the entire movies dataset
embed_job = co_client.create_embed_job(dataset_id=dataset.id, 
    input_type='search_document',
    model='embed-english-v3.0', 
    truncate='END')
embed_job.wait()
output_dataset = co_client.get_dataset(embed_job.output.id)
results = list(map(lambda x:{"text":x["text"], "embedding": x["embeddings"]["float"]},output_dataset))
len(results)
```

### How to initialize MongoDB Atlas and insert data to a MongoDB collection

Now that we have created the vector embeddings for our sample movies dataset, we can initialize the MongoDB client and insert the documents into our collection of choice by running the following lines of code in the Jupyter Notebook.

```python
from pymongo import MongoClient
mongo_client = MongoClient(MONGO_CONN_STR)
# Upload documents along with vector embeddings to MongoDB Atlas Collection
output_collection = mongo_client["sample_mflix"]["cohere_embed_movies"]
if output_collection.count_documents({})>0:
    output_collection.delete_many({})
e = output_collection.insert_many(results)
```

### Programmatically create vector search and full-text search index

With the latest update to the **Pymongo** Python package, you can now create your vector search index as well as full-text search indexes from the Python client itself. You can also create vector indexes using the MongoDB Atlas UI or `mongosh`.

Run the following lines of code in your Jupyter Notebook to create search and vector search indexes on your new collection.

```
output_collection.create_search_index({"definition":
        {"mappings":
         {"dynamic": true,
          "fields": {
            "embedding" : {
                "dimensions": 1024,
                "similarity": "cosine",
                "type": "vector"
                },
            "fullplot":
            }}},
     "name": "default"
    }
)
```

### Query MongoDB vector index using $vectorSearch

MongoDB Atlas brings the flexibility of using vector search alongside full-text search filters. Additionally, you can apply range, string, and numeric filters using the aggregation pipeline. This allows the end user to control the behavior of the semantic search response from the search engine. The below lines of code will demonstrate how you can perform vector search along with pre-filtering on the **year** field to get movies earlier than **1990.** Plus, you have better control over the relevance of returned results, so you can perform post-filtering on the response using the MongoDB Query API. In this demo, we are filtering on the **score** field generated as a result of performing the vector similarity between the query and respective documents, using a heuristic to retain only the accurate results.

Run the below lines of code in Jupyter Notebook to initialize a function that can help you achieve **vector search + pre-filter + post-filter**.

```python
def query_vector_search(q, prefilter = {}, postfilter = {},path="embedding",topK=2):
    ele = co_client.embed(model="embed-english-v3.0",input_type="search_query",texts=[q])
    query_embedding = ele.embeddings[0]
    vs_query = {
                "index": "default",
                "path": path,
                "queryVector": query_embedding,
                "numCandidates": 10,
                "limit": topK,
            }
    if len(prefilter)>0:
        vs_query["filter"] = prefilter
    new_search_query = {"$vectorSearch": vs_query}
    project = {"$project": {"score": {"$meta": "vectorSearchScore"},"_id": 0,"title": 1, "release_date": 1, "overview": 1,"year": 1}}
    if len(postfilter.keys())>0:
        postFilter = {"$match":postfilter}
        res = list(output_collection.aggregate([new_search_query, project, postFilter]))
    else:
        res = list(output_collection.aggregate([new_search_query, project]))
    return res
```

#### Vector search query example

Run the below lines of code in Jupyter Notebook cell and you can see the following results.

```python
query_vector_search("romantic comedy movies", topK=5)
```

![Vector Search Query Example Results][3]

#### Vector search query example with prefilter

```python
query_vector_search("romantic comedy movies", prefilter={"year":{"$lt": 1990}}, topK=5)
```

![Vector Search with Prefilter Example Results][4]

#### Vector search query example with prefilter and postfilter to control the semantic search relevance and behaviour

```python
query_vector_search("romantic comedy movies", prefilter={"year":{"$lt": 1990}}, postfilter={"score": {"$gt":0.76}},topK=5)
```

![Vector Search with Prefilter and Postfilter Example Results][5]

### Leverage MongoDB Atlas full-text search with Cohere Rerank module

[Cohere Rerank is a module in the Cohere suite of offerings that enhances the quality of search results by leveraging semantic search. This helps elevate the traditional search engine performance, which relies solely on keywords. Rerank goes a step further by ranking results retrieved from the search engine based on their semantic relevance to the input query. This pass of re-ranking search results helps achieve more appropriate and contextually similar search results.

To demonstrate how the Rerank module can be leveraged with MongoDB Atlas full-text search, we can follow along by running the following line of code in your Jupyter Notebook.

```python
# sample search query using $search operator in aggregation pipeline
def query_fulltext_search(q,topK=25):
    v = {"$search": {
      "text": {
        "query": q,
        "path":"overview"
      }
    }}
    project = {"$project": {"score": {"$meta": "searchScore"},"_id": 0,"title": 1, "release-date": 1, "overview": 1}}
    docs = list(output_collection.aggregate(v,project, {"$limit":topK}]))
    return docs
# results before re ranking
docs = query_fulltext_search("romantic comedy movies", topK=10)
docs
```

![Cohere Rerank Model Sample Results][6]

```python
# After passing the search results through the Cohere rerank module
q = "romantic comedy movies"
docs = query_fulltext_search(q)
results = co_client.rerank(query=q, documents=list(map(lambda x:x["overview"], docs)), top_n=5, model='rerank-english-v2.0') # Change top_n to change the number of results returned. If top_n is not passed, all results will be returned.
for idx, r in enumerate(results):
    print(f"Document Rank: {idx + 1}, Document Index: {r.index}")
    print(f"Document Title: {docs[r.index]['title']}")
    print(f"Document: {r.document['text']}")
    print(f"Relevance Score: {r.relevance_score:.2f}")
    print("\n")
```

Output post reranking the full-text search results:

```
Document Rank: 1, Document Index: 22
Document Title: Love Finds Andy Hardy
Document: A 1938 romantic comedy film which tells the story of a teenage boy who becomes entangled with three different girls all at the same time.
Relevance Score: 0.99

Document Rank: 2, Document Index: 12
Document Title: Seventh Heaven
Document: Seventh Heaven or De zevende zemel is a 1993 Dutch romantic comedy film directed by Jean-Paul Lilienfeld.
Relevance Score: 0.99

Document Rank: 3, Document Index: 19
Document Title: Shared Rooms
Document: A new romantic comedy feature film that brings together three interrelated tales of gay men seeking family, love and sex during the holiday season.
Relevance Score: 0.97

Document Rank: 4, Document Index: 3
Document Title: Too Many Husbands
Document: Romantic comedy adapted from a Somerset Maugham play.
Relevance Score: 0.97

Document Rank: 5, Document Index: 20
Document Title: Walking the Streets of Moscow
Document: "I Am Walking Along Moscow" aka "Ya Shagayu Po Moskve" (1963) is a charming lyrical comedy directed by Georgi Daneliya in 1963 that was nominated for Golden Palm at Cannes Film Festival. Daneliya proved that it is possible to create a masterpiece in the most difficult genre of romantic comedy. Made by the team of young and incredibly talented artists that besides Daneliya included writer/poet Gennady Shpalikov, composer Andrei Petrov, and cinematographer Vadim Yusov (who had made four films with Andrei Tarkovski), and the dream cast of the talented actors even in the smaller cameos, "I Am Walking Along Moscow" keeps walking victoriously through the decades remaining deservingly one of the best and most beloved Russian comedies and simply one of the best Russian movies ever made. Funny and gentle, dreamy and humorous, romantic and realistic, the film is blessed with the eternal youth and will always take to the walk on the streets of Moscow new generations of the grateful viewers.
Relevance Score: 0.96
```

## Summary

In this tutorial, we were able to demonstrate the following:

1. Using the Cohere embedding along with MongoDB Vector Search, we were able to show how easy it is to achieve semantic search functionality alongside your operational data functions.
2. With Cohere Rerank, we were able to search results using full-text search capabilities in MongoDB and then rank them by semantic relevance, thereby delivering richer, more relevant results without replacing your existing search architecture setup.
3. The implementations were achieved with minimal lines of code and showcasing ease of use.
4. Leveraging Cohere Embeddings and Rerank does not need a team of ML experts to develop and maintain. So the monthly costs of maintenance were kept to a minimum.
5. Both solutions are cloud-agnostic and, hence, can be set up on any cloud platform.

The same can be found on a [notebook which will help reduce the time and effort following the steps in this blog.

## What's next?

To learn more about how MongoDB Atlas is helping build application-side ML integration in real-world applications, you can visit the MongoDB for AI page.

  [1]: https://images.contentstack.io/v3/assets/blt39790b633ee0d5a7/blte8f8f2d8681106dd/660c5dfcdd5b9e752ba8949a/1.png
  [2]: https://images.contentstack.io/v3/assets/blt39790b633ee0d5a7/blt11b31c83a7a30a85/660c5e236c4a398354e46705/2.png
  [3]: https://images.contentstack.io/v3/assets/blt39790b633ee0d5a7/bltf09db7ce89c89f05/660c5e4a3110d0a96d069608/3.png
  [4]: https://images.contentstack.io/v3/assets/blt39790b633ee0d5a7/blt5707998b8d57764c/660c5e75c3bc8bfdfbdd1fc1/4.png
  [5]: https://images.contentstack.io/v3/assets/blt39790b633ee0d5a7/blt533d00bfde1ec48f/660c5e94c3bc8b26dedd1fcd/5.png
  [6]: https://images.contentstack.io/v3/assets/blt39790b633ee0d5a7/bltc67a9ac477d5029e/660c5eb0df43aaed1cf11e70/6.png