Upload 5 files

retrieval.py

@@ -3,27 +3,29 @@ import numpy as np
 from langchain.schema import Document
 import faiss
 
-from data_processing import embedding_model #, index, actual_docs
+from data_processing import embedding_model , index, actual_docs
 
 retrieved_docs = None
 
 # Retrieval Function
 def retrieve_documents(query, top_k=5):
-    faiss_index_path = f"data_local/rag7_index.faiss"
+    query_dataset = find_query_dataset(query)
+    #index, chunk_docs = load_data_from_faiss(query)
+
+    faiss_index_path = f"data_local/{query_dataset}__quantized.faiss"
     index = faiss.read_index(faiss_index_path)
 
     query_embedding = np.array(embedding_model.embed_documents([query]), dtype=np.float32)
 
     _, nearest_indices = index.search(query_embedding, top_k)
 
-    with open(f"data_local/rag7_docs.json", "r") as f:
+    with open( f"data_local/{query_dataset}_chunked_docs.json", "r") as f:
         documents = json.load(f)  # Contains all documents for this dataset
 
     retrieved_docs = [Document(page_content=documents[i]) for i in nearest_indices[0]]
 
     return retrieved_docs
 
-
 def remove_duplicate_documents(documents):
     unique_documents = []
     seen_documents = set()  # To keep track of seen documents
@@ -36,12 +38,13 @@ def remove_duplicate_documents(documents):
     return unique_documents
 
 def find_query_dataset(query):
-    index = faiss.read_index("question_index.faiss")
+    index = faiss.read_index("data_local/question_quantized.faiss")
 
-    with open("dataset_mapping.json", "r") as f:
+    with open("data_local/dataset_mapping.json", "r") as f:
         dataset_names = json.load(f)
 
     question_embedding = np.array(embedding_model.embed_documents([query]), dtype=np.float32)
     _, nearest_index = index.search(question_embedding, 1)  
     best_dataset = dataset_names[nearest_index[0][0]]
-    return best_dataset
+    return best_dataset
+