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

src/kpaa/embeddings/embedder.py

@@ -0,0 +1,81 @@
+"""sentence-transformers 기반 임베더 wrapper (lazy singleton).
+
+기본 모델: BAAI/bge-m3 (1024 dim, multilingual). 한국어 SOTA급.
+
+디바이스 자동 감지 (CUDA → MPS → CPU). KPAA_EMBED_DEVICE 로 강제 가능.
+모델 로드는 첫 호출 시 lazy — import 만으로는 다운로드 발생 X.
+"""
+from __future__ import annotations
+
+import logging
+import os
+from functools import cached_property
+from typing import ClassVar, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import numpy as np
+    from sentence_transformers import SentenceTransformer
+
+logger = logging.getLogger("kpaa.embeddings.embedder")
+
+_DEFAULT_MODEL = "BAAI/bge-m3"
+_DIM_BY_MODEL: dict[str, int] = {
+    "BAAI/bge-m3": 1024,
+}
+
+
+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"
+
+
+class Embedder:
+    """BGE-M3 (또는 KPAA_EMBEDDER 지정 모델) singleton."""
+
+    _instance: ClassVar["Embedder | None"] = None
+
+    def __init__(self, model_name: str | None = None, device: str | None = None) -> None:
+        self.model_name = model_name or os.environ.get("KPAA_EMBEDDER", _DEFAULT_MODEL)
+        self.device = device or _detect_device()
+
+    @classmethod
+    def default(cls) -> "Embedder":
+        if cls._instance is None:
+            cls._instance = cls()
+        return cls._instance
+
+    @cached_property
+    def model(self) -> "SentenceTransformer":
+        from sentence_transformers import SentenceTransformer
+        logger.info("Loading embedding model %s on %s ...", self.model_name, self.device)
+        return SentenceTransformer(self.model_name, device=self.device)
+
+    @property
+    def dim(self) -> int:
+        return _DIM_BY_MODEL.get(self.model_name) or self.model.get_sentence_embedding_dimension()
+
+    def encode_chunks(self, texts: list[str], *, batch: int = 32, show_progress: bool = True) -> "np.ndarray":
+        """문서 측 임베딩. cosine 검색 위해 정규화."""
+        return self.model.encode(
+            texts,
+            batch_size=batch,
+            normalize_embeddings=True,
+            show_progress_bar=show_progress,
+            convert_to_numpy=True,
+        )
+
+    def encode_query(self, text: str) -> "np.ndarray":
+        """쿼리 측 임베딩."""
+        return self.model.encode(
+            text,
+            normalize_embeddings=True,
+            convert_to_numpy=True,
+            show_progress_bar=False,
+        )