Implementing Robust RAG Pipelines: Integrating Google's Gemma 2 (2B) Open Model, MongoDB, and LLM Evaluation Techniques

1
!pip install --upgrade --quiet datasets pandas pymongo sentence_transformers deepevals

1
import os
2
import getpass
3
# Make sure you have an Hugging Face token(HF_TOKEN) in your development environment before running the code below
4
# How to get a token: https://huggingface.co/docs/hub/en/security-tokens
5
os.environ["HF_TOKEN"] = getpass.getpass("Enter your Hugging Face token: ")
6
os.environ["OPENAI_API_KEY"] = getpass.getpass("Enter your OpenAI API key: ")  # Need this for the evaluation step

1
import pandas as pd
2
from datasets import load_dataset
3
4
# Make sure you have an Hugging Face token(HF_TOKEN) in your development environment before running the code below
5
# How to get a token: https://huggingface.co/docs/hub/en/security-tokens
6
# https://huggingface.co/datasets/MongoDB/fake_tech_companies_market_reports
7
dataset = load_dataset("MongoDB/fake_tech_companies_market_reports", split="train", streaming=True)
8
dataset_df = dataset.take(100)
9
10
# Convert the dataset to a pandas dataframe
11
dataset_df = pd.DataFrame(dataset_df)
12
dataset_df.head(5)

1
# Data Preparation
2
def combine_attributes(row):
3
  combined = f"{row['company']} {row['sector']} "
4
5
  # Add report information
6
  for report in row['reports']:
7
    combined += f"{report['year']} {report['title']} {report['author']} {report['content']} "
8
9
  # Add recent news information
10
  for news in row['recent_news']:
11
    combined += f"{news['headline']} {news['summary']} "
12
13
  return combined.strip()

1
# Add the new column 'combined_attributes'
2
dataset_df['combined_attributes'] = dataset_df.apply(combine_attributes, axis=1)

1
from sentence_transformers import SentenceTransformer
2
from tqdm import tqdm
3
import numpy as np
4
5
6
# https://huggingface.co/thenlper/gte-large
7
# embedding_model = SentenceTransformer('thenlper/gte-large')
8
embedding_model = SentenceTransformer('Alibaba-NLP/gte-large-en-v1.5', trust_remote_code=True)
9
10
11
# Determine the maximum sequence length for the model
12
max_seq_length = embedding_model.max_seq_length
13
14
15
16
17
def chunk_text(text, tokenizer, max_length=8192, overlap=50):
18
   """
19
   Split the text into overlapping chunks based on token length.
20
   """
21
   tokens = tokenizer.tokenize(text)
22
   chunks = []
23
   for i in range(0, len(tokens), max_length - overlap):
24
       chunk_tokens = tokens[i:i + max_length]
25
       chunk = tokenizer.convert_tokens_to_string(chunk_tokens)
26
       chunks.append(chunk)
27
   return chunks
28
29
30
def get_embedding(text: str) -> list[float]:
31
   if not text.strip():
32
       print("Attempted to get embedding for empty text.")
33
       return []
34
35
36
   # Get the tokenizer from the model
37
   tokenizer = embedding_model.tokenizer
38
39
40
   # Split text into chunks if it's too long
41
   chunks = chunk_text(text, tokenizer, max_length=max_seq_length)
42
  
43
   if len(chunks) == 1:
44
       # If text fits in one chunk, embed as usual
45
       embedding = embedding_model.encode(text)
46
   else:
47
       # If text was split, embed each chunk and average the results
48
       chunk_embeddings = embedding_model.encode(chunks)
49
       embedding = np.mean(chunk_embeddings, axis=0)
50
51
   return embedding.tolist()
52
53
# Apply the embedding function with a progress bar
54
tqdm.pandas(desc="Generating embeddings")
55
dataset_df["embedding"] = dataset_df['combined_attributes'].progress_apply(get_embedding)

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

1
os.environ["MONGO_URI"] = getpass.getpass("Enter your MongoDB URI: ")

1
import pymongo
2
3
4
def get_mongo_client(mongo_uri):
5
 """Establish and validate connection to the MongoDB."""
6
7
8
 client = pymongo.MongoClient(mongo_uri, appname="devrel.showcase.rag.cohere_mongodb.python")
9
10
11
 # Validate the connection
12
 ping_result = client.admin.command('ping')
13
 if ping_result.get('ok') == 1.0:
14
   # Connection successful
15
   print("Connection to MongoDB successful")
16
   return client
17
 else:
18
   print("Connection to MongoDB failed")
19
 return None
20
21
22
MONGO_URI = os.environ["MONGO_URI"]
23
24
25
if not MONGO_URI:
26
 print("MONGO_URI not set in environment variables")
27
28
29
mongo_client = get_mongo_client(MONGO_URI)
30
31
32
DB_NAME = "asset_management_use_case"
33
COLLECTION_NAME = "market_reports"
34
35
36
db = mongo_client.get_database(DB_NAME)
37
collection = db.get_collection(COLLECTION_NAME)

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

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
6
 Args:
