Update src/embeddings.py

src/embeddings.py

@@ -1,12 +1,11 @@
 """
 embeddings.py — Local embedding generation and FAISS index management.
 
-all-MiniLM-L6-v2 runs entirely on CPU in ~200ms per batch.
-No API key, no cost, no rate limits.
+all-MiniLM-L6-v2 runs on CPU with no API key or cost.
+384-dimensional embeddings, ~80MB model, loads in ~3 seconds on first call.
 
-FAISS IndexFlatIP (inner product on normalized vectors = cosine similarity)
-is exact search — appropriate for < 100k chunks. For millions of chunks
-you'd switch to IndexIVFFlat with nprobe tuning.
+FAISS IndexFlatIP with L2-normalized vectors = exact cosine similarity search.
+Correct choice for < 500k chunks (enterprise PDF use case).
 """
 
 import logging
@@ -17,74 +16,62 @@ from src.utils import get_env
 
 logger = logging.getLogger("enterprise-rag.embeddings")
 
-# Global model instance — loaded once, reused across queries
-# Loading takes ~2-3 seconds on first call; cached in memory after that
 _embedding_model = None
-
-EMBEDDING_DIM = 384  # fixed for all-MiniLM-L6-v2
+EMBEDDING_DIM = 384
 
 
 def get_embedding_model() -> SentenceTransformer:
     """
-    Lazy-load the embedding model (singleton pattern).
-    On Hugging Face Spaces, the model is downloaded on first use
-    and cached in the Space's persistent storage.
+    Lazy-load the embedding model once and reuse across all calls.
+    Downloading happens on first use; cached in HF Spaces persistent storage.
     """
     global _embedding_model
     if _embedding_model is None:
-        model_name = get_env("EMBEDDING_MODEL", "sentence-transformers/all-MiniLM-L6-v2")
+        model_name = get_env(
+            "EMBEDDING_MODEL",
+            "sentence-transformers/all-MiniLM-L6-v2"
+        )
         logger.info(f"Loading embedding model: {model_name}")
         _embedding_model = SentenceTransformer(model_name)
-        logger.info("Embedding model loaded successfully")
+        logger.info("Embedding model ready")
     return _embedding_model
 
 
 def embed_texts(texts: list) -> np.ndarray:
     """
-    Generate normalized embeddings for a list of texts.
+    Generate L2-normalized embeddings for a list of texts.
 
-    Normalization (L2) converts dot product to cosine similarity,
-    which is what IndexFlatIP computes. This is the standard approach
-    for semantic search with sentence-transformers.
+    Normalization converts dot product to cosine similarity,
+    which is what IndexFlatIP computes. Standard for semantic search.
 
-    Returns: float32 array of shape (len(texts), 384)
+    Returns float32 array of shape (len(texts), 384).
     """
     if not texts:
-        return np.array([], dtype=np.float32)
+        return np.zeros((0, EMBEDDING_DIM), dtype=np.float32)
 
     model = get_embedding_model()
-
-    # batch_size=32 balances memory and throughput on CPU
-    # show_progress_bar=False keeps logs clean in production
     embeddings = model.encode(
         texts,
         batch_size=32,
-        normalize_embeddings=True,   # L2 normalize for cosine similarity
+        normalize_embeddings=True,
         show_progress_bar=False,
         convert_to_numpy=True,
     )
-
     return embeddings.astype(np.float32)
 
 
 def build_faiss_index(embeddings: np.ndarray) -> faiss.IndexFlatIP:
     """
-    Build a FAISS flat inner-product index from embeddings.
-
-    IndexFlatIP with L2-normalized vectors = exact cosine similarity search.
-    No training required, no approximation, deterministic results.
-
-    For > 500k vectors, consider IndexIVFFlat:
-        quantizer = faiss.IndexFlatIP(EMBEDDING_DIM)
-        index = faiss.IndexIVFFlat(quantizer, EMBEDDING_DIM, nlist=1024)
-        index.train(embeddings)
+    Build a FAISS flat inner-product index.
+    With L2-normalized vectors this equals exact cosine similarity search.
+    No training required, deterministic results.
     """
     if embeddings.shape[0] == 0:
-        raise ValueError("Cannot build FAISS index from empty embeddings array")
+        raise ValueError("Cannot build FAISS index from empty embeddings.")
 
     index = faiss.IndexFlatIP(EMBEDDING_DIM)
     index.add(embeddings)
-    logger.info(f"FAISS index built: {index.ntotal} vectors, dim={EMBEDDING_DIM}")
+    logger.info(f"FAISS index built: {index.ntotal} vectors")
     return index
 
 
@@ -97,20 +84,19 @@ def search_index(
     Retrieve top-k most similar chunks.
 
     Returns:
-        scores: cosine similarity scores (float32 array, shape [top_k])
-        indices: chunk indices in original list (int64 array, shape [top_k])
+        scores  — cosine similarity scores, float32 array shape [top_k]
+        indices — chunk positions in original list, int64 array shape [top_k]
 
-    Score interpretation:
-        1.0 = identical (query matches chunk exactly)
+    Score guide:
+        1.0  = identical
         0.8+ = highly relevant
         0.5-0.8 = moderately relevant
-        < 0.5 = likely irrelevant — consider raising this threshold
+        < 0.5 = likely irrelevant
     """
     if index.ntotal == 0:
-        return np.array([]), np.array([])
+        return np.array([], dtype=np.float32), np.array([], dtype=np.int64)
 
     k = min(top_k, index.ntotal)
     query_2d = query_embedding.reshape(1, -1).astype(np.float32)
     scores, indices = index.search(query_2d, k)
-
     return scores[0], indices[0]