metrale_speculative/
spec_cost_fit.rs1use std::collections::BTreeMap;
20
21use super::{AcceptanceCalibration, DrafterKey};
22
23pub const MIN_OUTCOMES: u64 = 100;
26
27impl AcceptanceCalibration {
28 pub fn fit(
32 drafter: DrafterKey,
33 edges: &[f32],
34 accepted: &[u64],
35 rejected: &[u64],
36 steps: &BTreeMap<(usize, usize), u64>,
37 ) -> Result<Self, String> {
38 if edges.is_empty() || edges.len() != accepted.len() || edges.len() != rejected.len() {
39 return Err("acceptance fit: edges and counts differ in length".into());
40 }
41 let mut p_accept = vec![0.0; edges.len()];
42 let mut pools: Vec<(Vec<usize>, u64, u64)> = Vec::new();
43 let mut open: (Vec<usize>, u64, u64) = (Vec::new(), 0, 0);
44 for i in (0..edges.len()).rev() {
45 open.0.push(i);
46 open.1 += accepted[i];
47 open.2 += accepted[i] + rejected[i];
48 if open.2 >= MIN_OUTCOMES {
49 pools.push(std::mem::take(&mut open));
50 }
51 }
52 if !open.0.is_empty() {
53 let Some(last) = pools.last_mut() else {
54 return Err(format!(
55 "acceptance fit: {} reached drafts, fewer than {MIN_OUTCOMES}",
56 open.2
57 ));
58 };
59 last.0.extend(open.0);
60 last.1 += open.1;
61 last.2 += open.2;
62 }
63 for (members, acc, total) in &pools {
64 for &i in members {
65 p_accept[i] = *acc as f64 / *total as f64;
66 }
67 }
68 let max_drafts = steps.keys().map(|&(d, _)| d).max().unwrap_or(0);
69 let mut prior_by_position = Vec::new();
70 for j in 1..=max_drafts {
71 let sum = |min_a: usize| -> u64 {
72 steps
73 .iter()
74 .filter(|&(&(d, a), _)| d >= j && a >= min_a)
75 .map(|(_, &n)| n)
76 .sum()
77 };
78 let reached = sum(j - 1);
79 if reached < MIN_OUTCOMES {
80 break;
81 }
82 prior_by_position.push(sum(j) as f64 / reached as f64);
83 }
84 if prior_by_position.is_empty() {
85 return Err(format!(
86 "acceptance fit: fewer than {MIN_OUTCOMES} verify steps reached the first draft"
87 ));
88 }
89 Ok(Self {
90 drafter,
91 edges: edges.to_vec(),
92 p_accept,
93 prior_by_position,
94 })
95 }
96
97 pub fn render(&self) -> Result<String, String> {
99 let d = &self.drafter;
100 let text = format!(
101 "[drafter]\nweights_sha256 = {:?}\nvocab = {}\nquantization = {:?}\ncontext = {}\n\n\
102 [acceptance]\nedges = {:?}\np_accept = {:?}\nprior_by_position = {:?}\n",
103 d.weights_sha256,
104 d.vocab,
105 d.quantization,
106 d.context,
107 self.edges,
108 self.p_accept,
109 self.prior_by_position
110 );
111 if Self::parse(&text)? != *self {
112 return Err("acceptance calibration: the rendered file does not read back".into());
113 }
114 Ok(text)
115 }
116}
117
118#[cfg(test)]
119#[path = "spec_cost_fit_tests.rs"]
120mod tests;