Update retrieval_engine.py

retrieval_engine.py

@@ -3,10 +3,10 @@ from sentence_transformers import SentenceTransformer
 import faiss
 import numpy as np
 
-class RetrievalEngine:
 
+class RetrievalEngine:
     def __init__(self):
-        self.dataset = load_dataset("YOUR_USERNAME/gmat-quant-corpus", split="train")
+        self.dataset = load_dataset("YOUR_ACTUAL_HF_USERNAME/YOUR_ACTUAL_DATASET_NAME", split="train")
 
         self.model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
 
@@ -21,13 +21,10 @@ class RetrievalEngine:
         self.texts = texts
 
     def search(self, query, k=3):
-
         q_emb = self.model.encode([query], convert_to_numpy=True)
-
         distances, indices = self.index.search(q_emb, k)
 
         results = []
-
         for idx in indices[0]:
             results.append(self.texts[idx])