"""Cross-encoder reranking of a semantic pool (free: local GPU sidecar). Why this stage exists: IRIS measured the cross-encoder as the dominant ranking channel behind its semantic doubling (0.121 -> 0.244), fused at ~3.3x the weight of any retrieval channel. The sidecar costs nothing per query (owned/rented GPU, already running), so on the score-per-dollar objective it is free accuracy: the one lever that improves the numerator without touching the denominator. Contract: wraps any SemPool (duck-typed: .ranked(), .evidence(cid), .texts(cid), plus optional .title/.year accessors) and returns an object with the same surface whose .ranked() is CE-reordered. Soft-fails to the original pool on any transport error, so a dropped tunnel degrades rather than breaks. """ from __future__ import annotations import json import os import urllib.error import urllib.request # Docs sent to the CE. This was the real bottleneck, and it is not a ranking # problem at all -- measured on the full slice: # nDCG 0.705 | recall@K 0.1497 | recall@FULL 0.2145 # i.e. 69.8% of the golds we SUBMIT already reach the tiny top-K window, so # ordering inside the submission is fine. But pool gold recall is 0.56-0.69 # while only 21% of golds reach the 250 submitted: the pool -> submission # truncation throws away two thirds of what retrieval already found. # With a cap of 400 over pools of 2000-4400, the cross-encoder never saw ~90% # of the pool, and everything below kept raw fusion order. The CE runs on our # own GPU (600 docs in ~0.6s), so extra depth here is negligible under the # benchmark's cost accounting. RERANK_CAP = int((os.environ.get("PFBMAX_CE_CAP") or "400").strip() or 400) DOC_CHARS = 1400 # per-doc text budget TIMEOUT_S = 180 def _post(url: str, payload: dict, timeout: float) -> dict: req = urllib.request.Request( url.rstrip("/") + "/rerank", data=json.dumps(payload).encode("utf-8"), headers={"Content-Type": "application/json"}, ) with urllib.request.urlopen(req, timeout=timeout) as resp: return json.loads(resp.read().decode("utf-8")) def rerank_scores(url: str, query: str, docs: list[dict], timeout: float = TIMEOUT_S) -> dict[str, float]: """{id: score} from the sidecar; {} on any failure (caller keeps order).""" if not url or not docs: return {} try: body = _post(url, {"query": query, "documents": docs}, timeout) except (urllib.error.URLError, OSError, ValueError, TimeoutError): return {} out: dict[str, float] = {} for row in body.get("scores") or []: if isinstance(row, dict) and "id" in row: try: out[str(row["id"])] = float(row["score"]) except (TypeError, ValueError): continue return out class RerankedPool: """SemPool passthrough with a CE-reordered .ranked().""" def __init__(self, pool, order: list[str], trace: dict | None = None): self._pool = pool self._order = order self.ce_trace = trace or {} def ranked(self) -> list[str]: return list(self._order) def __getattr__(self, name): # evidence/texts/title/year/size/... return getattr(self._pool, name) def ce_rerank_pool(query: str, criteria: list[str], pool, url: str | None = None, cap: int = RERANK_CAP, trace: dict | None = None): """Reorder the pool's head by cross-encoder score, keeping the tail in fused order. The CE query carries the criteria: PFB relevance is conjunctive (a paper must satisfy every criterion to earn the scorer's top label), so scoring against the bare query alone under-weights the constraints that decide the label.""" url = url or os.environ.get("IRIS_ASTA_RERANKER_URL", "") tr = trace if trace is not None else {} order = list(pool.ranked()) if not url or not order: tr["ce"] = {"status": "skipped-no-url" if not url else "empty-pool"} return pool head = order[:cap] docs = [] for cid in head: try: text = pool.evidence(cid) or "" except Exception: text = "" title = "" for attr in ("title",): getter = getattr(pool, attr, None) if callable(getter): try: title = getter(cid) or "" except Exception: title = "" blob = (title + "\n" + text).strip()[:DOC_CHARS] if blob: docs.append({"id": cid, "text": blob}) # Criteria ARE appended to the CE query by default (PFBMAX_CE_CRITERIA # defaults to on; set it to 0 to score against the bare query). The # evidence cuts both ways, so the history is worth recording. Against # appending: bge is trained on natural (query, passage) pairs, and gluing # four criteria onto the query makes a long unnatural string that dilutes # the signal. Measured on held-out labelled pairs (recall@K, K=2P), 14 # queries: # bare query 0.2710 | concat query+criteria 0.2404 (-0.0306) # Scoring each criterion separately and taking the MIN, the true # conjunction, is worse still (0.2045): the per-criterion scores are # noisy and a min over five of them just selects the worst noise. # Per-criterion MEAN (0.2518) also loses to the bare query. However, # dropping the criteria, after winning that pair-level test (+0.031), # LOST end-to-end (-0.0365, 3up/4dn, n=8). That was the third time a # pair-level win failed to transfer end-to-end, after the fine-tuned CE # (-0.083) and its distilled successor. So the default stays ON, and the # lesson stands: trust only end-to-end A/Bs. use_crit = os.environ.get("PFBMAX_CE_CRITERIA", "1").strip() in ("1", "true", "yes") ce_query = (query + " || " + " || ".join(criteria[:4]) if (use_crit and criteria) else query) scores = rerank_scores(url, ce_query, docs) if not scores: tr["ce"] = {"status": "error", "sent": len(docs)} return pool # Optional second reranker, fused by reciprocal rank rather than replacing. # Measured motivation: the fine-tuned cross-encoder beats stock on 9 of 12 # held-out queries (+0.0302 excluding its worst) but occasionally produces a # pathological ordering that takes a query from 0.2660 to 0.0000 outright. # Letting either model alone decide the head is what makes that failure # total; RRF lets one model's confident ranking survive the other's # collapse, so the blend keeps the upside without the cliff. url2 = os.environ.get("PFBMAX_RERANKER_URL_2", "").strip() fused_order = None if url2: scores2 = rerank_scores(url2, ce_query, docs) if scores2: # LEXICOGRAPHIC, not additive. The fine-tuned model suffers score # COLLAPSE -- measured sd 0.72 vs stock's 3.23, with only 29-47% # distinct values against stock's 91-98% -- so large groups tie and # their order inside a tie is arbitrary. Since K ~ 2P, the scored # window sits entirely inside that arbitrary zone. # So: the primary model sets a coarse band, the secondary orders # WITHIN the band. Averaging the two (RRF / z-sum) instead lets a # degenerate ordering corrupt the head of both, which is exactly # what it did live (0.1014 vs 0.1588 stock). # Selected offline on held-out labelled pairs (recall@K, K=2P): # stock 0.2710 | rrf 0.2908 | z-sum 0.2858 # trained 0.3314 | lex band .25 0.3319 | lex band .10 0.3364 band = float(os.environ.get("PFBMAX_CE_BAND", "0.10") or 0.10) fused_order = sorted( (c for c in head if c in scores or c in scores2), key=lambda c: (-round(scores.get(c, -1e9) / band), -scores2.get(c, -1e9))) tr.setdefault("ce_fuse", {}).update( {"mode": "lex", "band": band, "n2": len(scores2)}) if fused_order is not None: scored = fused_order unscored = [c for c in head if c not in set(fused_order)] new_order = scored + unscored + order[cap:] tr["ce"] = {"status": "ok-lex", "sent": len(docs), "scored": len(scored)} return RerankedPool(pool, new_order, tr["ce"]) scored = [c for c in head if c in scores] unscored = [c for c in head if c not in scores] scored.sort(key=lambda c: -scores[c]) new_order = scored + unscored + order[cap:] tr["ce"] = {"status": "ok", "sent": len(docs), "scored": len(scored), "moved_into_top50": len([c for c in scored[:50] if order.index(c) >= 50])} return RerankedPool(pool, new_order, tr["ce"])