1use 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 pub split: Split,
44 pub config: EvaluatorConfigId,
45 pub sources: Vec<String>,
47 pub cases: Vec<String>,
48 pub snapshots: Vec<SnapshotId>,
49 pub grid: Vec<Decimal>,
51 pub fitted: u64,
53 pub excluded: BTreeMap<String, u64>,
54 pub temperature: Decimal,
55 pub mean_log_loss: Decimal,
57}
58
59eval_document!(CalibrationFit, CALIBRATION_FIT, REPLAY_V1, ContentDigest);
60
61pub 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#[derive(Debug, Clone, PartialEq, Eq)]
72pub enum FitError {
73 Grid,
75 WrongSplit {
77 case: String,
78 },
79 UnknownCase {
80 case: String,
81 },
82 Empty,
84 NonFinite {
85 case: String,
86 },
87 Inconsistent {
89 case: String,
90 },
91 ZeroSupport {
93 case: String,
94 },
95 NoFiniteLoss,
97 Label {
99 case: String,
100 },
101 Dataset(crate::dataset::DatasetError),
103 DuplicateSample {
105 case: String,
106 },
107 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
138pub 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#[derive(Debug, Clone, PartialEq, Eq)]
268pub enum Qualification {
269 Qualified,
270 Unknown {
272 reason: String,
273 },
274 Refused {
275 reason: String,
276 },
277}
278
279pub 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}