Update app/policy_vector_db.py

app/policy_vector_db.py

@@ -8,6 +8,9 @@ from sentence_transformers import SentenceTransformer
 import chromadb
 from chromadb.config import Settings
 import logging
+import multiprocessing as mp
+from concurrent.futures import ThreadPoolExecutor
+import numpy as np
 
 # --- Basic Logging Setup ---
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
@@ -16,20 +19,40 @@ logger = logging.getLogger(__name__)
 class PolicyVectorDB:
     """
     Enhanced vector database for policy documents with metadata-aware search capabilities.
+    Optimized for CPU utilization.
     """
     def __init__(self, persist_directory: str, top_k_default: int = 5, relevance_threshold: float = 0.5):
         self.persist_directory = persist_directory
         self.client = chromadb.PersistentClient(path=persist_directory, settings=Settings(allow_reset=True))
         self.collection_name = "neepco_dop_policies"
         
+        # Optimize CPU usage
+        self.cpu_count = mp.cpu_count()
+        torch.set_num_threads(self.cpu_count)
+        
+        logger.info(f"Detected {self.cpu_count} CPU cores, optimizing threading...")
         logger.info("Loading embedding model 'BAAI/bge-large-en-v1.5'. This may take a moment...")
-        self.embedding_model = SentenceTransformer('BAAI/bge-large-en-v1.5', device='cpu') 
+        
+        # Optimize model loading for CPU
+        self.embedding_model = SentenceTransformer(
+            'BAAI/bge-large-en-v1.5', 
+            device='cpu',
+            # Use all available CPU cores for inference
+            model_kwargs={'torch_dtype': torch.float32}
+        )
+        
+        # Set model to use optimized CPU inference
+        self.embedding_model.max_seq_length = 512  # Reduce context length for speed
+        
         logger.info("Embedding model loaded successfully.")
         
         self.collection = None
         self.top_k_default = top_k_default
         self.relevance_threshold = relevance_threshold
         
