| """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) |
| 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)) |
| 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 |
|
|
