Update src/qa.py

src/qa.py

@@ -90,18 +90,28 @@ REASONING_PROMPT = (
 # 5️⃣ Retrieve Top-K Chunks (Balanced speed)
 # ==========================================================
 def retrieve_chunks(query: str, index, chunks: list, top_k: int = 3):
-    """Retrieve top-K relevant chunks efficiently."""
+    """Retrieve top-K relevant chunks and re-rank by cosine similarity for better precision."""
     if not index or not chunks:
         return []
 
     try:
         q_emb = _query_model.encode([f"query: {query.strip()}"], convert_to_numpy=True, normalize_embeddings=True)[0]
-        distances, indices = index.search(np.array([q_emb]).astype("float32"), top_k)
-        return [chunks[i] for i in indices[0]]
+        distances, indices = index.search(np.array([q_emb]).astype("float32"), top_k * 2)
+
+        # Compute similarity scores for re-ranking
+        candidates = [chunks[i] for i in indices[0]]
+        cand_vecs = _query_model.encode(candidates, convert_to_numpy=True, normalize_embeddings=True)
+        sims = cosine_similarity([q_emb], cand_vecs)[0]
+
+        # Return top-K most semantically aligned
+        top_indices = np.argsort(sims)[::-1][:top_k]
+        return [candidates[i] for i in top_indices]
+
     except Exception as e:
         print(f"⚠️ Retrieval error: {e}")
         return []
 
+
 # ==========================================================
 # 6️⃣ Generate Answer (Reasoning or Strict Mode)
 # ==========================================================