701ae503
 tangwang
docs
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
|
"""
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)
|