Update src/embeddings.py

src/embeddings.py

@@ -2,6 +2,8 @@
 # Hugging Face cache bootstrap
 # ----------------------------
 import os
+import numpy as np
+from sentence_transformers import SentenceTransformer
 
 CACHE_DIR = "/tmp/hf_cache"
 os.makedirs(CACHE_DIR, exist_ok=True)
@@ -14,28 +16,37 @@ os.environ["HF_MODULES_CACHE"] = CACHE_DIR
 print(f"✅ Using Hugging Face cache at {CACHE_DIR}")
 
 # ----------------------------
-# Imports AFTER cache bootstrap
-# ----------------------------
-from sentence_transformers import SentenceTransformer
-
-# ----------------------------
-# Load embedding model once
-# ----------------------------
-_model = SentenceTransformer(
-    "sentence-transformers/all-MiniLM-L6-v2",
-    cache_folder=CACHE_DIR
-)
+# Load embedding model once (with fallback)
+# ----------------------------
+try:
+    _model = SentenceTransformer(
+        "intfloat/e5-small-v2",     # ✅ Better for document QA retrieval
+        cache_folder=CACHE_DIR
+    )
+    print("✅ Loaded model: intfloat/e5-small-v2")
+except Exception as e:
+    print(f"⚠️ Model load failed ({e}), falling back to MiniLM.")
+    _model = SentenceTransformer(
+        "sentence-transformers/all-MiniLM-L6-v2",
+        cache_folder=CACHE_DIR
+    )
+    print("✅ Loaded fallback model: all-MiniLM-L6-v2")
 
 # ----------------------------
 # Function: generate embeddings
 # ----------------------------
 def generate_embeddings(chunks: list) -> list:
     """
-    📌 Generate embeddings for a list of text chunks.
-    Args:
-        chunks (list): List of text chunks.
-    Returns:
-        list: List of embedding vectors (plain Python lists).
+    Generate normalized embeddings for a list of text chunks.
+    Normalization improves FAISS retrieval accuracy (cosine-based).
     """
-    embeddings = _model.encode(chunks, convert_to_numpy=True)   # numpy array
-    return embeddings.tolist()   # convert to lists for FAISS / JSON
+    # Add semantic prefix for e5 model to help it distinguish queries vs passages
+    prepared_chunks = [f"passage: {chunk.strip()}" for chunk in chunks]
+
+    vectors = _model.encode(
+        prepared_chunks,
+        convert_to_numpy=True,
+        normalize_embeddings=True  # ✅ ensures better FAISS matching
+    )
+
+    return vectors.tolist()