bge.py 1.13 KB
"""
BGE reranker backend - wraps BGEReranker for the unified backend protocol.
"""

from __future__ import annotations

from typing import Any, Dict, List

from reranker.bge_reranker import BGEReranker


class BGERerankerBackend:
    """BGE reranker backend; config from services.rerank.backends.bge."""

    def __init__(self, config: Dict[str, Any]) -> None:
        self._config = config or {}
        self._impl = BGEReranker(
            model_name=str(self._config.get("model_name") or "BAAI/bge-reranker-v2-m3"),
            device=self._config.get("device"),
            batch_size=int(self._config.get("batch_size", 64)),
            use_fp16=bool(self._config.get("use_fp16", True)),
            max_length=int(self._config.get("max_length", 512)),
            cache_dir=str(self._config.get("cache_dir") or "./model_cache"),
            enable_warmup=bool(self._config.get("enable_warmup", True)),
        )

    def score_with_meta(
        self,
        query: str,
        docs: List[str],
        normalize: bool = True,
    ) -> tuple[list[float], Dict[str, Any]]:
        return self._impl.score_with_meta(query, docs, normalize=normalize)