kime_core/
answer.rs

1//! Answers, built from a question's option logits and act logits the way Laya builds them.
2//!
3//! This is Laya's `Agent._decode_answers` (unchanged from 0.3.7 to 0.3.11), done in the same float
4//! types numpy uses so the rounded numbers match: the calibrated softmax and the entropy in f32,
5//! the expected score in f64, and sums in numpy's pairwise order. `exp` and `log` are numpy's own
6//! AVX2 routines rather than libm's, since they differ in the last bit often enough to move a
7//! rounded probability: 2 of the 400 parity responses changed in the fourth decimal with libm.
8
9use serde_json::{Map, Value, json};
10
11use crate::request::{Criteria, QType, Question};
12
13/// The `model` field of a compat response.
14pub const LAYA_MODEL: &str = "laya-rl-agent";
15
16/// The range Laya 0.3.9 and later confine a checkpoint's temperatures to. The English checkpoint
17/// ships 0.1006 for `choice:11+`, which would sharpen the logits tenfold.
18pub const TEMPERATURE_RANGE: (f64, f64) = (0.5, 5.0);
19
20/// A temperature Laya would use: `t` confined to [`TEMPERATURE_RANGE`], or 1 when it is not a
21/// finite number.
22#[must_use]
23pub fn clamp_temperature(t: f64) -> f64 {
24    if t.is_finite() { t.clamp(TEMPERATURE_RANGE.0, TEMPERATURE_RANGE.1) } else { 1.0 }
25}
26
27/// Laya's `temp_bucket`: the key of `temperature_by_options` for a type and an option count.
28#[must_use]
29pub fn temperature_bucket(t: QType, k: usize) -> String {
30    let size = match k {
31        0..=2 => "2",
32        3..=5 => "3-5",
33        6..=10 => "6-10",
34        _ => "11+",
35    };
36    format!("{}:{size}", t.as_str())
37}
38
39/// The calibration of a compat checkpoint, from its `rl_agent_config.json`, already clamped.
40#[derive(Debug, Clone, PartialEq)]
41pub struct Temperatures {
42    by_type: [f64; 3],
43    by_options: Vec<(String, f64)>,
44}
45
46impl Temperatures {
47    /// From the checkpoint's `temperature` (choice, score, noul) and `temperature_by_options`.
48    #[must_use]
49    pub fn new(by_type: [f64; 3], by_options: &[(String, f64)]) -> Self {
50        Self {
51            by_type: by_type.map(clamp_temperature),
52            by_options: by_options
53                .iter()
54                .map(|(k, t)| (k.clone(), clamp_temperature(*t)))
55                .collect(),
56        }
57    }
58
59    /// The temperature for a question of type `t` with `k` options.
60    #[must_use]
61    pub fn get(&self, t: QType, k: usize) -> f64 {
62        let bucket = temperature_bucket(t, k);
63        self.by_options
64            .iter()
65            .find(|(b, _)| *b == bucket)
66            .map_or(self.by_type[t.index()], |(_, t)| *t)
67    }
68}
69
70impl Default for Temperatures {
71    fn default() -> Self {
72        Self::new([1.0; 3], &[])
73    }
74}
75
76/// One question's answer.
77#[derive(Debug, Clone, PartialEq)]
78pub enum Answer {
79    /// A choice question.
80    Choice {
81        /// The most likely label, the first one on a tie.
82        choice: String,
83        /// Each label with its calibrated probability, in option order.
84        probabilities: Vec<(String, f64)>,
85        /// Normalized entropy confidence.
86        confidence: f64,
87        /// The act head's probability that acting on this answer is right.
88        act_probability: f64,
89    },
90    /// A score question.
91    Score {
92        /// The expected level, which can fall between levels.
93        score: f64,
94        /// The levels exactly as sent.
95        legend: Vec<Value>,
96        /// The probability of each level.
97        probabilities: Vec<f64>,
98        /// Normalized entropy confidence.
99        confidence: f64,
100        /// See [`Answer::Choice`].
101        act_probability: f64,
102    },
103    /// A noul question.
104    Noul {
105        /// The probability the statement is true.
106        noul: f64,
107        /// `max(noul, 1 - noul)`.
108        confidence: f64,
109        /// See [`Answer::Choice`].
110        act_probability: f64,
111    },
112}
113
114/// Rounds to `digits` places the way Python's `round` does: correctly, ties to even.
115#[must_use]
116pub fn py_round(x: f64, digits: usize) -> f64 {
117    format!("{x:.digits$}").parse().unwrap_or(x)
118}
119
120impl Answer {
121    /// The answer in Laya's JSON shape, with every number rounded to 4 places.
122    #[must_use]
123    pub fn to_json(&self) -> Value {
124        let r = |x: f64| py_round(x, 4);
125        match self {
126            Answer::Choice { choice, probabilities, confidence, act_probability } => json!({
127                "type": "choice",
128                "choice": choice,
129                "probabilities": probabilities.iter().map(|(k, p)| (k.clone(), json!(r(*p)))).collect::<Map<_, _>>(),
130                "confidence": r(*confidence),
131                "action": {"act_probability": r(*act_probability)},
132            }),
133            Answer::Score { score, legend, probabilities, confidence, act_probability } => json!({
134                "type": "score",
135                "score": r(*score),
136                "legend": legend.iter().enumerate().map(|(i, v)| (i.to_string(), v.clone())).collect::<Map<_, _>>(),
137                "probabilities": probabilities.iter().enumerate().map(|(i, p)| (i.to_string(), json!(r(*p)))).collect::<Map<_, _>>(),
138                "confidence": r(*confidence),
139                "action": {"act_probability": r(*act_probability)},
140            }),
141            Answer::Noul { noul, confidence, act_probability } => json!({
142                "type": "noul",
143                "noul": r(*noul),
144                "confidence": r(*confidence),
145                "action": {"act_probability": r(*act_probability)},
146            }),
147        }
148    }
149}
150
151/// The answers to one request.
152#[derive(Debug, Clone, PartialEq)]
153pub struct Response {
154    /// The model that answered.
155    pub model: String,
156    /// The answers, in request order.
157    pub answers: Vec<(String, Answer)>,
158    /// Tokens read over all the request's sequences.
159    pub input_tokens: usize,
160}
161
162impl Response {
163    /// The answer to question `id`.
164    #[must_use]
165    pub fn get(&self, id: &str) -> Option<&Answer> {
166        self.answers.iter().find(|(q, _)| q == id).map(|(_, a)| a)
167    }
168
169    /// The response in Laya's JSON shape.
170    #[must_use]
171    pub fn to_json(&self) -> Value {
172        let answers: Map<_, _> =
173            self.answers.iter().map(|(id, a)| (id.clone(), a.to_json())).collect();
174        json!({
175            "model": self.model,
176            "answers": answers,
177            "usage": {"input_tokens": self.input_tokens, "output_tokens": 0},
178        })
179    }
180}
181
182/// numpy's pairwise sum, which is the order `ndarray.sum` adds contiguous values in.
183fn pairwise<T: Copy + Default + std::ops::Add<Output = T>>(x: &[T]) -> T {
184    let n = x.len();
185    if n < 8 {
186        return x.iter().fold(T::default(), |a, &b| a + b);
187    }
188    if n <= 128 {
189        let mut r = [x[0], x[1], x[2], x[3], x[4], x[5], x[6], x[7]];
190        let body = n - n % 8;
191        for c in x[8..body].as_chunks::<8>().0 {
192            for j in 0..8 {
193                r[j] = r[j] + c[j];
194            }
195        }
196        let mut s = ((r[0] + r[1]) + (r[2] + r[3])) + ((r[4] + r[5]) + (r[6] + r[7]));
197        for &v in &x[body..] {
198            s = s + v;
199        }
200        return s;
201    }
202    let half = n / 2 - (n / 2) % 8;
203    pairwise(&x[..half]) + pairwise(&x[half..])
204}
205
206/// numpy's float32 `exp` on x86 with AVX2 and FMA (`simd_exp_FLOAT` in
207/// `loops_exponent_log.dispatch.c.src`), one lane at a time: Cody-Waite reduction by ln 2, a 5/2
208/// rational approximation, then the power of two added straight into the exponent bits.
209// The constants are numpy's, digits and all, so they can be checked against its source.
210#[allow(clippy::excessive_precision)]
211fn np_exp(x: f32) -> f32 {
212    const XMAX: f32 = 88.722_84;
213    const XMIN: f32 = -103.972_08;
214    const MAGIC: f32 = 12_582_912.0;
215    const P: [f32; 6] = [
216        9.999_999_999_980_870_924_916e-1,
217        7.257_664_613_233_124_478_488e-1,
218        2.473_615_434_895_520_810_817e-1,
219        5.114_512_081_637_298_353_406e-2,
220        6.757_896_990_527_504_603_057e-3,
221        5.082_762_527_590_693_718_096e-4,
222    ];
223    const Q: [f32; 3] = [1.0, -2.742_335_390_411_667_452_936e-1, 2.159_509_375_685_829_852_307e-2];
224    if x.is_nan() {
225        return x;
226    }
227    if x >= XMAX {
228        return f32::INFINITY;
229    }
230    if x <= XMIN {
231        return 0.0;
232    }
233    let q = (x * std::f32::consts::LOG2_E + MAGIC) - MAGIC;
234    let r = q.mul_add(-1.428_606_77e-6, q.mul_add(-6.931_457_52e-1, x));
235    let r = q.mul_add(0.0, r);
236    let num =
237        P[5].mul_add(r, P[4]).mul_add(r, P[3]).mul_add(r, P[2]).mul_add(r, P[1]).mul_add(r, P[0]);
238    let den = Q[2].mul_add(r, Q[1]).mul_add(r, Q[0]);
239    let poly = num / den;
240    // q is a whole number here, and at most 128 in size.
241    let shift =
242        |poly: f32, q: f32| f32::from_bits(poly.to_bits().wrapping_add(((q as i32) << 23) as u32));
243    if q <= -125.0 {
244        let diff = (-(q + 125.0)) as u32;
245        shift(poly, -125.0) / (1u32 << diff) as f32
246    } else {
247        shift(poly, q)
248    }
249}
250
251/// numpy's float32 `log` on x86 with AVX2 and FMA (`simd_log_FLOAT`), for a positive finite `x`:
252/// the mantissa scaled into (sqrt(1/2), sqrt(2)], a 5/5 rational approximation of `log(1 + m)`,
253/// plus the exponent times ln 2.
254#[allow(clippy::excessive_precision)]
255fn np_log(x: f32) -> f32 {
256    const P: [f32; 6] = [
257        0.0,
258        9.999_999_999_999_998_702_752e-1,
259        2.112_677_543_073_053_063_722,
260        1.480_000_633_576_506_585_156,
261        3.808_837_741_388_407_920_751e-1,
262        2.589_979_117_907_922_693_523e-2,
263    ];
264    const Q: [f32; 6] = [
265        1.0,
266        2.612_677_543_073_109_236_779,
267        2.453_006_071_784_736_363_091,
268        9.864_942_958_519_418_960_339e-1,
269        1.546_476_374_983_906_719_538e-1,
270        5.875_095_403_124_574_342_950e-3,
271    ];
272    if x.is_nan() || x < 0.0 {
273        return f32::NAN;
274    }
275    if x == 0.0 {
276        return f32::NEG_INFINITY;
277    }
278    if x.is_infinite() {
279        return x;
280    }
281    let (bits, bias) = if x < f32::MIN_POSITIVE {
282        ((x * f32::from_bits(0x7180_0000)).to_bits(), 100.0)
283    } else {
284        (x.to_bits(), 0.0)
285    };
286    let mut exponent = ((bits >> 23) as i32 - 126) as f32 - bias;
287    let mut m = f32::from_bits((bits & 0x7f_ffff) | (126 << 23));
288    if m <= std::f32::consts::FRAC_1_SQRT_2 {
289        m += m;
290        exponent -= 1.0;
291    }
292    let m = m - 1.0;
293    let num =
294        P[5].mul_add(m, P[4]).mul_add(m, P[3]).mul_add(m, P[2]).mul_add(m, P[1]).mul_add(m, P[0]);
295    let den =
296        Q[5].mul_add(m, Q[4]).mul_add(m, Q[3]).mul_add(m, Q[2]).mul_add(m, Q[1]).mul_add(m, Q[0]);
297    exponent.mul_add(std::f32::consts::LN_2, num / den)
298}
299
300/// The calibrated distribution over `logits` at temperature `t`, in f32 as numpy computes it.
301fn softmax(logits: &[f32], t: f64) -> Vec<f32> {
302    let t = t as f32;
303    let z: Vec<f32> = logits.iter().map(|&l| l / t).collect();
304    let top = z.iter().copied().fold(f32::NEG_INFINITY, f32::max);
305    let e: Vec<f32> = z.iter().map(|&v| np_exp(v - top)).collect();
306    let s = pairwise(&e);
307    e.iter().map(|&v| v / s).collect()
308}
309
310/// Laya's `confidence_from_probs`, in f32.
311fn entropy_confidence(p: &[f32]) -> f64 {
312    let k = p.len();
313    if k < 2 {
314        return 1.0;
315    }
316    let terms: Vec<f32> = p.iter().map(|&v| v * np_log(v.clamp(1e-12, 1.0))).collect();
317    let ent = -pairwise(&terms);
318    f64::from((1.0 - ent / (k as f64).ln() as f32).clamp(0.0, 1.0))
319}
320
321/// The act head's first probability, as `torch.softmax` gives it.
322#[must_use]
323pub fn act_probability(act: [f32; 2]) -> f64 {
324    let top = act[0].max(act[1]);
325    let e = act.map(|v| (v - top).exp());
326    f64::from(e[0] * (1.0 / (e[0] + e[1])))
327}
328
329/// The answer to `q` from its option logits, in option order, and its two act logits.
330///
331/// # Panics
332///
333/// If `logits` does not have one value per option.
334#[must_use]
335pub fn laya_answer(q: &Question, logits: &[f32], act: [f32; 2], temps: &Temperatures) -> Answer {
336    let k = q.criteria.len();
337    assert_eq!(logits.len(), k, "one logit per option");
338    let p = softmax(logits, temps.get(q.qtype, k));
339    let act_probability = act_probability(act);
340    match &q.criteria {
341        Criteria::Choice(opts) => {
342            let mut best = 0;
343            for (i, &v) in p.iter().enumerate() {
344                if v > p[best] {
345                    best = i;
346                }
347            }
348            Answer::Choice {
349                choice: opts[best].label.clone(),
350                probabilities: opts
351                    .iter()
352                    .zip(&p)
353                    .map(|(o, &v)| (o.label.clone(), f64::from(v)))
354                    .collect(),
355                confidence: entropy_confidence(&p),
356                act_probability,
357            }
358        }
359        Criteria::Score(levels) => {
360            let weighted: Vec<f64> =
361                p.iter().enumerate().map(|(i, &v)| i as f64 * f64::from(v)).collect();
362            Answer::Score {
363                score: pairwise(&weighted),
364                legend: levels.clone(),
365                probabilities: p.iter().map(|&v| f64::from(v)).collect(),
366                confidence: entropy_confidence(&p),
367                act_probability,
368            }
369        }
370        Criteria::Noul { .. } => {
371            let yes = f64::from(p[1]);
372            Answer::Noul { noul: yes, confidence: yes.max(1.0 - yes), act_probability }
373        }
374    }
375}
376
377#[cfg(test)]
378// These values are exact, so comparing them exactly is the point.
379#[allow(clippy::float_cmp)]
380mod tests {
381    use super::*;
382
383    #[test]
384    fn pairwise_order() {
385        // Nine values: eight lanes plus one left over, which a left fold adds in another order.
386        let x = [1e8f32, 1.0, -1e8, 1.0, 0.5, 0.25, 3.0, 7.0, 1.0];
387        let lanes = ((1e8f32 + 1.0) + (-1e8 + 1.0)) + ((0.5 + 0.25) + (3.0 + 7.0)) + 1.0;
388        assert_eq!(pairwise(&x).to_bits(), lanes.to_bits());
389    }
390
391    #[test]
392    fn numpy_exp_log() {
393        // Values numpy 2.5 gives on an AVX2 machine, where libm's differ in the last bit.
394        for (x, want) in [(-0.0, 1.0f32), (-1.0, 0.367_879_43), (-20.0, 2.061_153_6e-9)] {
395            assert!((np_exp(x) - want).abs() <= want * 3e-7, "exp {x}");
396        }
397        assert_eq!(np_exp(0.0), 1.0);
398        assert_eq!(np_log(1.0), 0.0);
399        assert!((np_log(1e-12) - 1e-12f32.ln()).abs() < 1e-5);
400        assert!((np_log(0.3) - 0.3f32.ln()).abs() < 1e-6);
401    }
402
403    #[test]
404    fn rounding() {
405        assert_eq!(py_round(0.25, 1), 0.2);
406        assert_eq!(py_round(0.965_449_999, 4), 0.9654);
407        assert_eq!(py_round(1.0, 4), 1.0);
408    }
409
410    #[test]
411    fn buckets_and_clamp() {
412        let t = Temperatures::new([1.5, 1.2, 2.0], &[("choice:11+".into(), 0.1)]);
413        assert_eq!(t.get(QType::Choice, 12), 0.5);
414        assert_eq!(t.get(QType::Choice, 4), 1.5);
415        assert_eq!(temperature_bucket(QType::Score, 6), "score:6-10");
416        assert_eq!(clamp_temperature(f64::NAN), 1.0);
417    }
418}
kime_core/answer.rs

kime_core/
answer.rs
