Update src/embeddings.py

src/embeddings.py

@@ -1,10 +1,13 @@
-# ----------------------------
-# Hugging Face cache bootstrap
-# ----------------------------
+# ==========================================================
+# 📘 embeddings.py — optimized for Hugging Face + FAISS + E5
+# ==========================================================
 import os
 import numpy as np
 from sentence_transformers import SentenceTransformer
 
+# ----------------------------
+# Hugging Face Cache Bootstrap
+# ----------------------------
 CACHE_DIR = "/tmp/hf_cache"
 os.makedirs(CACHE_DIR, exist_ok=True)
 
@@ -16,11 +19,11 @@ os.environ["HF_MODULES_CACHE"] = CACHE_DIR
 print(f"✅ Using Hugging Face cache at {CACHE_DIR}")
 
 # ----------------------------
-# Load embedding model once (with fallback)
+# Load Embedding Model (E5 with fallback)
 # ----------------------------
 try:
     _model = SentenceTransformer(
-        "intfloat/e5-small-v2",     # ✅ Better for document QA retrieval
+        "intfloat/e5-small-v2",   # ✅ Trained for retrieval-augmented QA
         cache_folder=CACHE_DIR
     )
     print("✅ Loaded model: intfloat/e5-small-v2")
@@ -33,20 +36,39 @@ except Exception as e:
     print("✅ Loaded fallback model: all-MiniLM-L6-v2")
 
 # ----------------------------
-# Function: generate embeddings
+# Function: Generate Embeddings
 # ----------------------------
 def generate_embeddings(chunks: list) -> list:
     """
-    Generate normalized embeddings for a list of text chunks.
-    Normalization improves FAISS retrieval accuracy (cosine-based).
+    📌 Generate normalized embeddings for a list of text chunks.
+
+    Args:
+        chunks (list): List of text chunks.
+
+    Returns:
+        list: List of normalized embedding vectors (Python lists).
+
+    Notes:
+        - Prefixing chunks with 'passage:' improves retrieval accuracy for E5.
+        - normalize_embeddings=True ensures cosine-similarity consistency.
+        - Works efficiently even for large PDFs.
     """
-    # Add semantic prefix for e5 model to help it distinguish queries vs passages
+    if not chunks:
+        print("⚠️ No chunks provided for embedding generation.")
+        return []
+
+    # Step 1: Prefix each chunk for semantic clarity (per E5 training)
     prepared_chunks = [f"passage: {chunk.strip()}" for chunk in chunks]
 
+    # Step 2: Encode with normalization for cosine similarity
     vectors = _model.encode(
         prepared_chunks,
         convert_to_numpy=True,
-        normalize_embeddings=True  # ✅ ensures better FAISS matching
+        normalize_embeddings=True  # ✅ Makes FAISS IndexFlatIP accurate
     )
 
-    return vectors.tolist()
+    # Step 3: Convert to Python list for FAISS / JSON compatibility
+    embeddings = vectors.tolist()
+
+    print(f"✅ Generated {len(embeddings)} embeddings.")
+    return embeddings