Spaces:
Sleeping
Sleeping
Upd reranker declaration
Browse files- 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 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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.
|
| 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]
|