steve
6dfc13be1d
Implements five of the six enhancements from docs/kb-enhancements-proposal.htm, closing the retrieval-quality gap identified in the qmd review. - Cross-encoder reranking: new kb/reranker.py loads an optional reranking model at startup (KB_RERANK_ENABLED, KB_RERANKER_MODEL, KB_RERANK_CANDIDATES). Search degrades gracefully to plain hybrid retrieval when the model is absent. Exposed via a "rerank" block in /status, a rerank flag on search, and --no-rerank in the CLI. - RRF rank fusion: FTS and vector lists now merge by reciprocal rank fusion with a top-rank bonus, replacing the old score blend. Scores are comparable across queries. - Bench harness and explain traces: kb bench runs a query fixture against each backend and reports precision@k, recall and MRR. --explain returns a per-result score breakdown. - Tag context descriptions: tags carry an optional one-line description (kb tag-describe), returned as tag_contexts with search results. Adds a tags.description column migration. - Structured data ingestion: .json/.yaml/.toml files ingest as text via the new "data" doc type, pretty-printing minified JSON before chunking. Query expansion (proposal item 5) is deliberately left out pending bench results. Requires engine v3.3.0. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
64 lines
2.0 KiB
Python
64 lines
2.0 KiB
Python
"""Cross-encoder reranker management.
|
|
|
|
Mirrors the embeddings module: one module-level model, loaded eagerly at
|
|
startup when reranking is enabled. Reranking is strictly optional — search
|
|
degrades gracefully to plain hybrid retrieval when the model is absent.
|
|
"""
|
|
|
|
import logging
|
|
from typing import Optional
|
|
|
|
logger = logging.getLogger("kb.reranker")
|
|
|
|
_reranker: Optional[object] = None
|
|
_model_name: Optional[str] = None
|
|
|
|
|
|
def load_reranker(model_name: str, device: str = "cpu") -> None:
|
|
"""Load a cross-encoder reranking model.
|
|
|
|
Args:
|
|
model_name: HuggingFace model name or local path. Must have a
|
|
sequence-classification head (e.g. BAAI/bge-reranker-v2-m3).
|
|
device: Target device — "cpu", "cuda", or "auto".
|
|
"""
|
|
global _reranker, _model_name
|
|
|
|
from sentence_transformers import CrossEncoder
|
|
|
|
from kb.embeddings import _resolve_device
|
|
|
|
resolved_device = _resolve_device(device)
|
|
logger.info("Loading reranker '%s' on device '%s'", model_name, resolved_device)
|
|
|
|
_reranker = CrossEncoder(model_name, device=resolved_device)
|
|
_model_name = model_name
|
|
|
|
logger.info("Reranker loaded: %s", model_name)
|
|
|
|
|
|
def is_available() -> bool:
|
|
"""Return True if a reranker model is loaded and usable."""
|
|
return _reranker is not None
|
|
|
|
|
|
def rerank_scores(query: str, texts: list[str]) -> list[float]:
|
|
"""Score (query, text) pairs with the cross-encoder.
|
|
|
|
Returns:
|
|
One relevance score per text, sigmoid-normalised to 0-1.
|
|
|
|
Raises:
|
|
RuntimeError: If no reranker has been loaded.
|
|
"""
|
|
if _reranker is None:
|
|
raise RuntimeError("Reranker not loaded. Call load_reranker() first.")
|
|
|
|
import numpy as np
|
|
|
|
scores = _reranker.predict([(query, t) for t in texts], convert_to_numpy=True)
|
|
# CrossEncoder heads may emit raw logits; squash to 0-1 so scores blend
|
|
# predictably with normalised retrieval scores. Sigmoid is monotonic, so
|
|
# ordering is unaffected for models that already output probabilities.
|
|
return (1.0 / (1.0 + np.exp(-np.asarray(scores, dtype="float64")))).tolist()
|