Building a RAG System With Google's Gemma, Hugging Face and MongoDB

1
!pip install datasets pandas pymongo sentence_transformers
2
!pip install -U transformers
3
# Install below if using GPU
4
!pip install accelerate

1
# Load Dataset
2
from datasets import load_dataset
3
import pandas as pd
4
# https://huggingface.co/datasets/MongoDB/embedded_movies
5
dataset = load_dataset("MongoDB/embedded_movies")
6
# Convert the dataset to a pandas DataFrame
7
dataset_df = pd.DataFrame(dataset['train'])

1
# Remove data point where plot column is missing
2
dataset_df = dataset_df.dropna(subset=['fullplot'])
3
print("\nNumber of missing values in each column after removal:")
4
print(dataset_df.isnull().sum())
5
6
# Remove the plot_embedding from each data point in the dataset as we are going to create new embeddings with an open-source embedding model from Hugging Face: gte-large
7
dataset_df = dataset_df.drop(columns=['plot_embedding'])

1
    from sentence_transformers import SentenceTransformer
2
    # https://huggingface.co/thenlper/gte-large
3
    embedding_model = SentenceTransformer("thenlper/gte-large")
4
5
    def get_embedding(text: str) -> list[float]:
6
        if not text.strip():
7
            print("Attempted to get embedding for empty text.")
8
            return []
9
10
        embedding = embedding_model.encode(text)
11
12
        return embedding.tolist()
13
14
    dataset_df["embedding"] = dataset_df["fullplot"].apply(get_embedding)

1
import pymongo
2
from google.colab import userdata
3
4
def get_mongo_client(mongo_uri):
5
    """Establish connection to the MongoDB."""
6
    try:
7
        client = pymongo.MongoClient(mongo_uri)
8
        print("Connection to MongoDB successful")
9
        return client
10
    except pymongo.errors.ConnectionFailure as e:
11
        print(f"Connection failed: {e}")
12
        return None
13
14
mongo_uri = userdata.get("MONGO_URI")
15
if not mongo_uri:
16
    print("MONGO_URI not set in environment variables")
17
18
mongo_client = get_mongo_client(mongo_uri)
19
20
# Ingest data into MongoDB
21
db = mongo_client["movies"]
22
collection = db["movie_collection_2"]

1
# Delete any existing records in the collection
2
collection.delete_many({})

1
{
2
 "fields": [{
3
     "numDimensions": 1024,
4
     "path": "embedding",
5
     "similarity": "cosine",
6
     "type": "vector"
7
   }]
8
}

1
documents = dataset_df.to_dict('records')
2
collection.insert_many(documents)
3
print("Data ingestion into MongoDB completed")

1
def vector_search(user_query, collection):
2
    """
3
    Perform a vector search in the MongoDB collection based on the user query.
4
5
    Args:
6
    user_query (str): The user's query string.
7
    collection (MongoCollection): The MongoDB collection to search.
8
9
    Returns:
10
    list: A list of matching documents.
11
    """
12
13
    # Generate embedding for the user query
14
    query_embedding = get_embedding(user_query)
15
16
    if query_embedding is None:
17
        return "Invalid query or embedding generation failed."
18
19
    # Define the vector search pipeline
20
    pipeline = [
21
        {
22
            "$vectorSearch": {
23
                "index": "vector_index",
24
                "queryVector": query_embedding,
25
                "path": "embedding",
26
                "numCandidates": 150,  # Number of candidate matches to consider
27
                "limit": 4,  # Return top 4 matches
28
            }
29
        },
30
        {
31
            "$project": {
32
                "_id": 0,  # Exclude the _id field
33
                "fullplot": 1,  # Include the plot field
34
                "title": 1,  # Include the title field
35
                "genres": 1,  # Include the genres field
36
                "score": {"$meta": "vectorSearchScore"},  # Include the search score
37
            }
38
        },
39
    ]
40
41
    # Execute the search
42
    results = collection.aggregate(pipeline)
43
    return list(results)

1
def get_search_result(query, collection):
2
3
    get_knowledge = vector_search(query, collection)
4
5
    search_result = ""
6
    for result in get_knowledge:
7
        search_result += f"Title: {result.get('title', 'N/A')}, Plot: {result.get('fullplot', 'N/A')}\n"
8
9
    return search_result

1
# Conduct query with retrieval of sources
2
query = "What is the best romantic movie to watch and why?"
3
source_information = get_search_result(query, collection)
4
combined_information = f"Query: {query}\nContinue to answer the query by using the Search Results:\n{source_information}."
5
print(combined_information)

1
from transformers import AutoTokenizer, AutoModelForCausalLM
2
3
tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it")
4
# CPU Enabled uncomment below 👇🏽
5
# model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it")
6
# GPU Enabled use below 👇🏽
7
model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it", device_map="auto")

1
# Moving tensors to GPU
2
input_ids = tokenizer(combined_information, return_tensors="pt").to("cuda")
3
response = model.generate(**input_ids, max_new_tokens=500)
4
print(tokenizer.decode(response[0]))