Skip to main content

rustev_eval/
fit.rs

1//! Temperature fitting on the calibration split (spec 004, 3.6).
2//!
3//! A bounded grid search: mean log loss is evaluated at each caller-supplied
4//! temperature with the core's own `temperature/1` transformation, and the
5//! lowest loss wins, ties to the smaller temperature. It claims only the
6//! best evaluated candidate, never a global optimum. The fitted
7//! `rustev.calibration/1` artifact comes with a `rustev.calibration-fit/1`
8//! record binding it to its dataset, split, sources, cases and
9//! configuration, because the calibration schema has no split field and its
10//! dataset id alone never proves a disjoint evaluation. Synthetic fits
11//! establish mechanics only.
12
13use std::collections::{BTreeMap, BTreeSet};
14use std::fmt;
15
16use rustev_contract::calibration::{
17    CalibrationArtifact, CalibrationBinding, CalibrationMethod, CalibrationParameters,
18};
19use rustev_contract::decimal::Decimal;
20use rustev_contract::eval_report::Split;
21use rustev_contract::ids::{
22    ArtifactId, CalibrationId, ContentDigest, DatasetId, EvaluatorConfigId, SnapshotId,
23};
24use rustev_contract::limits::REPLAY_V1;
25use rustev_contract::output::RawOutput;
26use rustev_contract::{Identified, schema};
27use rustev_core::calibrate::{self, BindingCheck, CalibrationInput};
28use serde::{Deserialize, Serialize};
29
30use crate::config::{EvaluatorConfig, MAX_GRID};
31use crate::dataset::{DatasetManifest, Label};
32use crate::doc::eval_document;
33
34pub const CALIBRATION_FIT: &str = "rustev.calibration-fit/1";
35
36#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
37#[serde(deny_unknown_fields)]
38pub struct CalibrationFit {
39    pub schema: String,
40    pub calibration: CalibrationId,
41    pub dataset: DatasetId,
42    /// Always `calibration`.
43    pub split: Split,
44    pub config: EvaluatorConfigId,
45    /// Sorted; the fit's sources, cases and snapshots, for leakage checks.
46    pub sources: Vec<String>,
47    pub cases: Vec<String>,
48    pub snapshots: Vec<SnapshotId>,
49    /// The grid evaluated, in order.
50    pub grid: Vec<Decimal>,
51    /// Cases fitted, and cases of the split left out by reason.
52    pub fitted: u64,
53    pub excluded: BTreeMap<String, u64>,
54    pub temperature: Decimal,
55    /// Mean log loss at the chosen temperature, half-even at 9 digits.
56    pub mean_log_loss: Decimal,
57}
58
59eval_document!(CalibrationFit, CALIBRATION_FIT, REPLAY_V1, ContentDigest);
60
61/// The step being calibrated.
62pub struct FitTarget<'a> {
63    pub artifact: &'a ArtifactId,
64    pub task: &'a str,
65    pub question: &'a str,
66    pub options: &'a [String],
67    pub tolerance: Decimal,
68}
69
70/// Why nothing was fitted.
71#[derive(Debug, Clone, PartialEq, Eq)]
72pub enum FitError {
73    /// No grid, or one the configuration refuses.
74    Grid,
75    /// A sample from a split other than calibration.
76    WrongSplit {
77        case: String,
78    },
79    UnknownCase {
80        case: String,
81    },
82    /// No labeled sample to fit.
83    Empty,
84    NonFinite {
85        case: String,
86    },
87    /// The raw output does not answer exactly the target's options.
88    Inconsistent {
89        case: String,
90    },
91    /// A distribution with no mass on the true label: no temperature helps.
92    ZeroSupport {
93        case: String,
94    },
95    /// Every temperature gave an infinite loss.
96    NoFiniteLoss,
97    /// A label that is not an option.
98    Label {
99        case: String,
100    },
101    /// The manifest itself is refused (leakage, duplicates, provenance).
102    Dataset(crate::dataset::DatasetError),
103    /// A case sampled twice would weigh twice.
104    DuplicateSample {
105        case: String,
106    },
107    /// The chosen loss has no Decimal form.
108    Loss,
109}
110
111impl fmt::Display for FitError {
112    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
113        write!(f, "{self:?}")
114    }
115}
116
117impl std::error::Error for FitError {}
118
119fn artifact_at(target: &FitTarget<'_>, dataset: &DatasetId, t: Decimal) -> CalibrationArtifact {
120    CalibrationArtifact {
121        schema: schema::CALIBRATION.into(),
122        binding: CalibrationBinding {
123            artifact: target.artifact.clone(),
124            task: target.task.into(),
125            question: target.question.into(),
126            dataset: dataset.clone(),
127            method: CalibrationMethod::Temperature1,
128        },
129        options: target.options.to_vec(),
130        parameters: CalibrationParameters::Temperature { temperature: t },
131    }
132}
133
134fn to_decimal(x: f64) -> Result<Decimal, FitError> {
135    Decimal::parse(&format!("{x:.9}")).map_err(|_| FitError::Loss)
136}
137
138/// Fit `temperature/1` on the calibration split. `samples` are raw logits or
139/// distributions by case id.
140pub fn fit_temperature(
141    manifest: &DatasetManifest,
142    config: &EvaluatorConfig,
143    target: &FitTarget<'_>,
144    samples: &[(String, RawOutput)],
145) -> Result<(CalibrationArtifact, CalibrationFit), FitError> {
146    let grid = &config.temperature_grid;
147    if grid.is_empty() || grid.len() > MAX_GRID || config.check().is_err() {
148        return Err(FitError::Grid);
149    }
150    manifest.check_structure().map_err(FitError::Dataset)?;
151    let dataset = manifest.id().map_err(|_| FitError::Grid)?;
152    let config_id = config.id().map_err(|_| FitError::Grid)?;
153    let by_id: BTreeMap<&str, _> = manifest.cases.iter().map(|c| (c.id.as_str(), c)).collect();
154    let options: BTreeSet<&str> = target.options.iter().map(String::as_str).collect();
155    let mut rows = vec![];
156    let mut excluded = BTreeMap::new();
157    let mut seen = BTreeSet::new();
158    for (id, raw) in samples {
159        if !seen.insert(id.as_str()) {
160            return Err(FitError::DuplicateSample { case: id.clone() });
161        }
162        let case = by_id
163            .get(id.as_str())
164            .ok_or_else(|| FitError::UnknownCase { case: id.clone() })?;
165        if case.split != Split::Calibration {
166            return Err(FitError::WrongSplit { case: id.clone() });
167        }
168        let Label::Value(label) = &case.label else {
169            *excluded.entry("unlabeled".to_string()).or_insert(0u64) += 1;
170            continue;
171        };
172        if !options.contains(label.as_str()) {
173            return Err(FitError::Label { case: id.clone() });
174        }
175        let m = match raw {
176            RawOutput::Logits(m) | RawOutput::Distribution(m) => m,
177            _ => return Err(FitError::Inconsistent { case: id.clone() }),
178        };
179        if m.keys().map(String::as_str).collect::<BTreeSet<_>>() != options {
180            return Err(FitError::Inconsistent { case: id.clone() });
181        }
182        if m.values().any(|v| !v.is_finite()) {
183            return Err(FitError::NonFinite { case: id.clone() });
184        }
185        if matches!(raw, RawOutput::Distribution(_)) && m[label] <= 0.0 {
186            return Err(FitError::ZeroSupport { case: id.clone() });
187        }
188        rows.push((*case, label.clone(), raw.clone()));
189    }
190    let unsampled = manifest
191        .split(Split::Calibration)
192        .filter(|c| !seen.contains(c.id.as_str()))
193        .count() as u64;
194    if unsampled > 0 {
195        excluded.insert("no-sample".to_string(), unsampled);
196    }
197    if rows.is_empty() {
198        return Err(FitError::Empty);
199    }
200    let check = BindingCheck {
201        artifact: target.artifact,
202        task: target.task,
203        question: target.question,
204        options: target.options,
205    };
206    let mut best: Option<(Decimal, f64)> = None;
207    for t in grid {
208        let art = artifact_at(target, &dataset, *t);
209        let Ok(cid) = art.id() else {
210            return Err(FitError::Grid);
211        };
212        let mut sum = 0.0;
213        for (case, label, raw) in &rows {
214            let input = match raw {
215                RawOutput::Logits(m) => CalibrationInput::Logits(m),
216                RawOutput::Distribution(m) => CalibrationInput::Distribution(m),
217                _ => unreachable!("checked above"),
218            };
219            let p =
220                calibrate::apply(&art, &cid, &check, input, target.tolerance).map_err(|_| {
221                    FitError::Inconsistent {
222                        case: case.id.clone(),
223                    }
224                })?;
225            let q = p.mass(label).unwrap_or(0.0);
226            sum += if q > 0.0 {
227                -rustev_core::math::ln(q)
228            } else {
229                f64::INFINITY
230            };
231        }
232        let loss = sum / rows.len() as f64;
233        if loss.is_finite() && best.is_none_or(|(_, b)| loss < b) {
234            best = Some((*t, loss));
235        }
236    }
237    let (t, loss) = best.ok_or(FitError::NoFiniteLoss)?;
238    let artifact = artifact_at(target, &dataset, t);
239    let mut sources: Vec<String> = rows.iter().map(|(c, ..)| c.source.clone()).collect();
240    let mut cases: Vec<String> = rows.iter().map(|(c, ..)| c.id.clone()).collect();
241    let mut snapshots: Vec<SnapshotId> = rows.iter().map(|(c, ..)| c.snapshot.clone()).collect();
242    for v in [&mut sources, &mut cases] {
243        v.sort();
244        v.dedup();
245    }
246    snapshots.sort();
247    snapshots.dedup();
248    let fit = CalibrationFit {
249        schema: CALIBRATION_FIT.into(),
250        calibration: artifact.id().map_err(|_| FitError::Grid)?,
251        dataset,
252        split: Split::Calibration,
253        config: config_id,
254        sources,
255        cases,
256        snapshots,
257        grid: grid.clone(),
258        fitted: rows.len() as u64,
259        excluded,
260        temperature: t,
261        mean_log_loss: to_decimal(loss)?,
262    };
263    Ok((artifact, fit))
264}
265
266/// Whether a calibration artifact is qualified for evaluating a split.
267#[derive(Debug, Clone, PartialEq, Eq)]
268pub enum Qualification {
269    Qualified,
270    /// No or mismatched lineage: never a pass.
271    Unknown {
272        reason: String,
273    },
274    Refused {
275        reason: String,
276    },
277}
278
279/// Qualification needs the companion fit record and a different, disjoint
280/// split with no case, source or snapshot shared with the fit.
281pub fn qualify(
282    artifact: &CalibrationArtifact,
283    fit: Option<&CalibrationFit>,
284    evaluation: &DatasetManifest,
285    split: Split,
286) -> Qualification {
287    if let Err(e) = evaluation.check_structure() {
288        return Qualification::Refused {
289            reason: format!("the evaluation manifest is refused: {e}"),
290        };
291    }
292    let Some(fit) = fit else {
293        return Qualification::Unknown {
294            reason: "no fit record: lineage is missing".into(),
295        };
296    };
297    if artifact.id().ok().as_ref() != Some(&fit.calibration) {
298        return Qualification::Unknown {
299            reason: "the fit record is for another artifact".into(),
300        };
301    }
302    let same_dataset = evaluation.id().ok().as_ref() == Some(&fit.dataset);
303    if same_dataset && split == fit.split {
304        return Qualification::Refused {
305            reason: "evaluated on the split it was fitted on".into(),
306        };
307    }
308    for c in evaluation.split(split) {
309        if (same_dataset && fit.cases.contains(&c.id))
310            || fit.sources.contains(&c.source)
311            || fit.snapshots.contains(&c.snapshot)
312        {
313            return Qualification::Refused {
314                reason: format!("case {} shares a fit source", c.id),
315            };
316        }
317    }
318    Qualification::Qualified
319}