week_5: Module C (The Librarian) — C.3 confidence calibration (temperature scaling) - #974
week_5: Module C (The Librarian) — C.3 confidence calibration (temperature scaling)#974PRAteek-singHWY wants to merge 10 commits into
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
Summary by CodeRabbit
WalkthroughAdds temperature-scaling calibration utilities, deterministic tests, module documentation, dataset ID validation, and a live evaluation gate based on calibrated ECE. ChangesTemperature Calibration Feature
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
c03df58 to
fc68304
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (1)
scripts/evaluate_librarian.py (1)
182-184: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winLive pipeline is built (and positives reranked) twice per run.
With
--use_live_embeddings,report_retrieval_recallandreport_calibrationeach call_build_live_pipeline(...)independently, so the DB connection, embedding model, and cross-encoder are loaded twice, and every positive row is retrieved+reranked in both reports. Since the cross-encoder load and per-pair inference are the expensive steps here, building the pipeline once inmainand passing(retriever, reranker)into both reports would roughly halve the live cost. This also matches the intent stated in_build_live_pipeline's docstring that the heavy hub + model load happens once.♻️ Sketch: build once in main, pass into reports
if args.use_live_embeddings: + 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, - args.cache_file, - args.top_k_retrieval, - args.threshold, - args.top_k_rerank, - cfg.crossencoder_model, - ) - calib_status = report_calibration( - rows, - 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)(Adjust the two report signatures to accept the prebuilt
retriever, reranker.)Also applies to: 238-251
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@scripts/evaluate_librarian.py` around lines 182 - 184, The live pipeline is being constructed twice because both report_retrieval_recall and report_calibration call _build_live_pipeline independently. Build the pipeline once in main by calling _build_live_pipeline(cache_file, top_k, threshold, top_n_rerank, crossencoder_model) a single time, then pass the resulting retriever and reranker into both report functions. Update the signatures of report_retrieval_recall and report_calibration to accept the prebuilt pipeline objects and remove their internal _build_live_pipeline calls.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@scripts/evaluate_librarian.py`:
- Around line 182-184: The live pipeline is being constructed twice because both
report_retrieval_recall and report_calibration call _build_live_pipeline
independently. Build the pipeline once in main by calling
_build_live_pipeline(cache_file, top_k, threshold, top_n_rerank,
crossencoder_model) a single time, then pass the resulting retriever and
reranker into both report functions. Update the signatures of
report_retrieval_recall and report_calibration to accept the prebuilt pipeline
objects and remove their internal _build_live_pipeline calls.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
Run ID: e64443ca-ff6c-46c9-bff4-b4820fd64a97
📒 Files selected for processing (5)
application/tests/librarian/temperature_test.pyapplication/utils/librarian/__init__.pyapplication/utils/librarian/calibration/__init__.pyapplication/utils/librarian/calibration/temperature.pyscripts/evaluate_librarian.py
… once report_retrieval_recall and report_calibration each built the live pipeline (DB + embedding model + cross-encoder) independently, loading it twice and reranking every positive row twice per --use_live_embeddings run. Build it once in main and pass (retriever, reranker) into both reports, matching _build_live_pipeline's stated intent. Behavior-preserving: recall@20 285/292, rerank top-1 220/292, ECE 0.046 PASS unchanged.
northdpole
left a comment
There was a problem hiding this comment.
Maintainer review — Module C Week 5 (#974)
Temperature scaling (softmax-over-shortlist + NLL fit + ECE) looks correct and the hermetic tests are strong. One gate integrity issue before merge.
Please fix the inline comment (calibration skip must not exit 0 when --use_live_embeddings is on). After that, this is merge-ready as the base of the Module C stack (#974 → #990 → #991).
| "calibration (C.3): need both outcomes in the selection (positive + " | ||
| "hard_negative slices) to fit temperature; skipped" | ||
| ) | ||
| return 0 |
There was a problem hiding this comment.
High — calibration skip reports success
When the live C.3 calibration set is degenerate (all top-1 labels 0 or all 1, empty after dropping empty shortlists, or --slice with only one class), report_calibration prints "skipped" and returns 0. main then exits successfully even though the ECE < 0.10 gate never ran.
With --use_live_embeddings, a skipped gate should fail the run (non-zero exit), e.g. return 1, so CI/harness cannot greenwash a missing calibration. Only return 0 when calibration actually ran and ECE passed.
if len(set(labels)) < 2:
print(
"calibration (C.3): need both outcomes in the selection (positive + "
"hard_negative slices) to fit temperature; FAILED (gate did not run)"
)
return 1Please fix before merge.
There was a problem hiding this comment.
Fixed in 62abf2d2. A degenerate calibration set now returns 1 instead of 0, so --use_live_embeddings cannot exit 0 without the ECE gate having run. The message also prints the row and class counts, so a failure says why the set was degenerate rather than just that it was.
Verified the hermetic path is unaffected: evaluate_librarian.py --dataset application/tests/librarian/fixtures/golden_dataset.json still exits 0, since calibration is skipped upstream of this check when there are no live CRE vectors.
…on gate must fail
report_calibration returned 0 ("skipped") when the live calibration set was
degenerate (single-class labels, empty after dropping empty shortlists, or a
--slice with one class). Under --use_live_embeddings that let a run exit 0 without
the ECE < 0.10 gate ever running, so CI could greenwash a live run in which
calibration was never checked. A skipped gate now returns 1 (fail), with a message
stating the row/class counts, so exit 0 means the gate actually ran and passed.
…hortlists, dedupe softmax Three nitpicks from the bot review, none of which changed a metric: - evaluate_librarian: the live retrieve+rerank was recomputed per report. The models were already built once in main, but report_retrieval_recall and report_calibration each re-ran the pipeline over the positive slice, so every positive row paid for two cross-encoder passes. live_audits() now retrieves and reranks each row once, keyed by golden row id, and both reports read the same shortlists. Rows without an audit no longer count toward a report's denominator, so the printed fractions cannot divide by unscored rows. - temperature: _softmax_top duplicated the empty-shortlist guard and the softmax already in TemperatureScaler.probabilities. Both now route through one _softmax_at helper, and confidence() derives from probabilities() so the two can never disagree. - temperature is now clean under the --strict mypy the coding guidelines ask for: annotated _paired's return and the bounds tuple, re-asserted the array type over untyped scipy, and hoisted the label conversion out of the fit objective (it was re-validated on every optimiser iteration anyway). Behaviour is unchanged: 113 librarian tests pass and the hermetic harness run still exits 0.
…hortlists, dedupe softmax Three nitpicks from the bot review, none of which changed a metric: - evaluate_librarian: the live retrieve+rerank was recomputed per report. The models were already built once in main, but report_retrieval_recall and report_calibration each re-ran the pipeline over the positive slice, so every positive row paid for two cross-encoder passes. live_audits() now retrieves and reranks each row once, keyed by golden row id, and both reports read the same shortlists. Rows without an audit no longer count toward a report's denominator, so the printed fractions cannot divide by unscored rows. - temperature: _softmax_top duplicated the empty-shortlist guard and the softmax already in TemperatureScaler.probabilities. Both now route through one _softmax_at helper, and confidence() derives from probabilities() so the two can never disagree. - temperature is now clean under the --strict mypy the coding guidelines ask for: annotated _paired's return and the bounds tuple, re-asserted the array type over untyped scipy, and hoisted the label conversion out of the fit objective (it was re-validated on every optimiser iteration anyway). Behaviour is unchanged: 113 librarian tests pass and the hermetic harness run still exits 0. (cherry picked from commit 2bbc76e)
…on gate must fail
report_calibration returned 0 ("skipped") when the live calibration set was
degenerate (single-class labels, empty after dropping empty shortlists, or a
--slice with one class). Under --use_live_embeddings that let a run exit 0 without
the ECE < 0.10 gate ever running, so CI could greenwash a live run in which
calibration was never checked. A skipped gate now returns 1 (fail), with a message
stating the row/class counts, so exit 0 means the gate actually ran and passed.
(cherry picked from commit 62abf2d)
…hortlists, dedupe softmax Three nitpicks from the bot review, none of which changed a metric: - evaluate_librarian: the live retrieve+rerank was recomputed per report. The models were already built once in main, but report_retrieval_recall and report_calibration each re-ran the pipeline over the positive slice, so every positive row paid for two cross-encoder passes. live_audits() now retrieves and reranks each row once, keyed by golden row id, and both reports read the same shortlists. Rows without an audit no longer count toward a report's denominator, so the printed fractions cannot divide by unscored rows. - temperature: _softmax_top duplicated the empty-shortlist guard and the softmax already in TemperatureScaler.probabilities. Both now route through one _softmax_at helper, and confidence() derives from probabilities() so the two can never disagree. - temperature is now clean under the --strict mypy the coding guidelines ask for: annotated _paired's return and the bounds tuple, re-asserted the array type over untyped scipy, and hoisted the label conversion out of the fit objective (it was re-validated on every optimiser iteration anyway). Behaviour is unchanged: 113 librarian tests pass and the hermetic harness run still exits 0. (cherry picked from commit 2bbc76e) (cherry picked from commit ae6bb69)
…on gate must fail
report_calibration returned 0 ("skipped") when the live calibration set was
degenerate (single-class labels, empty after dropping empty shortlists, or a
--slice with one class). Under --use_live_embeddings that let a run exit 0 without
the ECE < 0.10 gate ever running, so CI could greenwash a live run in which
calibration was never checked. A skipped gate now returns 1 (fail), with a message
stating the row/class counts, so exit 0 means the gate actually ran and passed.
(cherry picked from commit 62abf2d)
(cherry picked from commit cb5577a)
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@scripts/evaluate_librarian.py`:
- Around line 176-179: Update load_dataset to detect and reject duplicate
GoldenDatasetRow.id values before live_audits, recall, or calibration compute
audits; preserve the existing non-empty ID validation and raise a clear
validation error identifying the duplicate ID.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro Plus
Run ID: af532b61-e0e4-4706-b0f2-db03cd46c26f
📒 Files selected for processing (2)
application/utils/librarian/calibration/temperature.pyscripts/evaluate_librarian.py
🚧 Files skipped from review as they are similar to previous changes (1)
- application/utils/librarian/calibration/temperature.py
|
Pushed
113 librarian tests pass and the hermetic harness run still exits 0. No metric moved. |
Follow-up to the shared-audit refactor. live_audits keys the reranked shortlists by row.id and the reports read them back with audits[row.id], so two rows sharing an id would collapse in that dict: the surviving audit gets reused for the earlier row and its scores are reported against the wrong text. Before the refactor each report recomputed per row, so a duplicate id was harmless; keying by it turned a harmless quirk into a silently wrong number. GoldenDatasetRow only requires an id to be non-empty, so uniqueness is enforced at load time and the harness refuses the file instead of printing a wrong metric. The committed dataset has 319 distinct ids, so nothing changes today; this closes the trap the refactor opened. Adds a test that a forced collision raises and names the offending id.
Hi @northdpole — Week 5 of Module C. This one turns the cross-encoder's raw scores into an honest probability (temperature scaling), so Week 6 can threshold "auto-link vs. send-to-human" on a number that actually means what it says.
Overview
Week 3 built the search step (C.1) — a bi-encoder that cosine-ranks the whole CRE hub down to a top-20 shortlist. Week 4 built the rerank step (C.2) — a cross-encoder that reads each
(section, candidate-CRE)pair together and re-sorts the shortlist by a raw relevance logit. Those logits are great for ordering but are not confidence: a raw+1.5does not mean "82% sure this is the right CRE," and cross-encoders are systematically over-confident.The problem: Week 6's decision — auto-link vs. send to a human — is a threshold on confidence ("link if ≥ 90%"). That is only safe if 90% really means 90%. A raw score cannot be thresholded honestly.
This PR's role: build the calibration step (C.3) — temperature scaling. Learn a single scalar temperature
Ton the golden set and map the reranked shortlist to a real probabilityp = softmax(logits / T); the top-1 candidate's share of that probability mass is the confidence — "how likely is the top candidate the correct CRE?" Prove it honest with Expected Calibration Error < 0.10. Temperature only flattens/sharpens the distribution — it never changes the ranking — so calibration can't hurt C.1/C.2 recall; it only makes the confidence trustworthy.One design note worth surfacing (the golden set caught it): calibrating the single absolute top-1 logit with
sigmoid(z / T)does not work — a cross-encoder's absolute logit has no fixed zero point (its 50/50 boundary isn't atz=0), and one temperature can only squash toward 0.5, never shift the boundary, so ECE stays stuck at ~0.27. Calibrating the softmax over the whole shortlist is the fix (and the textbook Guo et al. temperature scaling): the candidates' relative logits are what a cross-encoder's scores actually mean, and their softmax is a genuine distribution over "which candidate is right." It is near-calibrated even before fitting (ECE 0.053 atT=1), andTpolishes it to 0.046. (CALIBRATOR_NAME = "temperature-scaling/0.2.0"records this over the rejected single-logit0.1.0.)Scope: 1 new package (1 module + 1 test) + harness wiring for the ECE gate.
UpdateDetection(the other W5 schema deliverable) already exists inschemas.py— verified, nothing to add. No frontend, no behaviour change to OpenCRE proper.What changed
calibration/temperature.py(new),calibration/__init__.py(new)TemperatureScaler(T)maps a reranked shortlist to a calibrated distribution (.probabilities=softmax(logits/T)) and a top-1.confidence(the mass Week 6 thresholds on);fit_temperaturelearns the single scalarTby minimising NLL (scipy.optimize.minimize_scalar, bounded), guarding single-outcome data (DegenerateLabelsError);expected_calibration_erroris the 10-bin ECE gate;negative_log_likelihoodis the fit objective (exposed for testing/reporting). Model-free (numpy +scipy.special.softmax) so it stays import-light and hermetically testable — mirrors the C.1embed_fn/ C.2score_fnseams.CALIBRATOR_NAMEaudit tag (mirrorsRETRIEVER_NAME/RERANKER_NAME).evaluate_librarian.pyreport_calibrationbuilds a(shortlist, label)set from the live C.1→C.2 pipeline over the positive + hard_negative slices (each row's reranked logits; label1iff its top-1 is an expected CRE — hard_negatives supply the0outcome), fitsT, and reports ECE at T=1 vs. the fitted T with the < 0.10 gate._build_live_pipelineextracted so recall and calibration share one hub + model load. Offline path unchanged (no live number is faked).__init__.pytemperature_test.py(new)Tand reduces NLL vs. T=1, ECE on perfectly-/mis-calibrated data + a hand-checked two-bin value, peaked-vs-flat shortlist confidence, and every guard (single-outcome, non-binary, length mismatch, empty,n_bins < 1, non-positive/non-finiteT).How the pieces connect
flowchart TB audit["RetrievalAudit from C.2 (W4)<br/>reranked[] shortlist with score_rerank logits"] subgraph FIT["fit once on the golden set"] g["(shortlist, is-top1-correct) pairs<br/>positive + hard_negative slices"] nll["fit_temperature: argmin NLL(T)"] Tstar["T-star (one scalar)"] ece["expected_calibration_error<br/>gate: ECE below 0.10"] g --> nll --> Tstar --> ece end subgraph APPLY["apply per shortlist"] s["reranked logits"] p["confidence = top-1 mass of<br/>softmax(logits / T-star)"] s --> p end audit --> g audit --> s Tstar -.-> p p --> dec["W6 decision engine<br/>confidence at or above tau: LinkProposal, else ReviewItem"]Results
Reading it: the reranked top-1 is correct 75% of the time, and after calibration the model's stated confidence tracks that — mean confidence 0.767 vs. actual accuracy 0.757, ECE 0.046 < 0.10. So a "90%" from C.3 can be trusted as ~90%, which is exactly what the W6 auto-link threshold needs. (Calibration doesn't touch the 75% top-1 itself — that's the reranker's job, with W6/W7 the levers to lift it toward 0.80.)
What is intentionally not here
LinkProposal(≥ τ) vs. aReviewItem, and writing it intoProposedLink.confidence. This PR produces the honest number; W6 thresholds it.Tfor the live decision path (loaded by the W6 engine).--use_live_embeddingsrun.How to verify locally