diff --git a/application/cmd/cre_main.py b/application/cmd/cre_main.py index 4b814640a..f42a4a5c9 100644 --- a/application/cmd/cre_main.py +++ b/application/cmd/cre_main.py @@ -1073,10 +1073,10 @@ def run_librarian( For each knowledge-queue section: try the deterministic explicit-CRE fast path (C.0.5); on no/ambiguous reference, run the semantic retriever (C.1) - and log the top-K candidate shortlist. The cross-encoder rerank (C.2, W4), - decision/threshold routing (C.3-C.4, W5) and graph writes (W8) are not - built yet, so this is dry-run only: it never writes a link. ``--run_librarian`` - without writes behaves identically and warns. + and rerank its top-K shortlist with the cross-encoder (C.2, W4), then log + the reranked candidates. Decision/threshold routing (C.3-C.4, W5) and graph + writes (W8) are not built yet, so this is dry-run only: it never writes a + link. ``--run_librarian`` without writes behaves identically and warns. Ops note: this is opt-in CLI only — it is not on the ``Procfile`` and is wired into neither the web app nor the background worker (that lands W8). @@ -1089,6 +1089,10 @@ def run_librarian( build_retriever, ) from application.utils.librarian.config_loader import load_config + from application.utils.librarian.cross_encoder import ( + CrossEncoderReranker, + build_cross_encoder_score_fn, + ) from application.utils.librarian.explicit_link_resolver import ( ResolutionOutcome, resolve, @@ -1141,6 +1145,16 @@ def run_librarian( connection=database.session.connection(), ) + # C.2 reranker: reads each (section, candidate-CRE) pair together and + # re-sorts C.1's shortlist. Scored against the CRE's embeddings_content — + # the same text the hub vectors were built from. + cre_texts = database.get_embedding_contents_by_doc_type(defs.Credoctypes.CRE.value) + reranker = CrossEncoderReranker( + score_fn=build_cross_encoder_score_fn(cfg.crossencoder_model), + top_n=cfg.top_k_rerank, + cre_texts=cre_texts, + ) + source = FixtureKnowledgeSource(source_jsonl or _DEFAULT_LIBRARIAN_SOURCE) sections = explicit = semantic = rejected = 0 for item in source.items(): @@ -1158,15 +1172,26 @@ def run_librarian( logger.info("[explicit] %s -> %s", section.chunk_id, resolution.cre_ids[0]) continue + try: + audit = retriever.retrieve(section.text) + audit = reranker.rerank(section.text, audit) + # one bad section must not abort the batch + except Exception as exc: # noqa: BLE001 + rejected += 1 + logger.warning( + "[semantic] %s skipped: retrieval/rerank failed: %s", + section.chunk_id, + exc, + ) + continue + semantic += 1 - audit = retriever.retrieve(section.text) - top = ", ".join( - f"{c.cre_id}:{c.score_vector:.3f}" for c in audit.candidates[:5] - ) + top = ", ".join(f"{c.cre_id}:{c.score_rerank:.3f}" for c in audit.reranked) logger.info( - "[semantic] %s -> %d candidates (top5: %s)", + "[semantic] %s -> %d candidates, reranked top%d: %s", section.chunk_id, len(audit.candidates), + len(audit.reranked), top or "", ) diff --git a/application/database/db.py b/application/database/db.py index 6600c3920..e3290e5bc 100644 --- a/application/database/db.py +++ b/application/database/db.py @@ -2376,6 +2376,27 @@ def get_embeddings_by_doc_type(self, doc_type: str) -> Dict[str, List[float]]: res[entry.node_id] = [float(e) for e in entry.embeddings.split(",")] return res + def get_embedding_contents_by_doc_type(self, doc_type: str) -> Dict[str, str]: + """Return ``{id -> embeddings_content}`` for a doc type. + + ``embeddings_content`` is the text a vector was built from — the pair + text the Module C.2 cross-encoder scores a candidate against. Keyed by + cre_id for CREs, node_id otherwise, mirroring + ``get_embeddings_by_doc_type``. Rows with no content are skipped. + """ + res: Dict[str, str] = {} + embeddings = ( + self.session.query(Embeddings).filter(Embeddings.doc_type == doc_type).all() + ) + for entry in embeddings: + if not entry.embeddings_content: + continue + if doc_type == cre_defs.Credoctypes.CRE.value: + res[entry.cre_id] = entry.embeddings_content + else: + res[entry.node_id] = entry.embeddings_content + return res + def get_embeddings_by_doc_type_paginated( self, doc_type: str, page: int = 1, per_page: int = 100 ) -> Tuple[Dict[str, List[float]], int, int]: diff --git a/application/tests/librarian/config_loader_test.py b/application/tests/librarian/config_loader_test.py index 1b7abb294..27000a76c 100644 --- a/application/tests/librarian/config_loader_test.py +++ b/application/tests/librarian/config_loader_test.py @@ -1,6 +1,7 @@ import os import unittest from dataclasses import FrozenInstanceError +from typing import ClassVar, Dict from unittest import mock from application.utils.librarian.config_loader import LibrarianConfig, load_config @@ -27,7 +28,7 @@ def test_config_is_frozen(self): class TestConfigLoaderOverrides(unittest.TestCase): - OVERRIDES = { + OVERRIDES: ClassVar[Dict[str, str]] = { "CRE_LIBRARIAN_RETRIEVER_BACKEND": "pgvector", "CRE_LIBRARIAN_CROSSENCODER_MODEL": "cross-encoder/other", "CRE_LIBRARIAN_TOP_K_RETRIEVAL": "50", diff --git a/application/tests/librarian/cross_encoder_test.py b/application/tests/librarian/cross_encoder_test.py new file mode 100644 index 000000000..8813554d3 --- /dev/null +++ b/application/tests/librarian/cross_encoder_test.py @@ -0,0 +1,119 @@ +"""Tests for the C.2 cross-encoder reranker (Week 4). + +Hermetic: the cross-encoder score function and the CRE text map are injected as +controlled values, so re-ordering, top-N truncation, tie stability, the audit +shape, and the failure modes are all assertable without torch, an LLM, or a DB. +""" + +import unittest + +from application.utils.librarian.cross_encoder import ( + CrossEncoderReranker, + MissingCandidateTextError, + RerankerError, +) +from application.utils.librarian.schemas import CreCandidate, RetrievalAudit + +# The text each candidate CRE is scored against (its embeddings_content). +CRE_TEXTS = { + "111-111": "text-a", + "222-222": "text-b", + "333-333": "text-c", +} + + +def make_audit(threshold=0.8): + """A C.1 shortlist in cosine order: a (0.90) > b (0.80) > c (0.70).""" + return RetrievalAudit( + retriever="in-memory-cosine/0.1.0", + candidates=[ + CreCandidate(cre_id="111-111", score_vector=0.90), + CreCandidate(cre_id="222-222", score_vector=0.80), + CreCandidate(cre_id="333-333", score_vector=0.70), + ], + reranked=[], + threshold=threshold, + ) + + +# The cross-encoder disagrees with cosine: it likes b most, then c, then a. +_PAIR_SCORES = { + ("q", "text-a"): 0.10, + ("q", "text-b"): 0.90, + ("q", "text-c"): 0.50, +} + + +def fake_score(pairs): + return [_PAIR_SCORES[(q, cand)] for q, cand in pairs] + + +def make_reranker(top_n=5, cre_texts=None, score_fn=fake_score): + return CrossEncoderReranker( + score_fn=score_fn, + top_n=top_n, + cre_texts=CRE_TEXTS if cre_texts is None else cre_texts, + ) + + +class RerankTest(unittest.TestCase): + def test_reranks_by_cross_encoder_score_not_cosine(self) -> None: + audit = make_reranker().rerank("q", make_audit()) + # Cosine order was a,b,c; the cross-encoder reorders to b,c,a. + self.assertEqual( + [c.cre_id for c in audit.reranked], ["222-222", "333-333", "111-111"] + ) + scores = [c.score_rerank for c in audit.reranked] + self.assertEqual(scores, sorted(scores, reverse=True)) + self.assertAlmostEqual(audit.reranked[0].score_rerank, 0.90) + + def test_candidates_preserved_and_audit_shape(self) -> None: + audit = make_reranker().rerank("q", make_audit(threshold=0.8)) + # candidates[] (the pre-rerank shortlist) is untouched, still cosine order. + self.assertEqual( + [c.cre_id for c in audit.candidates], ["111-111", "222-222", "333-333"] + ) + # score_rerank is only set on reranked[], never on candidates[]. + self.assertTrue(all(c.score_rerank is None for c in audit.candidates)) + self.assertEqual(audit.retriever, "in-memory-cosine/0.1.0") + self.assertEqual(audit.threshold, 0.8) + + def test_top_n_truncates_to_best(self) -> None: + audit = make_reranker(top_n=2).rerank("q", make_audit()) + self.assertEqual([c.cre_id for c in audit.reranked], ["222-222", "333-333"]) + + def test_top_n_larger_than_shortlist_keeps_all(self) -> None: + audit = make_reranker(top_n=20).rerank("q", make_audit()) + self.assertEqual(len(audit.reranked), 3) + + def test_ties_preserve_cosine_order(self) -> None: + # Equal cross-encoder scores -> the stable sort keeps C.1's cosine order. + audit = make_reranker(score_fn=lambda pairs: [1.0] * len(pairs)).rerank( + "q", make_audit() + ) + self.assertEqual( + [c.cre_id for c in audit.reranked], ["111-111", "222-222", "333-333"] + ) + + def test_empty_candidates_yields_empty_reranked(self) -> None: + empty = RetrievalAudit(retriever="r", candidates=[], reranked=[], threshold=0.8) + out = make_reranker().rerank("q", empty) + self.assertEqual(out.reranked, []) + + +class FailureModeTest(unittest.TestCase): + def test_missing_cre_text_raises(self) -> None: + with self.assertRaises(MissingCandidateTextError): + make_reranker(cre_texts={"111-111": "text-a"}).rerank("q", make_audit()) + + def test_non_positive_top_n_rejected(self) -> None: + with self.assertRaises(RerankerError): + make_reranker(top_n=0) + + def test_score_count_mismatch_rejected(self) -> None: + with self.assertRaises(RerankerError): + make_reranker(score_fn=lambda pairs: [0.5]).rerank("q", make_audit()) + + +if __name__ == "__main__": + unittest.main() diff --git a/application/utils/librarian/__init__.py b/application/utils/librarian/__init__.py index f6b8440a5..f168f2bb4 100644 --- a/application/utils/librarian/__init__.py +++ b/application/utils/librarian/__init__.py @@ -15,7 +15,8 @@ re-normalizing text) and ExplicitLinkResolver (fail-safe explicit-link resolution). W3 (C.1): candidate retriever (in-memory + pgvector) + pipeline switch. -Cross-encoder rerank (C.2, W4) onward is not built yet. + W4 (C.2): cross-encoder reranker — re-sorts the C.1 shortlist, fills reranked[]. +Calibration + decision routing (C.3-C.4, W5-W6) onward is not built yet. Vendored RFC JSON schemas live under ``_rfc_schemas/``. They are pinned to upstream/owasp-graph @ 2b1437987768d5ed20fe9ee721ab9a898c4b84af (PR #734). diff --git a/application/utils/librarian/candidate_retriever.py b/application/utils/librarian/candidate_retriever.py index d4e36c12a..7e1c6f93c 100644 --- a/application/utils/librarian/candidate_retriever.py +++ b/application/utils/librarian/candidate_retriever.py @@ -157,7 +157,9 @@ def retrieve(self, text: str) -> RetrievalAudit: # Top-K by descending score. argsort is ascending, so take the tail # and reverse; cap at pool size when the hub is smaller than K. k = min(self._top_k, len(self._pool.cre_ids)) - top_idx = np.argsort(scores)[-k:][::-1] + # Stable descending sort so tied cosine scores keep hub (index) order — + # deterministic across runs. argsort(-scores) descends; kind="stable". + top_idx = np.argsort(-scores, kind="stable")[:k] candidates: List[CreCandidate] = [ CreCandidate( diff --git a/application/utils/librarian/cross_encoder.py b/application/utils/librarian/cross_encoder.py new file mode 100644 index 000000000..2e3f6991c --- /dev/null +++ b/application/utils/librarian/cross_encoder.py @@ -0,0 +1,128 @@ +"""Module C.2 — cross-encoder reranker (Week 4). The careful re-reader. + +C.1 (the bi-encoder, W3) fingerprints the section and each CRE *separately* and +cosine-ranks the whole hub — fast enough to scan every CRE, but it never reads a +section and a candidate *together*, so the ordering inside the top-K shortlist is +rough: the right CRE can sit at #7, not #1. C.2 fixes that ordering. It reads +each ``(section text, candidate CRE text)`` *pair* together as one input, scores +"do these two actually match?", re-sorts the shortlist by that score, and keeps +the best N. Slow per pair, so it runs only over the K candidates C.1 already +narrowed to — never the whole hub. + +Like C.1, the reranker is a thin dependency-injected seam over its model: + + - ``score_fn(pairs) -> Sequence[float]`` — scores a batch of + ``(query_text, candidate_text)`` pairs, higher = better match. Prod wires a + pinned cross-encoder (``ms-marco-MiniLM-L-6-v2``); the harness and tests + inject a deterministic stub. C.2 never imports the model directly, so it + stays import-light and hermetically testable (mirrors C.1's ``embed_fn``). + +The text a candidate CRE is scored against is its ``embeddings_content`` — the +same signal the hub vectors were built from — supplied as a ``{cre_id -> text}`` +map. The RFC is silent on ranking tech; it mandates only the +``candidates[]``/``reranked[]`` audit trail. C.2 fills ``reranked[]`` (the slot +C.1 deliberately left empty), populating ``score_rerank`` and re-ordering, while +leaving ``candidates[]`` untouched so the pre-rerank shortlist stays auditable. +""" + +from typing import Callable, List, Mapping, Sequence, Tuple + +from application.utils.librarian.schemas import RetrievalAudit + +# A function that scores a batch of (query_text, candidate_text) pairs. +RerankFn = Callable[[Sequence[Tuple[str, str]]], Sequence[float]] + +# Identify the reranker in the RFC audit trail. Bumped when the model or the +# scoring changes so a stored proposal is traceable to the code that ranked it. +RERANKER_NAME = "cross-encoder-ms-marco-MiniLM-L-6-v2/0.1.0" +DEFAULT_CROSSENCODER_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2" + + +class RerankerError(ValueError): + """Base class for reranker construction/usage failures.""" + + +class MissingCandidateTextError(RerankerError): + """A shortlisted CRE has no text to score against — cannot rerank fairly.""" + + +class CrossEncoderReranker: + """Re-score a C.1 shortlist by reading each (section, CRE) pair together. + + ``top_n`` is ``CRE_LIBRARIAN_TOP_K_RERANK`` (default 5). ``rerank`` reads + ``audit.candidates`` (C.1's shortlist), scores every pair, re-sorts by the + cross-encoder score, keeps the best ``top_n``, and returns a copy of the + audit with ``reranked`` filled and ``candidates`` preserved. + """ + + def __init__( + self, + score_fn: RerankFn, + top_n: int, + *, + cre_texts: Mapping[str, str], + ) -> None: + if top_n <= 0: + raise RerankerError(f"top_n must be > 0, got {top_n}") + self._score_fn = score_fn + self._top_n = top_n + self._cre_texts = dict(cre_texts) + + def rerank(self, text: str, audit: RetrievalAudit) -> RetrievalAudit: + """Return a copy of ``audit`` with ``reranked`` filled from ``candidates``. + + Raises ``MissingCandidateTextError`` if a shortlisted CRE has no text to + score (a silent-quality trap otherwise) and ``RerankerError`` if the + model returns the wrong number of scores. + """ + candidates = audit.candidates + if not candidates: + # Nothing to rerank (e.g. an empty hub upstream); keep the audit + # shape consistent — an explicit, empty reranked list. + return audit.model_copy(update={"reranked": []}) + + pairs: List[Tuple[str, str]] = [] + for c in candidates: + cre_text = self._cre_texts.get(c.cre_id) + if not cre_text: + raise MissingCandidateTextError( + f"no text for candidate CRE {c.cre_id!r}; the reranker needs " + "the CRE's embeddings_content to score the pair" + ) + pairs.append((text, cre_text)) + + scores = list(self._score_fn(pairs)) + if len(scores) != len(candidates): + raise RerankerError( + f"score_fn returned {len(scores)} scores for {len(candidates)} " + "candidates; the reranker expects exactly one score per pair" + ) + + reranked = [ + c.model_copy(update={"score_rerank": float(s)}) + for c, s in zip(candidates, scores, strict=True) + ] + # Highest cross-encoder score first, then keep only the best top_n. + # Python's sort is stable, so ties preserve C.1's cosine order. + reranked.sort(key=lambda c: c.score_rerank, reverse=True) + return audit.model_copy(update={"reranked": reranked[: self._top_n]}) + + +def build_cross_encoder_score_fn( + model_name: str = DEFAULT_CROSSENCODER_MODEL, +) -> RerankFn: + """Load a sentence-transformers CrossEncoder and adapt it to ``RerankFn``. + + ``sentence_transformers`` (and torch) is imported lazily so this module — + and CI/tests that inject a stub score_fn — never need the heavy ML stack + loaded (mirrors C.1 keeping the embedding model out of module import). + """ + from sentence_transformers import CrossEncoder # lazy, heavy + + model = CrossEncoder(model_name) + + def score_fn(pairs: Sequence[Tuple[str, str]]) -> List[float]: + # CrossEncoder.predict takes a list of [a, b] pairs, returns an ndarray. + return [float(s) for s in model.predict([list(p) for p in pairs])] + + return score_fn diff --git a/application/utils/librarian/knowledge_source.py b/application/utils/librarian/knowledge_source.py index 9f04aece7..9a87ce104 100644 --- a/application/utils/librarian/knowledge_source.py +++ b/application/utils/librarian/knowledge_source.py @@ -38,9 +38,11 @@ def items(self) -> Iterator[KnowledgeQueueItem]: try: yield KnowledgeQueueItem.model_validate_json(line) except ValidationError as exc: + # Log only error locations/types, never the raw input — + # queue rows can carry sensitive content. logger.warning( "Skipping malformed knowledge_queue row at line %d: %s", line_no, - exc, + exc.errors(include_input=False), ) continue diff --git a/application/utils/librarian/schemas.py b/application/utils/librarian/schemas.py index d0a41983d..12dc0a56d 100644 --- a/application/utils/librarian/schemas.py +++ b/application/utils/librarian/schemas.py @@ -22,6 +22,15 @@ _SCHEMA_VERSION_RE = re.compile(r"^0\.\d+\.\d+$") +def _require_schema_version(value: str) -> None: + """Shared envelope check — every RFC envelope pins schema_version to 0.x.y. + + Centralized so the regex and message can't drift across envelopes. + """ + if not _SCHEMA_VERSION_RE.match(value): + raise ValueError(r"schema_version must match ^0\.\d+\.\d+$") + + # ---------- Enums (RFC) ---------- @@ -212,8 +221,7 @@ class KnowledgeItem(BaseModel): @model_validator(mode="after") def _rfc_rules(self) -> "KnowledgeItem": - if not _SCHEMA_VERSION_RE.match(self.schema_version): - raise ValueError(r"schema_version must match ^0\.\d+\.\d+$") + _require_schema_version(self.schema_version) if self.status == KnowledgeStatus.accepted and self.content is None: raise ValueError("status=accepted requires content") if self.status == KnowledgeStatus.rejected and self.rejection is None: @@ -238,8 +246,7 @@ class LinkProposal(BaseModel): @model_validator(mode="after") def _schema_version_pattern(self) -> "LinkProposal": - if not _SCHEMA_VERSION_RE.match(self.schema_version): - raise ValueError(r"schema_version must match ^0\.\d+\.\d+$") + _require_schema_version(self.schema_version) return self @@ -262,8 +269,7 @@ class ReviewItem(BaseModel): @model_validator(mode="after") def _schema_version_pattern(self) -> "ReviewItem": - if not _SCHEMA_VERSION_RE.match(self.schema_version): - raise ValueError(r"schema_version must match ^0\.\d+\.\d+$") + _require_schema_version(self.schema_version) return self diff --git a/application/utils/librarian/section_validator.py b/application/utils/librarian/section_validator.py index 8d92eb082..22927ba0f 100644 --- a/application/utils/librarian/section_validator.py +++ b/application/utils/librarian/section_validator.py @@ -32,9 +32,9 @@ """ from dataclasses import dataclass -from typing import Any, Dict, Optional, Union +from typing import Any, Dict, Optional, Type, TypeVar, Union -from pydantic import ValidationError +from pydantic import BaseModel, ValidationError from application.utils.librarian.schemas import ( KnowledgeItem, @@ -46,6 +46,8 @@ SourceType, ) +_ModelT = TypeVar("_ModelT", bound=BaseModel) + KNOWLEDGE_LABEL = "KNOWLEDGE" # MVP scope: golden dataset and CRE hub vectors are English-only. @@ -86,6 +88,17 @@ class Section: locator: Locator +def _validate_or_raise(model_cls: Type[_ModelT], raw: Any) -> _ModelT: + """Coerce a raw row into ``model_cls``, converting Pydantic's ValidationError + into a typed MalformedKnowledgeItemError so it never escapes this module.""" + if isinstance(raw, model_cls): + return raw + try: + return model_cls.model_validate(raw) + except ValidationError as exc: + raise MalformedKnowledgeItemError(str(exc)) from exc + + def _require_text(text: str) -> str: if not text or not text.strip(): raise EmptyTextError("section text is empty or whitespace-only") @@ -111,11 +124,7 @@ def section_from_queue_row( Raises a SectionValidationError subclass on any rejection. """ - if not isinstance(row, KnowledgeQueueItem): - try: - row = KnowledgeQueueItem.model_validate(row) - except ValidationError as exc: - raise MalformedKnowledgeItemError(str(exc)) from exc + row = _validate_or_raise(KnowledgeQueueItem, row) if row.llm_label != KNOWLEDGE_LABEL: raise NotKnowledgeError( @@ -154,11 +163,7 @@ def section_from_knowledge_item( Raises a SectionValidationError subclass on any rejection. """ - if not isinstance(item, KnowledgeItem): - try: - item = KnowledgeItem.model_validate(item) - except ValidationError as exc: - raise MalformedKnowledgeItemError(str(exc)) from exc + item = _validate_or_raise(KnowledgeItem, item) if item.status != KnowledgeStatus.accepted: raise NotKnowledgeError( diff --git a/requirements.txt b/requirements.txt index 7fa6f6ba1..8edf363ed 100644 --- a/requirements.txt +++ b/requirements.txt @@ -96,6 +96,7 @@ redis ruamel.yaml ruamel.yaml.clib scikit-learn +sentence-transformers>=5.0.0,<6.0.0 Shapely six smmap diff --git a/scripts/build_golden_dataset.py b/scripts/build_golden_dataset.py index 7a371c836..fe8187e97 100644 --- a/scripts/build_golden_dataset.py +++ b/scripts/build_golden_dataset.py @@ -195,19 +195,27 @@ def _fetch_asvs_cre(conn: sqlite3.Connection, section_id: str) -> Optional[str]: - row = conn.execute( + rows = conn.execute( """ - SELECT c.external_id + SELECT DISTINCT c.external_id FROM node n JOIN cre_node_links l ON l.node = n.id JOIN cre c ON c.id = l.cre WHERE n.name LIKE '%ASVS%' AND n.section_id = ? ORDER BY c.external_id - LIMIT 1 """, (section_id,), - ).fetchone() - return row[0] if row else None + ).fetchall() + if not rows: + return None + if len(rows) > 1: + # Fail loud rather than silently pick one — an ambiguous section would + # bake a wrong CRE into the golden set (matches build_explicit/build_update). + raise ValueError( + f"ASVS section {section_id!r} maps to multiple CREs " + f"({', '.join(r[0] for r in rows)}); cannot pick a golden CRE" + ) + return rows[0][0] def build_positive_asvs(conn: sqlite3.Connection) -> List[Dict]: diff --git a/scripts/evaluate_librarian.py b/scripts/evaluate_librarian.py index 7c2ff3b03..89afa773d 100644 --- a/scripts/evaluate_librarian.py +++ b/scripts/evaluate_librarian.py @@ -95,22 +95,27 @@ def predict(section: Section, registry: Set[str], hub: List[HubRep]) -> List[str def report_retrieval_recall( - rows: List[GoldenDatasetRow], cache_file: str, top_k: int, threshold: float + rows: List[GoldenDatasetRow], + cache_file: str, + top_k: int, + threshold: float, + top_n_rerank: int, + crossencoder_model: str, ) -> None: - """Measure C.1 retrieval recall@k over the positive slice, live. + """Measure the live C.1 -> C.2 pipeline over the positive slice (v1). - Recall@k is the W3 metric — not the Jaccard link rule (that grades the - final auto-link, which needs the W4 reranker). It asks the only question - the search step controls: does the expected CRE id make it into the top-K - shortlist the reranker will later see? A miss here is unrecoverable - downstream, so it is the right thing to gate retrieval on. + Two metrics, both live — there is no honest offline value: the candidate + pool must be the real CRE-node vectors, and seeding it from the golden text + is exactly the leakage the hub firewall strips. - Live-only: there is no honest way to compute this offline. The candidate - pool must be the real CRE-node vectors, and seeding it from the golden - text itself is exactly the leakage the hub firewall strips. + - retrieval recall@k (C.1): does the expected CRE id make it into the top-K + shortlist the reranker will see? A miss here is unrecoverable downstream. + - rerank top-1 (C.2): after the cross-encoder re-reads each pair and re-sorts + the shortlist, is the #1 candidate an expected CRE? This is the first + end-to-end accuracy number for the search path (W4 target >= 0.80). """ - # Live deps are imported lazily so the offline harness needs neither a DB - # nor an embedding model. + # Live deps are imported lazily so the offline harness needs neither a DB, + # an embedding model, nor the cross-encoder stack. from application.cmd.cre_main import db_connect from application.defs import cre_defs from application.prompt_client import prompt_client @@ -119,15 +124,19 @@ def report_retrieval_recall( CandidateRetriever, ) from sqlalchemy import text as _sql_text + from application.utils.librarian.cross_encoder import ( + CrossEncoderReranker, + build_cross_encoder_score_fn, + ) database = db_connect(path=cache_file) ph = prompt_client.PromptHandler(database=database) # The embeddings pool is keyed by the CRE's internal UUID (add_embedding - # stores cre_id=db_object.id), but the golden dataset speaks external_ids - # ("616-305"). Translate the pool keys to external_id so recall compares in - # the same id-space. Keys not in the map (a DB that already stores - # external_ids) pass through unchanged. + # stores cre_id=db_object.id), but the golden dataset and the explicit + # resolver both speak external_ids ("616-305"). Translate the pool/reranker + # keys to external_id so recall/top-1 compare in the same id-space. Keys not + # in the map (a DB that already stores external_ids) pass through unchanged. id_to_ext = { row[0]: row[1] for row in database.session.execute( @@ -135,13 +144,12 @@ def report_retrieval_recall( ) if row[1] } + + def _to_ext(mapping): + return {id_to_ext.get(k, k): v for k, v in mapping.items()} + pool = CandidatePool.from_mapping( - { - id_to_ext.get(k, k): v - for k, v in database.get_embeddings_by_doc_type( - cre_defs.Credoctypes.CRE.value - ).items() - } + _to_ext(database.get_embeddings_by_doc_type(cre_defs.Credoctypes.CRE.value)) ) retriever = CandidateRetriever( embed_fn=ph.get_text_embeddings, @@ -149,25 +157,40 @@ def report_retrieval_recall( top_k=top_k, threshold=threshold, ) + reranker = CrossEncoderReranker( + score_fn=build_cross_encoder_score_fn(crossencoder_model), + top_n=top_n_rerank, + cre_texts=_to_ext( + database.get_embedding_contents_by_doc_type(cre_defs.Credoctypes.CRE.value) + ), + ) positives = [r for r in rows if r.slice.value == "positive" and r.expected.cre_ids] if not positives: print("retrieval recall: no positive rows with expected ids in this selection") return - any_hit = all_hit = 0 + any_hit = all_hit = top1_hit = 0 for row in positives: - retrieved = {c.cre_id for c in retriever.retrieve(row.input.text).candidates} + audit = retriever.retrieve(row.input.text) + audit = reranker.rerank(row.input.text, audit) + retrieved = {c.cre_id for c in audit.candidates} expected = set(row.expected.cre_ids or []) if expected & retrieved: any_hit += 1 if expected <= retrieved: all_hit += 1 + if audit.reranked and audit.reranked[0].cre_id in expected: + top1_hit += 1 n = len(positives) print( f"retrieval recall@{top_k} (live, {n} positive rows): " f"any-hit {any_hit}/{n} ({any_hit / n:.0%}), " f"all-hit {all_hit}/{n} ({all_hit / n:.0%})" ) + print( + f"rerank top-1 (C.2, top_n={top_n_rerank}): " + f"{top1_hit}/{n} ({top1_hit / n:.0%})" + ) def main(argv: List[str]) -> int: @@ -179,9 +202,6 @@ def main(argv: List[str]) -> int: parser.add_argument("--threshold", type=float, default=cfg.link_threshold) parser.add_argument("--top_k_retrieval", type=int, default=cfg.top_k_retrieval) parser.add_argument("--top_k_rerank", type=int, default=cfg.top_k_rerank) - parser.add_argument( - "--dry_run", action="store_true", help="no writes (always true pre-W8)" - ) parser.add_argument( "--no_hub_firewall", action="store_true", @@ -190,8 +210,9 @@ def main(argv: List[str]) -> int: parser.add_argument( "--use_live_embeddings", action="store_true", - help="connect to the OpenCRE DB + embedding model and measure semantic " - "retrieval recall@k over the positive slice (needs an LLM + populated DB)", + help="connect to the OpenCRE DB + embedding model and measure the live " + "C.1 retrieval recall@k and C.2 rerank top-1 over the positive slice " + "(needs an LLM + populated DB)", ) parser.add_argument( "--cache_file", @@ -256,12 +277,18 @@ def main(argv: List[str]) -> int: return 1 if args.use_live_embeddings: report_retrieval_recall( - rows, args.cache_file, args.top_k_retrieval, args.threshold + rows, + args.cache_file, + args.top_k_retrieval, + args.threshold, + args.top_k_rerank, + cfg.crossencoder_model, ) else: print( - "semantic retriever (C.1): wired; recall@k needs --use_live_embeddings " - "(no CRE vectors offline — seeding from golden text would be leakage)" + "semantic pipeline (C.1 retrieve + C.2 rerank): wired; recall@k and " + "rerank top-1 need --use_live_embeddings (no CRE vectors offline — " + "seeding from golden text would be leakage)" ) print(f"correct overall (semantic path still stubbed): {correct}/{len(rows)}") return 0