1use std::collections::hash_map::DefaultHasher;
19use std::hash::{Hash, Hasher};
20use std::path::Path;
21
22use mif_ontology::{CalibrationConfig, OntologyError};
23use serde::Serialize;
24
25use crate::error::MifRhError;
26use crate::resolve::{Basis, MapRecord, ResolveContext};
27use crate::suggest::{SUGGESTION_DEPTH, build_candidates, suggest_from_candidates};
28use crate::{Finding, index_text, review::list_finding_files};
29
30#[non_exhaustive]
37#[derive(Debug, Clone, Serialize)]
38pub struct CalibrationSample {
39 pub finding_id: String,
41 pub topic: String,
43 pub entity_type_gold: String,
45 pub entity_type_top1: String,
47 pub top1_score: f32,
49 pub top1_margin: Option<f32>,
53 pub top1_correct: bool,
55 pub gold_in_candidates: bool,
57}
58
59#[derive(Debug, Clone, Copy)]
61pub struct CalibrateOptions {
62 pub target_precision: f32,
64 pub tier2_target: f32,
66 pub sample: Option<usize>,
69 pub seed: u64,
71}
72
73impl Default for CalibrateOptions {
74 fn default() -> Self {
75 Self {
76 target_precision: 0.95,
77 tier2_target: 0.5,
78 sample: None,
79 seed: 0,
80 }
81 }
82}
83
84pub fn collect_topic_samples(
96 reports_dir: &Path,
97 ctx: &ResolveContext<'_>,
98 embedder: &mif_embed::Embedder,
99) -> Result<Vec<CalibrationSample>, MifRhError> {
100 let map_path = reports_dir.join(ctx.topic).join("ontology-map.json");
101 let contents = match std::fs::read_to_string(&map_path) {
102 Ok(contents) => contents,
103 Err(source) if source.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
105 Err(source) => {
108 return Err(MifRhError::Io {
109 path: map_path.display().to_string(),
110 source,
111 });
112 },
113 };
114 let records: Vec<MapRecord> =
115 serde_json::from_str(&contents).map_err(|source| MifRhError::Json {
116 path: map_path.display().to_string(),
117 source,
118 })?;
119
120 let findings_dir = reports_dir.join(ctx.topic).join("findings");
121 if !findings_dir.is_dir() {
122 return Ok(Vec::new());
123 }
124
125 let neutral = CalibrationConfig::default();
126 let candidates = build_candidates(ctx, embedder, &neutral)?;
130 let mut samples = Vec::new();
131 for file in list_finding_files(&findings_dir)? {
132 let Ok(finding) = Finding::load(&file) else {
133 continue; };
135 let Some(record) = records.iter().find(|r| r.finding_id == finding.id) else {
136 continue;
137 };
138 let stamped = record.valid && matches!(record.basis, Basis::Declared | Basis::Resolved);
139 let Some(gold) = record.entity_type.as_deref() else {
140 continue;
141 };
142 if !stamped {
143 continue;
144 }
145
146 let query = index_text(&finding);
147 if query.is_empty() {
148 continue;
149 }
150 let query_vector = embedder.embed(&query)?;
151 let ranked =
152 suggest_from_candidates(&query_vector, &candidates, &neutral, SUGGESTION_DEPTH);
153 let Some(top) = ranked.first() else {
154 continue; };
156 samples.push(CalibrationSample {
157 finding_id: finding.id.clone(),
158 topic: ctx.topic.to_string(),
159 entity_type_gold: gold.to_string(),
160 entity_type_top1: top.entity_type.clone(),
161 top1_score: top.score,
162 top1_margin: top.margin,
163 top1_correct: top.entity_type == gold,
164 gold_in_candidates: ranked.iter().any(|c| c.entity_type == gold),
165 });
166 }
167 Ok(samples)
168}
169
170#[must_use]
174pub fn subsample(
175 mut samples: Vec<CalibrationSample>,
176 opts: &CalibrateOptions,
177) -> Vec<CalibrationSample> {
178 let Some(cap) = opts.sample else {
179 return samples;
180 };
181 if samples.len() <= cap {
182 return samples;
183 }
184 samples.sort_by_key(|s| {
185 let mut hasher = DefaultHasher::new();
186 opts.seed.hash(&mut hasher);
187 s.finding_id.hash(&mut hasher);
188 hasher.finish()
189 });
190 samples.truncate(cap);
191 samples
192}
193
194pub fn sweep(
210 samples: &[CalibrationSample],
211 opts: &CalibrateOptions,
212 artifact_path: &Path,
213) -> Result<CalibrationConfig, MifRhError> {
214 let invalid = |detail: String| {
215 MifRhError::from(OntologyError::CalibrationInvalid {
216 path: artifact_path.display().to_string(),
217 detail,
218 })
219 };
220 if samples.is_empty() {
221 return Err(invalid(
222 "no stamped, valid findings with scorable entity types to calibrate from — \
223 review and stamp findings first"
224 .to_string(),
225 ));
226 }
227
228 let mut best: Option<(usize, u8, u8)> = None; for floor_pct in 0..=95_u8 {
234 for margin_pct in 0..=20_u8 {
235 let Some(accepted) =
236 accepted_meeting_target(samples, floor_pct, margin_pct, opts.target_precision)
237 else {
238 continue;
239 };
240 let candidate = (accepted, floor_pct, margin_pct);
241 if best.is_none_or(|current| gate_is_better(candidate, current)) {
242 best = Some(candidate);
243 }
244 }
245 }
246 let Some((_, tier1_floor_pct, tier1_margin_pct)) = best else {
247 return Err(invalid(format!(
248 "no (floor, margin) grid point reaches top-1 precision {} over {} samples — \
249 enrich entity types (aliases/exemplars) or lower --target-precision",
250 opts.target_precision,
251 samples.len()
252 )));
253 };
254
255 let mut tier2_floor_pct = tier1_floor_pct;
266 for floor_pct in 0..=tier1_floor_pct {
267 let floor = f32::from(floor_pct) / 100.0;
268 let (mut total, mut with_gold) = (0_usize, 0_usize);
269 for s in samples.iter().filter(|s| s.top1_score >= floor) {
270 total += 1;
271 with_gold += usize::from(s.gold_in_candidates);
272 }
273 if total > 0 && ratio(with_gold, total) >= opts.tier2_target {
274 tier2_floor_pct = floor_pct;
275 break; }
277 }
278
279 Ok(CalibrationConfig {
280 tier1_floor: f32::from(tier1_floor_pct) / 100.0,
281 tier1_margin: f32::from(tier1_margin_pct) / 100.0,
282 tier2_floor: f32::from(tier2_floor_pct) / 100.0,
283 calibrated: true,
284 calibrated_at: Some(now_rfc3339()),
285 sample_size: Some(u64::try_from(samples.len()).unwrap_or(u64::MAX)),
286 method: Some("stamped-quantile-v1".to_string()),
287 ..CalibrationConfig::default()
288 })
289}
290
291pub const CONFUSION_REPRESENTATIVES: usize = 5;
294
295#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
298pub struct ConfusionPair {
299 pub gold: String,
301 pub top1: String,
303 pub count: usize,
305 pub finding_ids: Vec<String>,
309}
310
311#[must_use]
320pub fn packs_carry_negatives<'a>(packs: impl IntoIterator<Item = &'a crate::OntologyPack>) -> bool {
321 packs
322 .into_iter()
323 .flat_map(|pack| &pack.entity_types)
324 .any(|entity_type| {
325 entity_type.embedding_doc().is_some()
329 && entity_type
330 .negative_examples
331 .iter()
332 .any(|negative| !negative.trim().is_empty())
333 })
334}
335
336#[derive(Debug, Clone, Serialize)]
355pub struct ConfusionReport {
356 pub version: &'static str,
358 pub sample_count: usize,
360 pub pairs: Vec<ConfusionPair>,
362}
363
364#[must_use]
371pub fn confusions(samples: &[CalibrationSample]) -> ConfusionReport {
372 let mut by_pair: std::collections::BTreeMap<(&str, &str), Vec<&str>> =
373 std::collections::BTreeMap::new();
374 for s in samples.iter().filter(|s| !s.top1_correct) {
375 by_pair
376 .entry((s.entity_type_gold.as_str(), s.entity_type_top1.as_str()))
377 .or_default()
378 .push(s.finding_id.as_str());
379 }
380 let mut pairs: Vec<ConfusionPair> = by_pair
381 .into_iter()
382 .map(|((gold, top1), mut ids)| {
383 ids.sort_unstable();
384 let count = ids.len();
385 ids.truncate(CONFUSION_REPRESENTATIVES);
386 ConfusionPair {
387 gold: gold.to_string(),
388 top1: top1.to_string(),
389 count,
390 finding_ids: ids.into_iter().map(str::to_string).collect(),
391 }
392 })
393 .collect();
394 pairs.sort_by_key(|p| std::cmp::Reverse(p.count));
397 ConfusionReport {
398 version: "confusions-v1",
399 sample_count: samples.len(),
400 pairs,
401 }
402}
403
404fn ratio(num: usize, den: usize) -> f32 {
407 #[allow(clippy::cast_precision_loss)]
408 {
409 num as f32 / den as f32
410 }
411}
412
413fn accepted_meeting_target(
417 samples: &[CalibrationSample],
418 floor_pct: u8,
419 margin_pct: u8,
420 target: f32,
421) -> Option<usize> {
422 let floor = f32::from(floor_pct) / 100.0;
423 let margin = f32::from(margin_pct) / 100.0;
424 let (mut accepted, mut correct) = (0_usize, 0_usize);
425 for s in samples {
426 let margin_ok = s.top1_margin.is_none_or(|m| m >= margin);
430 if s.top1_score >= floor && margin_ok {
431 accepted += 1;
432 correct += usize::from(s.top1_correct);
433 }
434 }
435 (accepted > 0 && ratio(correct, accepted) >= target).then_some(accepted)
436}
437
438const fn gate_is_better(candidate: (usize, u8, u8), current: (usize, u8, u8)) -> bool {
441 candidate.0 > current.0
442 || (candidate.0 == current.0
443 && (candidate.1 < current.1 || (candidate.1 == current.1 && candidate.2 < current.2)))
444}
445
446fn now_rfc3339() -> String {
448 chrono::Utc::now().format("%Y-%m-%dT%H:%M:%SZ").to_string()
449}
450
451#[cfg(test)]
452mod tests {
453 use std::path::Path;
454
455 use super::{
456 CONFUSION_REPRESENTATIVES, CalibrateOptions, CalibrationSample, confusions, subsample,
457 sweep,
458 };
459
460 fn sample(
461 id: &str,
462 score: f32,
463 margin: f32,
464 correct: bool,
465 gold_in: bool,
466 ) -> CalibrationSample {
467 typed_sample(
468 id,
469 "gold-type",
470 "other-type",
471 score,
472 margin,
473 correct,
474 gold_in,
475 )
476 }
477
478 #[allow(clippy::too_many_arguments)]
479 fn typed_sample(
480 id: &str,
481 gold: &str,
482 top1: &str,
483 score: f32,
484 margin: f32,
485 correct: bool,
486 gold_in: bool,
487 ) -> CalibrationSample {
488 CalibrationSample {
489 finding_id: id.to_string(),
490 topic: "topic".to_string(),
491 entity_type_gold: gold.to_string(),
492 entity_type_top1: if correct {
493 gold.to_string()
494 } else {
495 top1.to_string()
496 },
497 top1_score: score,
498 top1_margin: Some(margin),
499 top1_correct: correct,
500 gold_in_candidates: gold_in,
501 }
502 }
503
504 #[test]
505 fn sweep_prefers_the_loosest_gate_meeting_the_precision_target() {
506 let samples = [
509 sample("a", 0.90, 0.10, true, true),
510 sample("b", 0.88, 0.09, true, true),
511 sample("c", 0.60, 0.02, false, true),
512 ];
513 let cal = sweep(
514 &samples,
515 &CalibrateOptions {
516 target_precision: 1.0,
517 ..CalibrateOptions::default()
518 },
519 Path::new("test.json"),
520 )
521 .unwrap();
522
523 assert!(cal.calibrated);
524 assert!(
527 cal.tier1_floor > 0.60 || cal.tier1_margin > 0.02,
528 "gate must exclude the wrong sample (floor {}, margin {})",
529 cal.tier1_floor,
530 cal.tier1_margin
531 );
532 assert!(cal.tier1_floor <= 0.88 && cal.tier1_margin <= 0.09);
534 assert!(cal.tier2_floor <= cal.tier1_floor);
535 assert_eq!(cal.method.as_deref(), Some("stamped-quantile-v1"));
536 assert_eq!(cal.sample_size, Some(3));
537 }
538
539 #[test]
540 fn no_rival_samples_pass_margin_gates_vacuously_matching_runtime() {
541 let lone = CalibrationSample {
546 finding_id: "lone".to_string(),
547 topic: "topic".to_string(),
548 entity_type_gold: "gold-type".to_string(),
549 entity_type_top1: "gold-type".to_string(),
550 top1_score: 0.90,
551 top1_margin: None,
552 top1_correct: true,
553 gold_in_candidates: true,
554 };
555 let rivaled = sample("rivaled", 0.88, 0.10, true, true);
556 let cal = sweep(
557 &[lone, rivaled],
558 &CalibrateOptions {
559 target_precision: 1.0,
560 ..CalibrateOptions::default()
561 },
562 Path::new("test.json"),
563 )
564 .unwrap();
565 assert!(cal.calibrated);
566 assert_eq!(cal.sample_size, Some(2));
567 }
568
569 #[test]
570 fn sweep_fails_loud_on_an_empty_sample_set() {
571 let error = sweep(&[], &CalibrateOptions::default(), Path::new("test.json")).unwrap_err();
572 assert!(error.to_string().contains("no stamped"));
573 }
574
575 #[test]
576 fn sweep_fails_loud_when_no_grid_point_reaches_the_target() {
577 let samples = [
579 sample("a", 0.90, 0.10, false, true),
580 sample("b", 0.88, 0.09, false, false),
581 ];
582 let error = sweep(
583 &samples,
584 &CalibrateOptions::default(),
585 Path::new("test.json"),
586 )
587 .unwrap_err();
588 assert!(error.to_string().contains("precision"));
589 }
590
591 #[test]
592 fn confusions_rank_known_pairs_by_count_then_lexicographically() {
593 let samples = [
597 typed_sample("f-3", "curriculum", "title", 0.7, 0.05, false, true),
598 typed_sample("f-1", "curriculum", "title", 0.6, 0.04, false, true),
599 typed_sample("f-2", "curriculum", "title", 0.8, 0.06, false, true),
600 typed_sample("f-4", "title", "curriculum", 0.5, 0.02, false, true),
601 typed_sample("f-5", "title", "title", 0.9, 0.20, true, true),
602 ];
603 let report = confusions(&samples);
604
605 assert_eq!(report.version, "confusions-v1");
606 assert_eq!(report.sample_count, 5);
607 assert_eq!(report.pairs.len(), 2);
608 assert_eq!(report.pairs[0].gold, "curriculum");
609 assert_eq!(report.pairs[0].top1, "title");
610 assert_eq!(report.pairs[0].count, 3);
611 assert_eq!(report.pairs[0].finding_ids, ["f-1", "f-2", "f-3"]);
614 assert_eq!(report.pairs[1].count, 1);
615 assert_eq!(report.pairs[1].gold, "title");
616 }
617
618 #[test]
619 fn confusions_are_deterministic_and_cap_representatives() {
620 let build = |order: &[usize]| {
621 order
622 .iter()
623 .map(|i| typed_sample(&format!("f-{i}"), "a", "b", 0.5, 0.0, false, true))
624 .collect::<Vec<_>>()
625 };
626 let forward = confusions(&build(&[1, 2, 3, 4, 5, 6, 7, 8]));
628 let reverse = confusions(&build(&[8, 7, 6, 5, 4, 3, 2, 1]));
629 assert_eq!(forward.pairs, reverse.pairs);
630
631 assert_eq!(forward.pairs[0].count, 8);
632 assert_eq!(
633 forward.pairs[0].finding_ids.len(),
634 CONFUSION_REPRESENTATIVES
635 );
636 assert_eq!(forward.pairs[0].finding_ids[0], "f-1");
637 }
638
639 #[test]
640 fn confusions_on_an_all_correct_corpus_are_empty_but_versioned() {
641 let samples = [sample("a", 0.9, 0.1, true, true)];
642 let report = confusions(&samples);
643 assert_eq!(report.version, "confusions-v1");
644 assert_eq!(report.sample_count, 1);
645 assert!(report.pairs.is_empty());
646 }
647
648 #[test]
649 fn negatives_on_a_signal_less_type_never_claim_participation() {
650 let signal_less = crate::ontology_pack::parse_pack(
654 "ontology:\n id: p\n version: \"0.1.0\"\nentity_types:\n - name: ghost\n negative_examples: [a near miss]\n",
655 "p.yaml",
656 )
657 .unwrap();
658 assert!(!super::packs_carry_negatives([&signal_less]));
659
660 let scorable = crate::ontology_pack::parse_pack(
661 "ontology:\n id: q\n version: \"0.1.0\"\nentity_types:\n - name: real\n description: A described type\n negative_examples: [a near miss]\n",
662 "q.yaml",
663 )
664 .unwrap();
665 assert!(super::packs_carry_negatives([&scorable]));
666 }
667
668 #[test]
669 fn subsample_is_deterministic_and_seed_sensitive() {
670 let build = || {
671 (0..20)
672 .map(|i| sample(&format!("f-{i}"), 0.5, 0.0, true, true))
673 .collect::<Vec<_>>()
674 };
675 let opts_a = CalibrateOptions {
676 sample: Some(5),
677 seed: 1,
678 ..CalibrateOptions::default()
679 };
680 let first = subsample(build(), &opts_a);
681 let second = subsample(build(), &opts_a);
682 assert_eq!(
683 first.iter().map(|s| &s.finding_id).collect::<Vec<_>>(),
684 second.iter().map(|s| &s.finding_id).collect::<Vec<_>>()
685 );
686 assert_eq!(first.len(), 5);
687
688 let opts_b = CalibrateOptions {
689 sample: Some(5),
690 seed: 2,
691 ..CalibrateOptions::default()
692 };
693 let other_seed = subsample(build(), &opts_b);
694 assert_ne!(
695 first.iter().map(|s| &s.finding_id).collect::<Vec<_>>(),
696 other_seed.iter().map(|s| &s.finding_id).collect::<Vec<_>>()
697 );
698 }
699}