File size: 3,525 Bytes
66cf40d 190f0fc 66cf40d 190f0fc 66cf40d a41b023 66cf40d a41b023 66cf40d a41b023 66cf40d | 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 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | """Step 09 — cross-encoder rerank (granite-embedding-reranker-r2 ONNX): the most precise relevance signal.
A cross-encoder reads the JD query and each candidate's text TOGETHER (one transformer pass per pair) and
emits a single relevance logit — strictly more precise than the bi-encoder cosine of Step 08, and the most
expensive per item, so it runs only on the ~300 granite survivors. ModernBERT sequence-classification head;
onnxruntime + tokenizers (NO torch). Re-ranks by ``reranker_score`` (raw logit; monotonic for ordering).
"""
from __future__ import annotations
from functools import lru_cache
from typing import Any
import numpy as np
import polars as pl
from common.logging import step
from config import Settings, load_settings
from .onnx import load_onnx_session
@lru_cache(maxsize=1)
def _load_reranker(repo: str, onnx_file: str, max_length: int, cpu_threads: int = 0) -> tuple[Any, Any, set[str]]:
"""Load (downloading on first use) the cross-encoder ONNX session + tokenizer; cached for the process."""
return load_onnx_session(repo, onnx_file, max_length, "reranker.load", cpu_threads)
def _score_pairs(
session: Any, tokenizer: Any, input_names: set[str], query: str, docs: list[str], batch_size: int
) -> np.ndarray:
"""Cross-encode (query, doc) pairs → 1-D array of relevance logits, in doc order."""
logits: list[np.ndarray] = []
for start in range(0, len(docs), batch_size):
batch_pairs = [(query, doc) for doc in docs[start : start + batch_size]]
tokenizer.enable_padding()
encodings = tokenizer.encode_batch(batch_pairs) # pair encoding: [CLS] query [SEP] doc [SEP]
token_ids = np.array([enc.ids for enc in encodings], dtype=np.int64)
feed: dict[str, np.ndarray] = {
"input_ids": token_ids,
"attention_mask": np.array([enc.attention_mask for enc in encodings], dtype=np.int64),
}
if "token_type_ids" in input_names:
feed["token_type_ids"] = np.zeros_like(token_ids)
batch_logits = session.run(None, {k: v for k, v in feed.items() if k in input_names})[0]
logits.append(np.asarray(batch_logits, dtype=np.float64).reshape(-1)) # (batch, 1) → (batch,)
return np.concatenate(logits) if logits else np.zeros(0)
def cross_encoder_rerank(
candidates: pl.DataFrame, query_text: str, settings: Settings | None = None
) -> pl.DataFrame:
"""Add ``reranker_score`` (cross-encoder relevance) and re-sort by it. Empty/blank ⇒ passthrough at 0."""
active = settings or load_settings()
rerank = active.rerank
with step("rank.cross_encoder", rows_in=candidates.height) as metrics:
if candidates.height == 0 or not query_text.strip():
metrics["rows_out"] = candidates.height
metrics["skipped"] = True
return candidates.with_columns(pl.lit(0.0).alias("reranker_score"))
session, tokenizer, input_names = _load_reranker(
active.models.reranker, rerank.onnx_file, rerank.max_length, active.compute.cpu_threads
)
scores = _score_pairs(
session, tokenizer, input_names, query_text, candidates.get_column("text").to_list(), rerank.batch_size
)
reranked = candidates.with_columns(pl.Series("reranker_score", scores)).sort(
["reranker_score", "candidate_id"], descending=[True, False]
)
metrics["rows_out"] = reranked.height
metrics["score_max"] = round(float(scores.max()), 4)
return reranked
|