Update src/retrieval.py

src/retrieval.py

@@ -5,7 +5,6 @@ from sentence_transformers import SentenceTransformer
 from typing import List, Dict
 import logging
 
-# Configure logging
 logger = logging.getLogger(__name__)
 
 class DocumentRetriever:
@@ -26,7 +25,7 @@ class DocumentRetriever:
             raise
         self.data_path = data_path
         self.documents = self._load_documents()
-        self.doc_embeddings = self._embed_documents()
+        self.doc_embeddings = self._load_or_compute_embeddings()
 
     def _load_documents(self) -> List[Dict]:
         """Load documents from the JSON file."""
@@ -42,15 +41,40 @@ class DocumentRetriever:
             logger.warning(f"Invalid JSON in {self.data_path}, using empty documents")
             return []
 
-    def _embed_documents(self) -> np.ndarray:
-        """Embed document contents using the SentenceTransformer model."""
+    def _load_or_compute_embeddings(self) -> np.ndarray:
+        """Load cached embeddings or compute new ones."""
+        embedding_cache_path = 'data/doc_embeddings.npy'
         if not self.documents:
             logger.info("No documents to embed, returning empty embeddings")
             return np.array([])
+        
+        # Check for cached embeddings
+        if os.path.exists(embedding_cache_path):
+            try:
+                embeddings = np.load(embedding_cache_path)
+                if embeddings.shape[0] == len(self.documents):
+                    logger.info(f"Loaded {embeddings.shape[0]} cached embeddings from {embedding_cache_path}")
+                    return embeddings
+                else:
+                    logger.warning(f"Cached embeddings shape mismatch, recomputing...")
+            except Exception as e:
+                logger.warning(f"Failed to load cached embeddings: {str(e)}, recomputing...")
+        
+        # Compute new embeddings
         texts = [doc['content'] for doc in self.documents]
-        logger.info(f"Embedding {len(texts)} documents...")
-        embeddings = self.model.encode(texts)
-        logger.info("Document embeddings generated successfully")
+        logger.info(f"Computing embeddings for {len(texts)} documents...")
+        start_time = time.time()
+        embeddings = self.model.encode(texts, batch_size=32, show_progress_bar=True)
+        logger.info(f"Embedding {len(texts)} documents took {time.time() - start_time:.2f} seconds")
+        
+        # Cache embeddings
+        try:
+            os.makedirs('data', exist_ok=True)
+            np.save(embedding_cache_path, embeddings)
+            logger.info(f"Saved embeddings to {embedding_cache_path}")
+        except Exception as e:
+            logger.warning(f"Failed to save embeddings: {str(e)}")
+        
         return embeddings
 
     def retrieve(self, query: str, top_k: int = 3) -> List[Dict]: