Update retrieval.py

retrieval.py

@@ -1,8 +1,13 @@
-# retrieval.py — FAISS + optional BM25 (no reranker; CPU-friendly)
-import json, faiss, numpy as np, os, re
+import json, os, re
 from typing import List, Dict, Any
-from sentence_transformers import SentenceTransformer
+import numpy as np
 
+try:
+    import faiss
+except Exception:
+    faiss = None
+
+from sentence_transformers import SentenceTransformer
 try:
     from rank_bm25 import BM25Okapi
 except Exception:
@@ -20,11 +25,17 @@ class Retriever:
             for line in f:
                 self.chunks.append(json.loads(line))
 
-        # faiss
-        self.faiss = faiss.read_index(FAISS_PATH)
-        self.embed = SentenceTransformer(embed_model_name)
+        # try FAISS
+        self.faiss = None
+        self.embed = None
+        if faiss is not None and os.path.exists(FAISS_PATH):
+            try:
+                self.faiss = faiss.read_index(FAISS_PATH)
+                self.embed = SentenceTransformer(embed_model_name)
+            except Exception:
+                self.faiss = None
 
-        # bm25 (optional)
+        # BM25 is immediate
         if BM25Okapi is not None:
             tokenized = [self._tokenize(c["chunk"]) for c in self.chunks]
             self.bm25 = BM25Okapi(tokenized)
@@ -37,26 +48,33 @@ class Retriever:
     def _tokenize(self, s): 
         return self._normalize(s).split()
 
-    def search(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
-        q_norm = self._normalize(query)
-        q_emb = self.embed.encode([q_norm], normalize_embeddings=True)
+    def _faiss_hits(self, query, k):
+        if self.faiss is None or self.embed is None:
+            return []
+        q = self._normalize(query)
+        q_emb = self.embed.encode([q], normalize_embeddings=True)
         D, I = self.faiss.search(np.asarray(q_emb, dtype="float32"), max(k*4, k))
-        faiss_hits = [self.chunks[i] for i in I[0] if i >= 0]
+        return [self.chunks[i] for i in I[0] if i >= 0]
 
-        # optional BM25
-        bm25_hits = []
-        if self.bm25 is not None:
-            bm25_scores = self.bm25.get_scores(self._tokenize(q_norm))
-            bm25_ranked = np.argsort(-bm25_scores)[:k*2]
-            bm25_hits = [self.chunks[i] for i in bm25_ranked]
+    def _bm25_hits(self, query, k):
+        if self.bm25 is None:
+            return []
+        scores = self.bm25.get_scores(self._tokenize(query))
+        order = np.argsort(-scores)[:k*2]
+        return [self.chunks[i] for i in order]
 
-        # merge uniques (prefer FAISS order)
-        seen = set(); merged=[]
+    def search(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
+        faiss_hits = self._faiss_hits(query, k)
+        bm25_hits = self._bm25_hits(query, k)
+        pool = []
+        seen = set()
         for h in faiss_hits + bm25_hits:
             key = (h["source_row"], h["chunk_id"])
             if key in seen: 
                 continue
-            seen.add(key); merged.append(h)
-            if len(merged) >= k:
+            seen.add(key)
+            pool.append(h)
+            if len(pool) >= k:
                 break
-        return merged
+        # if everything missing, just return first k docs to avoid empty UI
+        return pool or self.chunks[:k]