How to Use Cohere Embeddings and Rerank Modules With MongoDB Atlas

1
!pip install cohere==4.57 pymongo pandas s3fs

1
import os
2
import getpass
3
# cohere api key
4
try:
5
    cohere_api_key = os.environ["COHERE_API_KEY"]
6
except KeyError:
7
    cohere_api_key = getpass.getpass("Please enter your COHERE API KEY (hit enter): ")
8
9
# MongoDB connection string
10
try:
11
    MONGO_CONN_STR = os.environ["MONGODB_CONNECTION_STR"]
12
except KeyError:
13
    MONGO_CONN = getpass.getpass("Please enter your MongoDB Atlas Connection String (hit enter): ")

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

1
import cohere  
2
co_client = cohere.Client(cohere_api_key, client_name='mongodb')
3
# create a dataset in Cohere Platform
4
dataset = co_client.create_dataset(name='movies',
5
                                   data=open("./movies_sample_dataset.jsonl",'r'),
6
                                   keep_fields=["overview","title","year"],
7
                                   dataset_type="embed-input").wait()
8
dataset.wait()
9
dataset
10
11
dataset.wait()
12
# Schedule an Embedding job to run on the entire movies dataset
13
embed_job = co_client.create_embed_job(dataset_id=dataset.id, 
14
    input_type='search_document',
15
    model='embed-english-v3.0', 
16
    truncate='END')
17
embed_job.wait()
18
output_dataset = co_client.get_dataset(embed_job.output.id)
19
results = list(map(lambda x:{"text":x["text"], "embedding": x["embeddings"]["float"]},output_dataset))
20
len(results)

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

1
output_collection.create_search_index({"definition":
2
        {"mappings":
3
         {"dynamic": true,
4
          "fields": {
5
            "embedding" : {
6
                "dimensions": 1024,
7
                "similarity": "cosine",
8
                "type": "vector"
9
                },
10
            "fullplot":
11
            }}},
12
     "name": "default"
13
    }
14
)

1
def query_vector_search(q, prefilter = {}, postfilter = {},path="embedding",topK=2):
2
    ele = co_client.embed(model="embed-english-v3.0",input_type="search_query",texts=[q])
3
    query_embedding = ele.embeddings[0]
4
    vs_query = {
5
                "index": "default",
6
                "path": path,
7
                "queryVector": query_embedding,
8
                "numCandidates": 10,
9
                "limit": topK,
10
            }
11
    if len(prefilter)>0:
12
        vs_query["filter"] = prefilter
13
    new_search_query = {"$vectorSearch": vs_query}
14
    project = {"$project": {"score": {"$meta": "vectorSearchScore"},"_id": 0,"title": 1, "release_date": 1, "overview": 1,"year": 1}}
15
    if len(postfilter.keys())>0:
16
        postFilter = {"$match":postfilter}
17
        res = list(output_collection.aggregate([new_search_query, project, postFilter]))
18
    else:
19
        res = list(output_collection.aggregate([new_search_query, project]))
20
    return res

1
query_vector_search("romantic comedy movies", topK=5)

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

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

1
# sample search query using $search operator in aggregation pipeline
2
def query_fulltext_search(q,topK=25):
3
    v = {"$search": {
4
      "text": {
5
        "query": q,
6
        "path":"overview"
7
      }
8
    }}
9
    project = {"$project": {"score": {"$meta": "searchScore"},"_id": 0,"title": 1, "release-date": 1, "overview": 1}}
10
    docs = list(output_collection.aggregate([v,project, {"$limit":topK}]))
11
    return docs
12
# results before re ranking
13
docs = query_fulltext_search("romantic comedy movies", topK=10)
14
docs

1
# After passing the search results through the Cohere rerank module
2
q = "romantic comedy movies"
3
docs = query_fulltext_search(q)
4
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.
5
for idx, r in enumerate(results):
6
    print(f"Document Rank: {idx + 1}, Document Index: {r.index}")
7
    print(f"Document Title: {docs[r.index]['title']}")
8
    print(f"Document: {r.document['text']}")
9
    print(f"Relevance Score: {r.relevance_score:.2f}")
10
    print("\n")

1
Document Rank: 1, Document Index: 22
2
Document Title: Love Finds Andy Hardy
3
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.
4
Relevance Score: 0.99
5
6
7
Document Rank: 2, Document Index: 12
8
Document Title: Seventh Heaven
9
Document: Seventh Heaven or De zevende zemel is a 1993 Dutch romantic comedy film directed by Jean-Paul Lilienfeld.
10
Relevance Score: 0.99
11
12
13
Document Rank: 3, Document Index: 19
14
Document Title: Shared Rooms
15
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.
16
Relevance Score: 0.97
17
18
19
Document Rank: 4, Document Index: 3
20
Document Title: Too Many Husbands
21
Document: Romantic comedy adapted from a Somerset Maugham play.
22
Relevance Score: 0.97
23
24
25
Document Rank: 5, Document Index: 20
26
Document Title: Walking the Streets of Moscow
27
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.
28
Relevance Score: 0.96