7
 user_query (str): The user's query string.
8
 collection (MongoCollection): The MongoDB collection to search.
9
10
11
 Returns:
12
 list: A list of matching documents.
13
 """
14
15
16
 # Generate embedding for the user query
17
 query_embedding = get_embedding(user_query)
18
19
20
 if query_embedding is None:
21
   return "Invalid query or embedding generation failed."
22
23
24
 # Define the vector search pipeline
25
 vector_search_stage = {
26
   "$vectorSearch": {
27
     "index": "vector_index",
28
     "queryVector": query_embedding,
29
     "path": "embedding",
30
     "numCandidates": 150,  # Number of candidate matches to consider
31
     "limit": 2  # Return top 4 matches
32
   }
33
 }
34
35
36
 unset_stage = {
37
   "$unset": "embedding"  # Exclude the 'embedding' field from the results
38
 }
39
40
41
 project_stage = {
42
   "$project": {
43
     "_id": 0,  # Exclude the _id field
44
     "company": 1,  # Include the plot field
45
     "reports": 1,  # Include the title field
46
     "combined_attributes": 1, # Include the genres field
47
     "score": {
48
       "$meta": "vectorSearchScore"  # Include the search score
49
     }
50
   }
51
 }
52
53
54
 pipeline = [vector_search_stage, unset_stage, project_stage]
55
56
57
 # Execute the search
58
 results = collection.aggregate(pipeline)
59
 return list(results)

1
def get_search_result(query, collection):
2
3
4
 get_knowledge = vector_search(query, collection)
5
6
7
 search_result = ''
8
 for result in get_knowledge:
9
     search_result += f"Company: {result.get('company', 'N/A')}, Combined Attributes: {result.get('combined_attributes', 'N/A')}\n"
10
11
12
 return search_result

1
# Conduct query with the retrieval of sources
2
query = "What companies have negative market reports or negative sentiment that might deter from investment in the long term"
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
6
print(combined_information)

1
import torch
2
from transformers import AutoTokenizer, AutoModelForCausalLM
3
4
5
tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b-it")
6
model = AutoModelForCausalLM.from_pretrained("google/gemma-2-2b-it", torch_dtype=torch.bfloat16)

1
def extract_model_response(response):
2
 # Split the response at the start of the model's turn
3
 parts = response.split("<start_of_turn>model")
4
  # If there's a model response, it will be in the last part
5
 if len(parts) > 1:
6
   model_response = parts[-1].strip()
7
  
8
   # Remove any potential end-of-turn markers
9
   model_response = model_response.split("<end_of_turn>")[0].strip()
10
  
11
   return model_response
12
 else:
13
   return "No model response found."

1
chat = [
2
   { "role": "user", "content": combined_information },
3
]
4
prompt = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
5
6
7
inputs = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
8
outputs = model.generate(input_ids=inputs.to(model.device), max_new_tokens=150, do_sample=True, temperature=0.7)
9
10
11
response = tokenizer.decode(outputs[0])

1
model_output = extract_model_response(response)
2
print(model_output)

1
faithfulness = (Number of truthful claims) / (Total number of claims in response)

1
from deepeval.test_case import LLMTestCase
2
from deepeval.metrics import FaithfulnessMetric

1
from deepeval.metrics import FaithfulnessMetric
2
3
4
actual_output = model_output
5
6
7
retrieval_context = [source_information]
8
9
10
metric = FaithfulnessMetric(
11
   threshold=0.7,
12
   model="gpt-4",
13
   include_reason=True
14
)
15
test_case = LLMTestCase(
16
   input=query,
17
   actual_output=actual_output,
18
   retrieval_context=retrieval_context
19
)
20
21
22
metric.measure(test_case)
23
print(metric.score)
24
print(metric.reason)

1
Event loop is already running. Applying nest_asyncio patch to allow async execution...
2
0.9333333333333333
3
The score is 0.93 because there is a minor discrepancy with the year of CDDY's strategic acquisitions. The actual output mistakenly indicates they occurred in 2023, while the context clearly states they happened in 2024.

1
from deepeval.metrics import AnswerRelevancyMetric
2
3
4
actual_output = model_output
5
6
7
metric = AnswerRelevancyMetric(
8
   threshold=0.7,
9
   model="gpt-4",
10
   include_reason=True
11
)
12
13
14
test_case = LLMTestCase(
15
   input=query,
16
   actual_output=actual_output
17
)
18
19
metric.measure(test_case)
20
print(metric.score)
21
print(metric.reason)

1
Event loop is already running. Applying nest_asyncio patch to allow async execution...
2
1.0
3
The score is 1.00 because the response accurately selected a company from the provided information and justified why it's a safe long-term investment, addressing all aspects of the input.