Update src/qa.py

src/qa.py

@@ -2,7 +2,7 @@
 qa.py — Retrieval + Generation Layer
 -------------------------------------
 Handles:
-• Query embedding (MPNet / E5 / MiniLM fallback)
+• Query embedding (SentenceTransformer / E5-compatible)
 • Chunk retrieval (FAISS)
 • Answer generation (Flan-T5)
 Optimized for Hugging Face Spaces & Streamlit.
@@ -29,36 +29,27 @@ os.environ.update({
 })
 
 # ==========================================================
-# 2️⃣ Query Embedding Model (MPNet → E5 → MiniLM)
+# 2️⃣ Query Embedding Model
 # ==========================================================
-# Try best retrieval model first, then gracefully degrade
+# Use E5-small-v2 for retrieval consistency with embeddings.py
 try:
     _query_model = SentenceTransformer(
-        "sentence-transformers/all-mpnet-base-v2",  # ✅ Best for QA and reasoning-heavy text
+        "intfloat/e5-small-v2",
         cache_folder=CACHE_DIR
     )
-    print("✅ Loaded query model: all-mpnet-base-v2")
-
-except Exception as e1:
-    print(f"⚠️ MPNet load failed ({e1}), trying E5-small-v2...")
-    try:
-        _query_model = SentenceTransformer(
-            "intfloat/e5-small-v2",
-            cache_folder=CACHE_DIR
-        )
-        print("✅ Loaded fallback model: e5-small-v2")
-    except Exception as e2:
-        print(f"⚠️ E5 load failed ({e2}), falling back to MiniLM...")
-        _query_model = SentenceTransformer(
-            "sentence-transformers/all-MiniLM-L6-v2",
-            cache_folder=CACHE_DIR
-        )
-        print("✅ Loaded fallback model: MiniLM-L6-v2")
+    print("✅ Loaded query model: intfloat/e5-small-v2")
+except Exception as e:
+    print(f"⚠️ Query model load failed ({e}), falling back to MiniLM.")
+    _query_model = SentenceTransformer(
+        "sentence-transformers/all-MiniLM-L6-v2",
+        cache_folder=CACHE_DIR
+    )
+    print("✅ Loaded fallback model: all-MiniLM-L6-v2")
 
 # ==========================================================
 # 3️⃣ LLM for Answer Generation (FLAN-T5)
 # ==========================================================
-MODEL_NAME = "google/flan-t5-base"   # use 'large' if you have enough memory
+MODEL_NAME = "google/flan-t5-base"   # switch to 'large' if RAM allows
 print(f"✅ Loading LLM: {MODEL_NAME}")
 
 _tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, cache_dir=CACHE_DIR)
@@ -68,7 +59,7 @@ _answer_model = pipeline(
     "text2text-generation",
     model=_model,
     tokenizer=_tokenizer,
-    device=-1  # CPU-safe for Hugging Face Spaces
+    device=-1  # CPU-safe for Spaces
 )
 
 # ==========================================================
@@ -96,16 +87,15 @@ Answer:
 def retrieve_chunks(query: str, index, chunks: list, top_k: int = 3):
     """
     Encodes the user query and retrieves top-k relevant chunks via FAISS.
-    Uses 'query:' prefix (E5 / instruction-tuned) for semantic alignment.
+    Uses 'query:' prefix (E5 training style) for semantic alignment.
     """
     if not index or not chunks:
         return []
 
     try:
-        # Prefix improves intent understanding (esp. for E5 / MPNet)
-        prefix = "query: " if "e5" in _query_model.name_or_path.lower() else ""
+        # E5 expects 'query:' prefix for better retrieval accuracy
         query_emb = _query_model.encode(
-            [f"{prefix}{query.strip()}"],
+            [f"query: {query.strip()}"],
             convert_to_numpy=True,
             normalize_embeddings=True
         )[0]
@@ -128,8 +118,10 @@ def generate_answer(query: str, retrieved_chunks: list):
     if not retrieved_chunks:
         return "Sorry, I couldn’t find relevant information in the document."
 
-    # Combine chunks as structured context
+    # Merge retrieved chunks for context
     context = "\n\n".join([f"[Chunk {i+1}]: {chunk}" for i, chunk in enumerate(retrieved_chunks)])
+
+    # Build structured prompt
     prompt = PROMPT_TEMPLATE.format(context=context, query=query)
 
     try:
@@ -146,7 +138,7 @@ def generate_answer(query: str, retrieved_chunks: list):
 
 
 # ==========================================================
-# 7️⃣ Optional: Local Test Run
+# 7️⃣ Optional Local Test
 # ==========================================================
 if __name__ == "__main__":
     dummy_chunks = [