+        # Thread pool for parallel processing
+        self.thread_pool = ThreadPoolExecutor(max_workers=self.cpu_count)
+        
         # Add monetary normalization for queries
         self.money_patterns = {
             r'(\d+(?:,\d+)*(?:\.\d+)?)\s*crore': lambda x: float(x.replace(',', '')) * 1e7,
@@ -97,8 +120,40 @@ class PolicyVectorDB:
         
         return entities
 
+    def _encode_batch_parallel(self, texts: List[str]) -> np.ndarray:
+        """Parallel encoding of text batches for better CPU utilization."""
+        # Split texts into smaller batches for parallel processing
+        batch_size = max(1, len(texts) // self.cpu_count)
+        if len(texts) <= batch_size:
+            return self.embedding_model.encode(
+                texts, 
+                normalize_embeddings=True, 
+                show_progress_bar=False,
+                batch_size=32,  # Optimize batch size for CPU
+                convert_to_numpy=True
+            )
+        
+        # Process in parallel batches
+        batches = [texts[i:i + batch_size] for i in range(0, len(texts), batch_size)]
+        
+        def encode_batch(batch):
+            return self.embedding_model.encode(
+                batch, 
+                normalize_embeddings=True, 
+                show_progress_bar=False,
+                batch_size=16,
+                convert_to_numpy=True
+            )
+        
+        # Use thread pool for parallel encoding
+        futures = [self.thread_pool.submit(encode_batch, batch) for batch in batches]
+        results = [future.result() for future in futures]
+        
+        # Concatenate results
+        return np.vstack(results) if results else np.array([])
+
     def add_chunks(self, chunks: List[Dict]):
-        """Enhanced chunk addition with better metadata handling."""
+        """Enhanced chunk addition with better metadata handling and parallel processing."""
         collection = self._get_collection()
         if not chunks:
             logger.info("No chunks provided to add.")
@@ -119,7 +174,9 @@ class PolicyVectorDB:
         
         logger.info(f"Adding {len(new_chunks)} new chunks to the vector database...")
         
-        batch_size = 32
+        # Optimized batch size for CPU processing
+        batch_size = min(64, max(16, len(new_chunks) // 4))
+        
         for i in range(0, len(new_chunks), batch_size):
             batch = new_chunks[i:i + batch_size]
             
@@ -127,7 +184,8 @@ class PolicyVectorDB:
             texts = [chunk['text'] for chunk in batch]
             metadatas = [self._flatten_metadata(chunk.get('metadata', {})) for chunk in batch]
             
-            embeddings = self.embedding_model.encode(texts, normalize_embeddings=True, show_progress_bar=False).tolist()
+            # Use parallel encoding
+            embeddings = self._encode_batch_parallel(texts).tolist()
             
             collection.add(ids=ids, embeddings=embeddings, documents=texts, metadatas=metadatas)
             logger.info(f"Added batch {i//batch_size + 1}/{(len(new_chunks) + batch_size - 1) // batch_size}")
@@ -137,6 +195,7 @@ class PolicyVectorDB:
     def search(self, query_text: str, top_k: int = None, filters: Dict = None) -> List[Dict]:
         """
         Enhanced search with metadata filtering and entity extraction.
+        Optimized for CPU performance.
         """
         collection = self._get_collection()
         
@@ -158,7 +217,15 @@ class PolicyVectorDB:
             where_conditions["section"] = {"$in": [s.split()[-1] for s in entities['sections']]}
         
         instructed_query = f"Represent this sentence for searching relevant passages: {query_text}"
-        query_embedding = self.embedding_model.encode([instructed_query], normalize_embeddings=True).tolist()
+        
+        # Optimized single query encoding
+        query_embedding = self.embedding_model.encode(
+            [instructed_query], 
+            normalize_embeddings=True,
+            show_progress_bar=False,
+            batch_size=1,
+            convert_to_numpy=True
+        ).tolist()
         
         k = top_k if top_k is not None else self.top_k_default
         
@@ -176,14 +243,13 @@ class PolicyVectorDB:
         
         search_results = []
         if results and results.get('documents') and results['documents'][0]:
-            for i, doc in enumerate(results['documents'][0]):  # Fixed: iterate over results['documents'][0]
-                # Fixed: Access distances correctly as results['distances'][0][i]
+            for i, doc in enumerate(results['documents'][0]):
                 relevance_score = 1 - results['distances'][0][i]
                 
                 if relevance_score >= self.relevance_threshold:
                     result = {
                         'text': doc,
-                        'metadata': results['metadatas'][0][i],  # Fixed: Access metadata correctly
+                        'metadata': results['metadatas'][0][i],
                         'relevance_score': relevance_score
                     }
                     
@@ -273,6 +339,11 @@ class PolicyVectorDB:
             logger.warning(f"Error in search_by_amount: {e}")
             return []
 
+    def __del__(self):
+        """Cleanup thread pool on deletion."""
+        if hasattr(self, 'thread_pool'):
+            self.thread_pool.shutdown(wait=False)
+
 def ensure_db_populated(db_instance: PolicyVectorDB, chunks_file_path: str) -> bool:
     """Checks if the DB is empty and populates it from a JSONL file if needed."""
     try:
@@ -297,7 +368,13 @@ def ensure_db_populated(db_instance: PolicyVectorDB, chunks_file_path: str) -> b
             logger.warning(f"Chunks file at '{chunks_file_path}' is empty or invalid. No data to add.")
             return False
 
-        db_instance.add_chunks(chunks_to_add)
+        # Process in batches to avoid memory issues
+        batch_size = 500
+        for i in range(0, len(chunks_to_add), batch_size):
+            batch = chunks_to_add[i:i + batch_size]
+            db_instance.add_chunks(batch)
+            logger.info(f"Processed batch {i//batch_size + 1}/{(len(chunks_to_add) + batch_size - 1) // batch_size}")
+
         logger.info("Vector database population attempt complete.")
         return True
     except Exception as e: