diff --git a/application/tests/librarian/temperature_test.py b/application/tests/librarian/temperature_test.py new file mode 100644 index 000000000..4b4dee1cd --- /dev/null +++ b/application/tests/librarian/temperature_test.py @@ -0,0 +1,169 @@ +"""Tests for C.3 temperature-scaling calibration (Week 5). + +Hermetic and deterministic: synthetic candidate-logit shortlists + 0/1 labels +only — no cross-encoder, DB, or embedding key. Covers the softmax-over-shortlist +confidence, temperature flatten/sharpen behaviour, the NLL fit recovering a known +T and reducing NLL, ECE on perfectly- and mis-calibrated data (hand-checked), and +every guard. +""" + +import math +import unittest + +import numpy as np +from scipy.special import softmax + +from application.utils.librarian.calibration.temperature import ( + CALIBRATOR_NAME, + CalibrationError, + DegenerateLabelsError, + TemperatureScaler, + expected_calibration_error, + fit_temperature, + negative_log_likelihood, +) + + +class TemperatureScalerTest(unittest.TestCase): + def test_confidence_is_softmax_top_mass(self) -> None: + logits = [2.0, 1.0, 0.0] + expected = float(softmax(np.array(logits)).max()) + self.assertAlmostEqual(TemperatureScaler(1.0).confidence(logits), expected, 9) + + def test_probabilities_sum_to_one_and_peak_at_top(self) -> None: + p = TemperatureScaler(1.0).probabilities([3.0, 1.0, -2.0]) + self.assertAlmostEqual(float(p.sum()), 1.0, places=9) + self.assertEqual(int(np.argmax(p)), 0) # highest logit -> highest prob + + def test_high_temperature_flattens_toward_uniform(self) -> None: + logits = [3.0, 1.0, 0.0] + hot = TemperatureScaler(1000.0).confidence(logits) + base = TemperatureScaler(1.0).confidence(logits) + self.assertLess(hot, base) + self.assertAlmostEqual(hot, 1.0 / 3.0, places=2) # ~uniform over 3 + + def test_low_temperature_sharpens_toward_one(self) -> None: + logits = [3.0, 1.0, 0.0] + cold = TemperatureScaler(0.1).confidence(logits) + base = TemperatureScaler(1.0).confidence(logits) + self.assertGreater(cold, base) + self.assertGreater(cold, 0.99) + + def test_single_candidate_confidence_is_one(self) -> None: + # softmax over one element is always 1.0, at any temperature. + self.assertAlmostEqual(TemperatureScaler(3.0).confidence([0.42]), 1.0, 9) + + def test_empty_shortlist_rejected(self) -> None: + with self.assertRaises(CalibrationError): + TemperatureScaler(1.0).confidence([]) + + def test_non_positive_or_nonfinite_temperature_rejected(self) -> None: + for bad in (0.0, -1.0, float("inf"), float("nan")): + with self.assertRaises(CalibrationError): + TemperatureScaler(bad) + + +class NegativeLogLikelihoodTest(unittest.TestCase): + def test_matches_hand_computed_single_pair(self) -> None: + # one shortlist, top-1 correct: loss = -log(softmax([2,0]).max()) + conf = float(softmax(np.array([2.0, 0.0])).max()) + self.assertAlmostEqual( + negative_log_likelihood([[2.0, 0.0]], [1.0], 1.0), -math.log(conf), 6 + ) + + def test_length_mismatch_rejected(self) -> None: + with self.assertRaises(CalibrationError): + negative_log_likelihood([[1.0, 0.0], [2.0, 0.0]], [1.0], 1.0) + + def test_non_positive_temperature_rejected(self) -> None: + with self.assertRaises(CalibrationError): + negative_log_likelihood([[1.0, 0.0]], [1.0], 0.0) + + +class FitTemperatureTest(unittest.TestCase): + def _synthetic(self, true_t: float, n: int = 4000, k: int = 5): + """Shortlists whose top-1 correctness is drawn from softmax(logits/true_t). + + Data generated with true_t flattens the softmax, so the raw (T=1) + confidence is over-confident and the fit should recover T ~ true_t. + Seeded -> stable. + """ + rng = np.random.default_rng(0) + sets, labels = [], [] + for _ in range(n): + z = rng.normal(0.0, 3.0, size=k) + z[::-1].sort() # descending so index 0 is the top-1 (argmax) + p_correct = softmax(z / true_t).max() + sets.append(z.tolist()) + labels.append(1.0 if rng.random() < p_correct else 0.0) + return sets, labels + + def test_recovers_known_temperature(self) -> None: + sets, labels = self._synthetic(true_t=2.0) + scaler = fit_temperature(sets, labels) + self.assertAlmostEqual(scaler.temperature, 2.0, delta=0.5) + + def test_fit_reduces_nll_versus_T1(self) -> None: + sets, labels = self._synthetic(true_t=2.0) + scaler = fit_temperature(sets, labels) + self.assertLess( + negative_log_likelihood(sets, labels, scaler.temperature), + negative_log_likelihood(sets, labels, 1.0), + ) + + def test_single_class_labels_raise_degenerate(self) -> None: + with self.assertRaises(DegenerateLabelsError): + fit_temperature([[2.0, 0.0], [3.0, 1.0]], [1.0, 1.0]) + + def test_non_binary_labels_rejected(self) -> None: + with self.assertRaises(CalibrationError): + fit_temperature([[2.0, 0.0], [3.0, 1.0]], [0.0, 2.0]) + + +class ExpectedCalibrationErrorTest(unittest.TestCase): + def test_perfectly_calibrated_is_near_zero(self) -> None: + # bin [0.5,0.6): five correct, five wrong -> acc 0.5 == conf 0.5. + self.assertAlmostEqual( + expected_calibration_error([0.5] * 10, [1.0, 0.0] * 5), 0.0, places=6 + ) + + def test_confidently_wrong_is_large(self) -> None: + self.assertAlmostEqual( + expected_calibration_error([0.9] * 10, [0.0] * 10), 0.9, places=6 + ) + + def test_hand_checked_two_bin_value(self) -> None: + # [0.2,0.3): conf .2 acc 0 gap .2 ; [0.8,0.9): conf .8 acc 1 gap .2 + # ECE = 0.5*0.2 + 0.5*0.2 = 0.20 + self.assertAlmostEqual( + expected_calibration_error([0.2, 0.2, 0.8, 0.8], [0.0, 0.0, 1.0, 1.0]), + 0.20, + places=6, + ) + + def test_length_mismatch_and_empty_rejected(self) -> None: + with self.assertRaises(CalibrationError): + expected_calibration_error([0.5, 0.6], [1.0]) + with self.assertRaises(CalibrationError): + expected_calibration_error([], []) + with self.assertRaises(CalibrationError): + expected_calibration_error([0.5], [1.0], n_bins=0) + + +class EndToEndTest(unittest.TestCase): + def test_flat_shortlists_are_low_confidence(self) -> None: + # A near-tie shortlist (bad match) -> low top-1 confidence. + self.assertLess(TemperatureScaler(1.0).confidence([0.1, 0.0, -0.1]), 0.45) + + def test_peaked_shortlist_is_high_confidence(self) -> None: + # A clear winner -> high top-1 confidence. + self.assertGreater(TemperatureScaler(1.0).confidence([6.0, -1.0, -3.0]), 0.9) + + +class MetadataTest(unittest.TestCase): + def test_calibrator_name_is_versioned(self) -> None: + self.assertEqual(CALIBRATOR_NAME, "temperature-scaling/0.2.0") + + +if __name__ == "__main__": + unittest.main() diff --git a/application/utils/librarian/__init__.py b/application/utils/librarian/__init__.py index f168f2bb4..1a23cc3cb 100644 --- a/application/utils/librarian/__init__.py +++ b/application/utils/librarian/__init__.py @@ -16,7 +16,9 @@ explicit-link resolution). W3 (C.1): candidate retriever (in-memory + pgvector) + pipeline switch. 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. + W5 (C.3): confidence calibration — temperature scaling maps a rerank logit to + an honest probability (fit by NLL on the golden set, gated ECE < 0.10). +Decision routing (C.4, 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/calibration/__init__.py b/application/utils/librarian/calibration/__init__.py new file mode 100644 index 000000000..305586a76 --- /dev/null +++ b/application/utils/librarian/calibration/__init__.py @@ -0,0 +1,13 @@ +"""Module C.3 — confidence calibration (Week 5). + +C.2 (the cross-encoder, W4) emits a raw ranking logit per candidate — great for +ordering, meaningless as confidence (a +1.5 is not "82% sure"). C.3 turns that +logit into an honest probability via **temperature scaling**: ``p = sigmoid(z/T)`` +with a single scalar ``T`` fit by negative-log-likelihood on the golden set, and +proves the result honest with **ECE < 0.10**. + +The W6 decision engine thresholds that probability (auto-link vs. human review), +so calibration is what makes the threshold trustworthy. Kept dependency-light +(numpy + scipy) and model-free so it stays hermetically testable — mirrors the +C.1/C.2 seams. +""" diff --git a/application/utils/librarian/calibration/temperature.py b/application/utils/librarian/calibration/temperature.py new file mode 100644 index 000000000..928dacf7a --- /dev/null +++ b/application/utils/librarian/calibration/temperature.py @@ -0,0 +1,207 @@ +"""Module C.3 — temperature-scaling calibration (Week 5). The truth-teller. + +C.2 hands up a reranked shortlist of candidate CREs, each with a raw cross-encoder +logit (unbounded, higher = better match). We need one honest number: "how likely +is the top candidate the correct CRE?" — so Week 6 can threshold auto-link vs. +human review. + +The naive attempt — ``sigmoid(top1_logit / T)`` on the single absolute logit — +does not work, and the golden set proves why: a cross-encoder's absolute logit has +no fixed zero point (its 50/50 boundary is not at z=0), and dividing by a single +temperature can only *squash* toward 0.5, never *shift* the boundary. So no T +makes it honest. + +The fix is **temperature scaling in the Guo et al. sense**: calibrate the *softmax +over the whole shortlist*, not one absolute logit. The candidates' *relative* +logits are what a cross-encoder's scores actually mean, so + + p = softmax(logits / T) confidence = p for the top-1 candidate + +is a genuine probability distribution over "which candidate is right", and its +top-1 mass answers exactly the question Week 6 asks. It is still **one knob T**: +T = 1 leaves the distribution unchanged, T > 1 flattens it (less confident), T < 1 +sharpens it. T is fit once by minimising negative-log-likelihood of "is the top-1 +correct?" on the golden set; honesty is measured by Expected Calibration Error, +and the Week 5 gate is ECE < 0.10. + +Like C.1/C.2 this is a thin, model-free seam (numpy + scipy only) so every branch +is hermetically testable. The fitted ``TemperatureScaler`` is a frozen, shareable +artifact the W6 decision engine loads to turn a reranked shortlist into the +confidence it thresholds on. +""" + +from dataclasses import dataclass +from typing import Sequence + +import numpy as np +from scipy.optimize import minimize_scalar +from scipy.special import softmax + +# Identify the calibrator in the RFC audit trail (mirrors RETRIEVER_NAME / +# RERANKER_NAME). 0.2.0 = softmax-over-shortlist (the single-logit sigmoid of +# 0.1.0 could not be calibrated by temperature alone — see module docstring). +CALIBRATOR_NAME = "temperature-scaling/0.2.0" + +# Clip probabilities off {0, 1} so the log in NLL stays finite. +_EPS = 1e-7 + + +class CalibrationError(ValueError): + """Base class for calibration construction/usage failures.""" + + +class DegenerateLabelsError(CalibrationError): + """Labels are single-class — temperature is unidentifiable from NLL. + + With only correct (or only incorrect) top-1s, NLL is monotonic in ``T`` and + the optimum runs to a bound: the fit is meaningless. The calibration set must + contain both outcomes (in the harness: the ``positive`` slice supplies + correct top-1s, ``hard_negative`` supplies incorrect ones). + """ + + +def _softmax_top(logits: Sequence[float], temperature: float) -> float: + """Top-1 probability mass of ``softmax(logits / T)`` over one shortlist.""" + z = np.asarray(list(logits), dtype=float) + if z.size == 0: + raise CalibrationError("cannot calibrate an empty candidate shortlist") + return float(softmax(z / temperature).max()) + + +def _validate_temperature(temperature: float) -> None: + if not np.isfinite(temperature) or temperature <= 0: + raise CalibrationError(f"temperature must be finite and > 0, got {temperature}") + + +def _paired(logit_sets: Sequence[Sequence[float]], labels: Sequence[float]): + """Validate matched (shortlist, label) inputs: non-empty and equal length.""" + sets = list(logit_sets) + y = np.asarray(list(labels), dtype=float) + if y.ndim != 1: + raise CalibrationError(f"labels must be 1-D, got shape {y.shape}") + if len(sets) != y.shape[0]: + raise CalibrationError( + f"{len(sets)} logit-sets and {y.shape[0]} labels must be the same length" + ) + if not sets: + raise CalibrationError("need at least one (shortlist, label) pair") + return sets, y + + +@dataclass(frozen=True) +class TemperatureScaler: + """Maps a reranked shortlist to a calibrated top-1 confidence. + + ``temperature`` is the single learned scalar; build one with + ``fit_temperature``. Frozen so a fitted scaler is a stable, shareable value + (like the retriever/reranker being constructed once and reused). + """ + + temperature: float + + def __post_init__(self) -> None: + _validate_temperature(self.temperature) + + def probabilities(self, logits: Sequence[float]) -> np.ndarray: + """The full ``softmax(logits / T)`` distribution over one shortlist.""" + z = np.asarray(list(logits), dtype=float) + if z.size == 0: + raise CalibrationError("cannot calibrate an empty candidate shortlist") + return softmax(z / self.temperature) + + def confidence(self, logits: Sequence[float]) -> float: + """P(the top candidate is correct) — the top-1 mass of the softmax. + + This is the number the W6 decision engine thresholds on. + """ + return _softmax_top(logits, self.temperature) + + +def negative_log_likelihood( + logit_sets: Sequence[Sequence[float]], + labels: Sequence[float], + temperature: float, +) -> float: + """Binary cross-entropy of the top-1 confidence against 0/1 correctness. + + For each shortlist the confidence is ``softmax(logits / T)``'s top-1 mass; + the label is 1 iff that top-1 candidate is the correct CRE. This is the + objective ``fit_temperature`` minimises over ``T``; exposed so the fit is + testable and a caller can compare NLL at ``T=1`` vs the fitted T. + """ + _validate_temperature(temperature) + sets, y = _paired(logit_sets, labels) + p = np.clip( + np.array([_softmax_top(s, temperature) for s in sets]), _EPS, 1.0 - _EPS + ) + return float(-np.sum(y * np.log(p) + (1.0 - y) * np.log(1.0 - p))) + + +def fit_temperature( + logit_sets: Sequence[Sequence[float]], + labels: Sequence[float], + *, + bounds: tuple = (1e-2, 1e2), +) -> TemperatureScaler: + """Fit ``T`` by minimising NLL over (shortlist, is-top1-correct) pairs. + + ``labels`` must be 0/1 and contain both outcomes (else + ``DegenerateLabelsError``). One free parameter over a bounded interval, so + the 1-D ``minimize_scalar`` fit is fast and barely over-fits — the golden set + is the calibration set by design. + """ + sets, y = _paired(logit_sets, labels) + values = set(np.unique(y).tolist()) + if not values <= {0.0, 1.0}: + raise CalibrationError(f"labels must be 0/1, got values {sorted(values)}") + if len(values) < 2: + raise DegenerateLabelsError( + "labels are single-class; temperature is unidentifiable — the " + "calibration set needs both correct and incorrect top-1s" + ) + + result = minimize_scalar( + lambda t: negative_log_likelihood(sets, y, t), + bounds=bounds, + method="bounded", + ) + if not result.success: + raise CalibrationError(f"temperature fit did not converge: {result.message}") + return TemperatureScaler(temperature=float(result.x)) + + +def expected_calibration_error( + confidences: Sequence[float], labels: Sequence[float], *, n_bins: int = 10 +) -> float: + """ECE — sample-weighted mean ``|accuracy - confidence|`` across equal bins. + + Partition [0, 1] into ``n_bins`` equal-width bins. Per bin, ``confidence`` is + the mean predicted top-1 probability and ``accuracy`` is the fraction whose + top-1 was actually correct; ECE weights each bin's gap by its share of the + sample. 0 = perfectly honest; the Week 5 gate is < 0.10. + """ + if n_bins < 1: + raise CalibrationError(f"n_bins must be >= 1, got {n_bins}") + p = np.asarray(list(confidences), dtype=float) + y = np.asarray(list(labels), dtype=float) + if p.shape != y.shape: + raise CalibrationError( + f"confidences {p.shape} and labels {y.shape} must be the same length" + ) + if p.size == 0: + raise CalibrationError("need at least one (confidence, label) pair") + + edges = np.linspace(0.0, 1.0, n_bins + 1) + idx = np.clip(np.digitize(p, edges[1:-1], right=False), 0, n_bins - 1) + + n = p.size + ece = 0.0 + for b in range(n_bins): + mask = idx == b + count = int(mask.sum()) + if count == 0: + continue + confidence = float(p[mask].mean()) + accuracy = float(y[mask].mean()) + ece += (count / n) * abs(accuracy - confidence) + return ece diff --git a/scripts/evaluate_librarian.py b/scripts/evaluate_librarian.py index 89afa773d..201497b23 100644 --- a/scripts/evaluate_librarian.py +++ b/scripts/evaluate_librarian.py @@ -94,28 +94,20 @@ def predict(section: Section, registry: Set[str], hub: List[HubRep]) -> List[str return [] -def report_retrieval_recall( - rows: List[GoldenDatasetRow], +def _build_live_pipeline( cache_file: str, top_k: int, threshold: float, top_n_rerank: int, crossencoder_model: str, -) -> None: - """Measure the live C.1 -> C.2 pipeline over the positive slice (v1). - - 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. +): + """Construct the live C.1 retriever + C.2 reranker against the OpenCRE DB. - - 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, an + embedding model, nor the cross-encoder stack. Shared by every live report + (recall/top-1 and calibration) so the heavy hub + model load happens once + per report and the id-space translation stays in one place. """ - # 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 @@ -164,7 +156,30 @@ def _to_ext(mapping): database.get_embedding_contents_by_doc_type(cre_defs.Credoctypes.CRE.value) ), ) + return retriever, reranker + +def report_retrieval_recall( + rows: List[GoldenDatasetRow], + retriever, + reranker, + top_k: int, + top_n_rerank: int, +) -> None: + """Measure the live C.1 -> C.2 pipeline over the positive slice (v1). + + Takes a prebuilt ``retriever``/``reranker`` (built once in ``main``) so the + DB, embedding model, and cross-encoder load once per run and are shared with + ``report_calibration``. 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. + + - 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). + """ 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") @@ -193,6 +208,62 @@ def _to_ext(mapping): ) +def report_calibration( + rows: List[GoldenDatasetRow], + retriever, + reranker, +) -> int: + """Fit temperature on the golden set and report the Week 5 ECE gate (< 0.10). + + Takes the prebuilt ``retriever``/``reranker`` shared with + ``report_retrieval_recall`` (built once in ``main``). Builds a + (shortlist, label) calibration set from the live C.1 -> C.2 pipeline over the + positive + hard_negative slices: each row's *reranked shortlist* of logits, + labelled 1 iff its top-1 candidate is an expected CRE (hard_negatives expect + none, so they contribute the 0 class). Both slices are needed so the fit sees + both outcomes (else it is degenerate). Confidence is the top-1 mass of + softmax(logits / T); prints ECE at T=1 vs the fitted T and PASS/FAIL on + ECE < 0.10; returns 1 on a failed gate so a live run can fail. + """ + from application.utils.librarian.calibration.temperature import ( + TemperatureScaler, + expected_calibration_error, + fit_temperature, + ) + + cal_rows = [r for r in rows if r.slice.value in ("positive", "hard_negative")] + logit_sets: List[List[float]] = [] + labels: List[float] = [] + for row in cal_rows: + audit = reranker.rerank(row.input.text, retriever.retrieve(row.input.text)) + reranked = [c for c in audit.reranked if c.score_rerank is not None] + if not reranked: + continue + expected = set(row.expected.cre_ids or []) + logit_sets.append([float(c.score_rerank) for c in reranked]) + labels.append(1.0 if reranked[0].cre_id in expected else 0.0) + + if len(set(labels)) < 2: + print( + "calibration (C.3): need both outcomes in the selection (positive + " + "hard_negative slices) to fit temperature; skipped" + ) + return 0 + + scaler = fit_temperature(logit_sets, labels) + conf_raw = [TemperatureScaler(1.0).confidence(s) for s in logit_sets] + conf_cal = [scaler.confidence(s) for s in logit_sets] + ece_raw = expected_calibration_error(conf_raw, labels) + ece_cal = expected_calibration_error(conf_cal, labels) + gate_ok = ece_cal < 0.10 + print( + f"calibration (C.3, {len(labels)} rows): fitted T={scaler.temperature:.3f}; " + f"ECE {ece_raw:.3f} (raw, T=1) -> {ece_cal:.3f} (calibrated); " + f"gate ECE<0.10: {'PASS' if gate_ok else 'FAIL'}" + ) + return 0 if gate_ok else 1 + + def main(argv: List[str]) -> int: cfg = load_config() parser = argparse.ArgumentParser(description="Module C eval harness (W2: C.0)") @@ -275,23 +346,31 @@ def main(argv: List[str]) -> int: ) if not gate_ok: return 1 + calib_status = 0 if args.use_live_embeddings: - report_retrieval_recall( - rows, + # Build the live pipeline once (DB + embedding model + cross-encoder) and + # share it across both live reports, so the heavy load and per-row rerank + # happen a single time per run. + retriever, reranker = _build_live_pipeline( args.cache_file, args.top_k_retrieval, args.threshold, args.top_k_rerank, cfg.crossencoder_model, ) + report_retrieval_recall( + rows, retriever, reranker, args.top_k_retrieval, args.top_k_rerank + ) + calib_status = report_calibration(rows, retriever, reranker) else: print( - "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)" + "semantic pipeline (C.1 retrieve + C.2 rerank) + calibration (C.3): " + "wired; recall@k, rerank top-1, and the ECE gate 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 + return calib_status if __name__ == "__main__":