Update src/qa.py

src/qa.py

@@ -85,71 +85,66 @@ REASONING_PROMPT = (
 # ==========================================================
 # 5️⃣ Retrieve Chunks (FAISS + Rerank + Neighbor Expansion)
 # ==========================================================
-def retrieve_chunks(
-    query: str,
-    index,
-    chunks: list,
-    top_k: int = 3,
-    topn_candidates: int = 20,
-    neighbor_threshold: float = 0.68,
-    expansion_window: int = 1,
-    max_context_chunks: int = 6,
-):
-    """Retrieve semantically relevant chunks with reranking and neighbor expansion."""
+def retrieve_chunks(query: str, index, chunks: list, top_k: int = 5, min_similarity: float = 0.6):
+    """
+    Hybrid retrieval:
+    1️⃣ Get semantic top-K chunks via FAISS.
+    2️⃣ Re-rank by cosine similarity and apply a minimum similarity filter.
+    3️⃣ If fewer than top_k remain, fill remaining seats with adjacent chunks (±1) for continuity.
+    """
     if not index or not chunks:
         return []
 
-    # 1️⃣ Encode query (normalized)
-    query_emb = _query_model.encode(
-        [f"query: {query.strip()}"],
-        convert_to_numpy=True,
-        normalize_embeddings=True
-    )[0].astype("float32")
-
-    # 2️⃣ FAISS search (initial candidates)
-    topn_candidates = min(topn_candidates, getattr(index, "ntotal", topn_candidates))
-    _, candidate_ids = index.search(np.array([query_emb]).astype("float32"), topn_candidates)
-    candidate_ids = [int(i) for i in candidate_ids[0] if i != -1]
-
-    # 3️⃣ Re-encode candidate chunks and compute cosine similarities
-    candidate_texts = [chunks[i] for i in candidate_ids]
-    candidate_vecs = np.array([
-        _query_model.encode([t], convert_to_numpy=True, normalize_embeddings=True)[0]
-        for t in candidate_texts
-    ])
-    sims = cosine_similarity([query_emb], candidate_vecs)[0]
-    sorted_idx = np.argsort(sims)[::-1]
-    reranked_ids = [candidate_ids[i] for i in sorted_idx]
-
-    # 4️⃣ Select top-k base chunks
-    selected, selected_set = [], set()
-    for rid in reranked_ids:
-        if len(selected) >= top_k:
-            break
-        selected.append(rid)
-        selected_set.add(rid)
-
-    # 5️⃣ Conditional neighbor expansion
-    final_order = list(selected)
-    for base_id in selected:
-        if len(final_order) >= max_context_chunks:
-            break
-        for offset in range(1, expansion_window + 1):
-            for neighbor in (base_id - offset, base_id + offset):
-                if neighbor < 0 or neighbor >= len(chunks) or neighbor in selected_set:
-                    continue
-                # Check semantic closeness
-                neighbor_vec = _query_model.encode([chunks[neighbor]], convert_to_numpy=True, normalize_embeddings=True)[0]
-                sim = float(cosine_similarity([query_emb], [neighbor_vec])[0][0])
-                if sim >= neighbor_threshold:
-                    final_order.append(neighbor)
-                    selected_set.add(neighbor)
-                    if len(final_order) >= max_context_chunks:
-                        break
-            if len(final_order) >= max_context_chunks:
-                break
-
-    return [chunks[i] for i in final_order]
+    try:
+        # Encode query
+        q_emb = _query_model.encode(
+            [f"query: {query.strip()}"],
+            convert_to_numpy=True,
+            normalize_embeddings=True
+        )[0]
+
+        # Step 1️⃣ — FAISS initial retrieval
+        distances, indices = index.search(np.array([q_emb]).astype("float32"), top_k * 2)
+        retrieved_indices = list(indices[0])
+
+        # Step 2️⃣ — Compute cosine similarity for re-ranking
+        retrieved_texts = [chunks[i] for i in retrieved_indices]
+        doc_embs = _query_model.encode(
+            [f"passage: {c}" for c in retrieved_texts],
+            convert_to_numpy=True,
+            normalize_embeddings=True
+        )
+
+        sims = cosine_similarity([q_emb], doc_embs)[0]
+        ranked = sorted(zip(retrieved_indices, sims), key=lambda x: x[1], reverse=True)
+
+        # Step 3️⃣ — Apply minimum similarity filter
+        filtered_indices = [idx for idx, score in ranked if score >= min_similarity]
+
+        # Step 4️⃣ — If not enough, add ±1 neighbors for continuity
+        if len(filtered_indices) < top_k:
+            extras_needed = top_k - len(filtered_indices)
+            expanded_indices = set(filtered_indices)
+            for idx in filtered_indices:
+                for neighbor in [idx - 1, idx + 1]:
+                    if 0 <= neighbor < len(chunks):
+                        expanded_indices.add(neighbor)
+                        if len(expanded_indices) >= top_k:
+                            break
+                if len(expanded_indices) >= top_k:
+                    break
+            filtered_indices = list(sorted(expanded_indices))[:top_k]
+
+        # Step 5️⃣ — Build final ordered list of chunks
+        final_chunks = [chunks[i] for i in filtered_indices]
+
+        print(f"✅ Retrieved {len(final_chunks)} chunks (semantic + neighbor fill).")
+        return final_chunks
+
+    except Exception as e:
+        print(f"⚠️ Retrieval error: {e}")
+        return []
+
 
 # ==========================================================
 # 6️⃣ Answer Generation