load dataset locally

data/load_dataset.py

@@ -1,9 +1,23 @@
+import os
 import logging
 from datasets import load_dataset
+import pickle  # For saving the dataset locally
 
-def load_data(data_set_name):
-    logging.info("Loading dataset")
-    dataset = load_dataset("rungalileo/ragbench", data_set_name, split="test")
+def load_data(data_set_name, local_path="local_datasets"):
+    os.makedirs(local_path, exist_ok=True)
+    dataset_file = os.path.join(local_path, f"{data_set_name}_test.pkl")
+    
+    if os.path.exists(dataset_file):
+        logging.info("Loading dataset from local storage")
+        with open(dataset_file, "rb") as f:
+            dataset = pickle.load(f)
+    else:
+        logging.info("Loading dataset from Hugging Face")
+        dataset = load_dataset("rungalileo/ragbench", data_set_name, split="test")
+        logging.info(f"Saving {data_set_name} dataset locally")
+        with open(dataset_file, "wb") as f:
+            pickle.dump(dataset, f)
+    
     logging.info("Dataset loaded successfully")
     logging.info(f"Number of documents found: {dataset.num_rows}")
-    return dataset
+    return dataset

generator/initialize_llm.py

@@ -3,14 +3,14 @@ import os
 from langchain_groq import ChatGroq
 
 def initialize_generation_llm():
-    os.environ["GROQ_API_KEY"] = "your_groq_api_key"
+    os.environ["GROQ_API_KEY"] = "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"
     model_name = "llama3-8b-8192"
     llm = ChatGroq(model=model_name, temperature=0.7)
     logging.info(f'Generation LLM {model_name} initialized')
     return llm
 
 def initialize_validation_llm():
-    os.environ["GROQ_API_KEY"] = "your_groq_api_key"
+    os.environ["GROQ_API_KEY"] = "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"
     model_name = "llama3-70b-8192"
     llm = ChatGroq(model=model_name, temperature=0.7)
     logging.info(f'Validation LLM {model_name} initialized')

main.py

@@ -12,7 +12,7 @@ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(
 
 def main():
     logging.info("Starting the RAG pipeline")
-    data_set_name = 'techqa'
+    data_set_name = 'covidqa'
 
     # Load the dataset
     dataset = load_data(data_set_name)
@@ -36,11 +36,11 @@ def main():
     val_llm = initialize_validation_llm()
 
     # Sample question
-    row_num = 7
+    row_num = 2
     query = dataset[row_num]['question']
 
     # Call generate_metrics for above sample question
-    generate_metrics(gen_llm, val_llm, vector_store, query)
+    #generate_metrics(gen_llm, val_llm, vector_store, query)
     
     #Compute RMSE and AUC-ROC for entire dataset
     compute_rmse_auc_roc_metrics(gen_llm, val_llm, dataset, vector_store, 10)

retriever/embed_documents.py

@@ -1,7 +1,17 @@
+import os
+import logging
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
 
-def embed_documents(documents):
+def embed_documents(documents, embedding_path="embeddings.faiss"):
     embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/paraphrase-MiniLM-L3-v2")
-    vector_store = FAISS.from_texts([doc['text'] for doc in documents], embedding_model)
+
+    if os.path.exists(embedding_path):
+        logging.info("Loading embeddings from local file")
+        vector_store = FAISS.load_local(embedding_path, embedding_model, allow_dangerous_deserialization=True)
+    else:
+        logging.info("Generating and saving embeddings")
+        vector_store = FAISS.from_texts([doc['text'] for doc in documents], embedding_model)
+        vector_store.save_local(embedding_path)
+    
     return vector_store