Hybrid RAG: BM25+Dense (sqlite-vec/BGE-M3) + cross-encoder reranker (bge-reranker-v2-m3)

src/kpaa/retrieval/reranker.py

@@ -0,0 +1,106 @@
+"""Cross-encoder reranker — Hybrid 후보(top-N)를 정밀 정렬해 LLM context 정확도 ↑.
+
+기본 모델: BAAI/bge-reranker-v2-m3 (568M, multilingual, 한국어 SOTA).
+
+사용:
+    rr = Reranker.default()           # 환경변수로 disable 시 None
+    if rr:
+        excerpts = rr.rerank(query, excerpts, text_fn=lambda e: e.content, top_k=5)
+
+환경변수:
+    KPAA_RERANKER          기본 BAAI/bge-reranker-v2-m3, "off" 면 비활성
+    KPAA_EMBED_DEVICE      embedder 와 공유 (auto / cuda / mps / cpu)
+
+비용 (BGE-reranker-v2-m3, 후보 20개 기준):
+    - M2 MPS:  ~50-100ms
+    - CPU:     ~250-450ms
+    - A10g:    ~30-60ms
+"""
+from __future__ import annotations
+
+import logging
+import os
+from functools import cached_property
+from typing import Callable, ClassVar, TypeVar
+
+logger = logging.getLogger("kpaa.retrieval.reranker")
+
+_DEFAULT_MODEL = "BAAI/bge-reranker-v2-m3"
+T = TypeVar("T")
+
+
+def _detect_device() -> str:
+    forced = os.environ.get("KPAA_EMBED_DEVICE", "auto").lower()
+    if forced != "auto":
+        return forced
+    import torch
+    if torch.cuda.is_available():
+        return "cuda"
+    if torch.backends.mps.is_available():
+        return "mps"
+    return "cpu"
+
+
+def _disabled() -> bool:
+    return os.environ.get("KPAA_RERANKER", "").lower() in ("off", "0", "false", "no", "disabled")
+
+
+class Reranker:
+    """sentence-transformers CrossEncoder wrapper (lazy singleton)."""
+
+    _instance: ClassVar["Reranker | None"] = None
+    _missing: ClassVar[bool] = False  # 한 번 실패하면 retry 안 함
+
+    def __init__(self, model_name: str | None = None, device: str | None = None) -> None:
+        self.model_name = model_name or os.environ.get("KPAA_RERANKER", _DEFAULT_MODEL)
+        self.device = device or _detect_device()
+
+    @classmethod
+    def default(cls) -> "Reranker | None":
+        """싱글톤 인스턴스 — disabled 또는 첫 로드 실패 시 None.
+
+        retriever 가 매 요청마다 호출하므로 None 반환 시 BM25+Dense 결과 그대로 사용.
+        """
+        if _disabled() or cls._missing:
+            return None
+        if cls._instance is None:
+            cls._instance = cls()
+        return cls._instance
+
+    @cached_property
+    def model(self):
+        from sentence_transformers import CrossEncoder
+
+        logger.info("Loading reranker %s on %s ...", self.model_name, self.device)
+        try:
+            return CrossEncoder(self.model_name, device=self.device)
+        except Exception as e:
+            logger.warning("Reranker load failed (%s) — disabling", e)
+            type(self)._missing = True
+            type(self)._instance = None
+            raise
+
+    def rerank(
+        self,
+        query: str,
+        candidates: list[T],
+        *,
+        text_fn: Callable[[T], str],
+        top_k: int = 5,
+    ) -> list[T]:
+        """`(query, text_fn(c))` 쌍을 cross-encoder 로 점수화해 정렬한 뒤 top_k.
+
+        cross-encoder 점수 ↑ = 더 관련성 높음. 후보가 top_k 이하면 그대로 반환.
+        """
+        if not candidates or top_k <= 0:
+            return list(candidates[:top_k])
+        if len(candidates) <= top_k:
+            return list(candidates)
+        try:
+            pairs = [(query, text_fn(c)) for c in candidates]
+            scores = self.model.predict(pairs, show_progress_bar=False)
+            ranked = sorted(zip(candidates, scores), key=lambda x: -float(x[1]))
+            return [c for c, _ in ranked[:top_k]]
+        except Exception as e:  # noqa: BLE001
+            logger.warning("Reranker.predict failed (%s) — falling back to original order", e)
+            return list(candidates[:top_k])