BinKhoaLe1812 commited on
Commit
3b787c4
·
verified ·
1 Parent(s): d7cdcb0

Upd reranker declaration

Browse files
Files changed (1) hide show
  1. api/retrieval.py +10 -3
api/retrieval.py CHANGED
@@ -14,7 +14,14 @@ logger = logging.getLogger("retrieval-bot")
14
  class RetrievalEngine:
15
  def __init__(self):
16
  self.db_manager = db_manager
17
- self._reranker = _NvidiaReranker()
 
 
 
 
 
 
 
18
 
19
  @staticmethod
20
  def _is_cpg_text(text: str) -> bool:
@@ -117,7 +124,7 @@ class RetrievalEngine:
117
  cpg_candidates = [t for t in kept if self._is_cpg_text(t)]
118
  if cpg_candidates:
119
  logger.info("[Retrieval] CPG content detected; invoking NVIDIA reranker")
120
- reranked = self._reranker.rerank(query, cpg_candidates)
121
  # Keep only valid high-scoring items
122
  filtered: List[Dict] = [r for r in reranked if r.get("score", 0) >= 0.3 and r.get("text")]
123
  # Limit to top 3 for prompt efficiency
@@ -220,4 +227,4 @@ class _NvidiaReranker:
220
  except Exception as e:
221
  logger.warning(f"[Reranker] Failed calling NVIDIA reranker: {e}")
222
  # On failure, return original order with neutral scores
223
- return [{"text": d, "score": 0.0} for d in documents]
 
14
  class RetrievalEngine:
15
  def __init__(self):
16
  self.db_manager = db_manager
17
+ # Lazy-init reranker to avoid NameError during module import ordering
18
+ self._reranker = None
19
+
20
+ def _get_reranker(self):
21
+ """Initialize the NVIDIA reranker on first use."""
22
+ if self._reranker is None:
23
+ self._reranker = _NvidiaReranker()
24
+ return self._reranker
25
 
26
  @staticmethod
27
  def _is_cpg_text(text: str) -> bool:
 
124
  cpg_candidates = [t for t in kept if self._is_cpg_text(t)]
125
  if cpg_candidates:
126
  logger.info("[Retrieval] CPG content detected; invoking NVIDIA reranker")
127
+ reranked = self._get_reranker().rerank(query, cpg_candidates)
128
  # Keep only valid high-scoring items
129
  filtered: List[Dict] = [r for r in reranked if r.get("score", 0) >= 0.3 and r.get("text")]
130
  # Limit to top 3 for prompt efficiency
 
227
  except Exception as e:
228
  logger.warning(f"[Reranker] Failed calling NVIDIA reranker: {e}")
229
  # On failure, return original order with neutral scores
230
+ return [{"text": d, "score": 0.0} for d in documents]