laya/
common.rs

1//! Core model-independent logic: state serialization (Python `json.dumps` parity),
2//! option rendering, token-sequence construction, collation, and calibration helpers.
3//!
4//! Ported from `laya/common.py` (source of truth) and mirrored against `laya-ts/src/common.ts`.
5
6use crate::error::{LayaError, Result};
7use serde_json::Value;
8
9/// The three question types. Numeric indices match Python `QTYPES`
10/// (`choice=0, score=1, noul=2`) — they select the model's type embedding row.
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub enum QType {
13    Choice,
14    Score,
15    Noul,
16}
17
18impl QType {
19    pub fn index(self) -> usize {
20        match self {
21            QType::Choice => 0,
22            QType::Score => 1,
23            QType::Noul => 2,
24        }
25    }
26
27    pub fn name(self) -> &'static str {
28        match self {
29            QType::Choice => "choice",
30            QType::Score => "score",
31            QType::Noul => "noul",
32        }
33    }
34
35    pub fn from_name(s: &str) -> Option<QType> {
36        match s {
37            "choice" => Some(QType::Choice),
38            "score" => Some(QType::Score),
39            "noul" => Some(QType::Noul),
40            _ => None,
41        }
42    }
43
44    pub fn from_index(i: usize) -> Option<QType> {
45        match i {
46            0 => Some(QType::Choice),
47            1 => Some(QType::Score),
48            2 => Some(QType::Noul),
49            _ => None,
50        }
51    }
52}
53
54/// A validated + normalized question, ready for sequence building and decoding.
55/// `crit` is kept as a JSON value in the shape the type expects (choice: object,
56/// score: array, noul: object-or-null), matching the normalization in `Agent::to_internal`.
57#[derive(Debug, Clone)]
58pub struct InternalQ {
59    pub t: QType,
60    pub ins: String,
61    pub crit: Value,
62    pub labels: Option<Value>,
63}
64
65/// Abstraction over a checkpoint tokenizer. Mirrors `TokenizerLike` in laya-ts.
66/// `encode` corresponds to HF `tokenizer(text, add_special_tokens=False)["input_ids"]`.
67pub trait Tokenizer {
68    fn cls_id(&self) -> u32;
69    fn sep_id(&self) -> u32;
70    fn mask_id(&self) -> u32;
71    fn pad_id(&self) -> u32;
72    fn mask_token(&self) -> &str;
73    fn encode(&self, text: &str) -> Vec<u32>;
74}
75
76const DEFAULT_NOUL_FALSE: &str = "false";
77const DEFAULT_NOUL_TRUE: &str = "true";
78
79/// Replica of Python `json.dumps(value, ensure_ascii=False, separators=(", ", ": "))`.
80///
81/// This differs from `serde_json::to_string` (which uses no spaces after `,`/`:`); the exact
82/// spacing matters because serialized states/criteria are tokenized, so a spacing difference
83/// changes token ids. Mirrors `pyJson` in `laya-ts/src/common.ts`.
84pub fn py_json(v: &Value) -> String {
85    match v {
86        Value::Null => "null".to_string(),
87        Value::Bool(b) => {
88            if *b {
89                "true".to_string()
90            } else {
91                "false".to_string()
92            }
93        }
94        Value::Number(n) => n.to_string(),
95        // serde_json escapes control chars / `"` / `\` and keeps non-ASCII raw — same as
96        // Python `ensure_ascii=False`.
97        Value::String(_) => serde_json::to_string(v).unwrap_or_else(|_| "\"\"".to_string()),
98        Value::Array(a) => {
99            let parts: Vec<String> = a.iter().map(py_json).collect();
100            format!("[{}]", parts.join(", "))
101        }
102        Value::Object(m) => {
103            let parts: Vec<String> = m
104                .iter()
105                .map(|(k, val)| {
106                    let key = serde_json::to_string(&Value::String(k.clone()))
107                        .unwrap_or_else(|_| "\"\"".to_string());
108                    format!("{}: {}", key, py_json(val))
109                })
110                .collect();
111            format!("{{{}}}", parts.join(", "))
112        }
113    }
114}
115
116/// A state as the model sees it: raw strings pass through, everything else becomes compact JSON.
117pub fn serialize_state(state: &Value) -> String {
118    match state {
119        Value::String(s) => s.clone(),
120        other => py_json(other),
121    }
122}
123
124/// Render one criterion value as text (strings pass through, structured values become JSON).
125pub fn render_criterion(v: &Value) -> String {
126    match v {
127        Value::String(s) => s.clone(),
128        other => py_json(other),
129    }
130}
131
132/// True when a criterion value means "no description" (only `null`/`""`; `0`/`false` are real).
133fn is_empty_criterion(v: Option<&Value>) -> bool {
134    match v {
135        None => true,
136        Some(Value::Null) => true,
137        Some(Value::String(s)) if s.is_empty() => true,
138        _ => false,
139    }
140}
141
142/// Resolve noul labels to `(false_label, true_label)`, validating the mapping.
143pub fn resolve_noul_labels(labels: Option<&Value>) -> Result<(String, String)> {
144    let err = || {
145        LayaError::InvalidQuestion(
146            "noul labels must map exactly 'false' and 'true' to distinct non-empty strings"
147                .to_string(),
148        )
149    };
150    let (f, t) = match labels {
151        None | Some(Value::Null) => {
152            return Ok((
153                DEFAULT_NOUL_FALSE.to_string(),
154                DEFAULT_NOUL_TRUE.to_string(),
155            ))
156        }
157        Some(Value::Object(m)) => {
158            let keys: std::collections::BTreeSet<&str> = m.keys().map(|s| s.as_str()).collect();
159            let want: std::collections::BTreeSet<&str> = ["false", "true"].into_iter().collect();
160            if keys != want {
161                return Err(err());
162            }
163            (m.get("false"), m.get("true"))
164        }
165        _ => return Err(err()),
166    };
167    let (f, t) = match (f, t) {
168        (Some(Value::String(f)), Some(Value::String(t))) => (f, t),
169        _ => return Err(err()),
170    };
171    let f = f.trim();
172    let t = t.trim();
173    if f.is_empty() || t.is_empty() || f == t {
174        return Err(err());
175    }
176    Ok((f.to_string(), t.to_string()))
177}
178
179/// Render option texts in label-index order. Noul semantic order is always `[false, true]`.
180pub fn render_options(q: &InternalQ) -> Result<Vec<String>> {
181    if q.t != QType::Noul && q.labels.is_some() {
182        return Err(LayaError::InvalidQuestion(
183            "labels is only supported for noul questions".to_string(),
184        ));
185    }
186    match q.t {
187        QType::Choice => {
188            let obj = q.crit.as_object().ok_or_else(|| {
189                LayaError::InvalidQuestion("choice criteria must be an object".to_string())
190            })?;
191            Ok(obj
192                .iter()
193                .map(|(k, v)| {
194                    if is_empty_criterion(Some(v)) {
195                        k.clone()
196                    } else {
197                        format!("{}: {}", k, render_criterion(v))
198                    }
199                })
200                .collect())
201        }
202        QType::Score => {
203            let arr = q.crit.as_array().ok_or_else(|| {
204                LayaError::InvalidQuestion("score criteria must be a list".to_string())
205            })?;
206            Ok(arr
207                .iter()
208                .enumerate()
209                .map(|(i, c)| format!("level {}: {}", i, render_criterion(c)))
210                .collect())
211        }
212        QType::Noul => {
213            let (false_label, true_label) = resolve_noul_labels(q.labels.as_ref())?;
214            let empty = serde_json::Map::new();
215            let crit = q.crit.as_object().unwrap_or(&empty);
216            let false_crit = crit.get("false");
217            let true_crit = crit.get("true");
218            let false_text = if is_empty_criterion(false_crit) {
219                "no, the statement does not hold".to_string()
220            } else {
221                render_criterion(false_crit.unwrap())
222            };
223            let true_text = if is_empty_criterion(true_crit) {
224                "yes, the statement holds".to_string()
225            } else {
226                render_criterion(true_crit.unwrap())
227            };
228            Ok(vec![
229                format!("{}: {}", false_label, false_text),
230                format!("{}: {}", true_label, true_text),
231            ])
232        }
233    }
234}
235
236/// Build one token sequence + option marker positions.
237///
238/// Format: `[CLS] <type> question: <instructions> [SEP] [MASK] opt0 [MASK] opt1 ... [SEP] state [SEP]`.
239/// Ported verbatim from `build_sequence` in `laya/common.py`.
240pub fn build_sequence<T: Tokenizer + ?Sized>(
241    tok: &T,
242    state: &Value,
243    q: &InternalQ,
244    max_len: usize,
245    head_max_len: usize,
246    option_order: Option<&[usize]>,
247    truncate_left: bool,
248) -> Result<(Vec<u32>, Vec<usize>)> {
249    let mask_tok = tok.mask_token();
250    let opts = render_options(q)?;
251    let order: Vec<usize> = match option_order {
252        Some(o) => o.to_vec(),
253        None => (0..opts.len()).collect(),
254    };
255
256    let ins = q.ins.replace(mask_tok, " ");
257    let head_text = format!("{} question: {}", q.t.name(), ins);
258    let mut head_ids = tok.encode(&head_text);
259
260    let mut opt_ids: Vec<Vec<u32>> = Vec::with_capacity(order.len());
261    for &i in &order {
262        let mut o = vec![tok.mask_id()];
263        let enc = tok.encode(&format!(" {}", opts[i].replace(mask_tok, " ")));
264        o.extend(enc.into_iter().take(48));
265        opt_ids.push(o);
266    }
267
268    let sum_opt = |v: &[Vec<u32>]| -> usize { v.iter().map(|o| o.len()).sum() };
269    let mut opt_budget: isize = head_max_len as isize - sum_opt(&opt_ids) as isize;
270    if opt_budget < 16 {
271        let denom = std::cmp::max(1, opt_ids.len()) as isize;
272        let per = std::cmp::max(4, (head_max_len as isize - 16) / denom) as usize;
273        for o in opt_ids.iter_mut() {
274            o.truncate(per);
275        }
276        opt_budget = head_max_len as isize - sum_opt(&opt_ids) as isize;
277    }
278    let head_keep = opt_budget.max(8) as usize;
279    head_ids.truncate(head_keep.min(head_ids.len()));
280
281    let mut ids: Vec<u32> = Vec::new();
282    ids.push(tok.cls_id());
283    ids.extend_from_slice(&head_ids);
284    ids.push(tok.sep_id());
285
286    let mut markers: Vec<usize> = Vec::with_capacity(opt_ids.len());
287    for o in &opt_ids {
288        markers.push(ids.len());
289        ids.extend_from_slice(o);
290    }
291    ids.push(tok.sep_id());
292
293    let room = (max_len as isize - ids.len() as isize - 1).max(0) as usize;
294    let st_all = tok.encode(&serialize_state(state).replace(mask_tok, " "));
295    let st: Vec<u32> = if truncate_left {
296        let start = st_all.len().saturating_sub(room);
297        st_all[start..].to_vec()
298    } else {
299        st_all.into_iter().take(room).collect()
300    };
301    ids.extend_from_slice(&st);
302    ids.push(tok.sep_id());
303    ids.truncate(max_len);
304
305    let markers = markers.into_iter().filter(|&m| m < max_len).collect();
306    Ok((ids, markers))
307}
308
309/// Round to `ndigits` decimals using round-half-to-even, matching Python's built-in `round`.
310pub fn round_half_even(x: f64, ndigits: i32) -> f64 {
311    if !x.is_finite() {
312        return x;
313    }
314    let factor = 10f64.powi(ndigits);
315    let scaled = x * factor;
316    let floor = scaled.floor();
317    let diff = scaled - floor;
318    let rounded = if (diff - 0.5).abs() < 1e-9 {
319        // exact half → round to even
320        if (floor as i64) % 2 == 0 {
321            floor
322        } else {
323            floor + 1.0
324        }
325    } else {
326        scaled.round()
327    };
328    rounded / factor
329}
330
331/// Round to 4 decimals like Python `round(x, 4)` (the precision Laya publishes).
332pub fn round4(x: f64) -> f64 {
333    round_half_even(x, 4)
334}
335
336/// Numerically stable softmax.
337pub fn softmax(z: &[f64]) -> Vec<f64> {
338    if z.is_empty() {
339        return Vec::new();
340    }
341    let m = z.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
342    let exps: Vec<f64> = z.iter().map(|&v| (v - m).exp()).collect();
343    let sum: f64 = exps.iter().sum();
344    exps.into_iter().map(|v| v / sum).collect()
345}
346
347/// Probability mass on the answer being reported: `max(p[:k])`, clamped to `[0, 1]`.
348///
349/// This is the quantity temperature scaling fits and every calibration figure is computed on
350/// (`conf = max(probs)`), so it is the one confidence with the calibration guarantee the gating
351/// docs rely on. [`confidence_from_probs`] reports a different quantity on a different scale and
352/// carries no such guarantee, so the two must not be compared against the same threshold.
353pub fn answer_confidence(p: &[f64], k: usize) -> f64 {
354    if k < 1 {
355        return 1.0;
356    }
357    let p = &p[..k.min(p.len())];
358    let max = p.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
359    max.clamp(0.0, 1.0)
360}
361
362/// Normalized Shannon-entropy confidence: `1 - H(p)/log(k)`, clamped to `[0, 1]`.
363/// `p` is expected to already hold the `k` option probabilities.
364pub fn confidence_from_probs(p: &[f64], k: usize) -> f64 {
365    if k < 2 {
366        return 1.0;
367    }
368    let p = &p[..k.min(p.len())];
369    let ent: f64 = -p.iter().map(|&v| v * v.clamp(1e-12, 1.0).ln()).sum::<f64>();
370    (1.0 - ent / (k as f64).ln()).clamp(0.0, 1.0)
371}
372
373pub const TEMP_MIN: f64 = 0.5;
374pub const TEMP_MAX: f64 = 5.0;
375
376/// A usable temperature: `t` confined to `[TEMP_MIN, TEMP_MAX]`, falling back to 1.0 if not finite.
377pub fn clamp_temperature(t: &Value) -> f64 {
378    let f = match t {
379        Value::Number(n) => n.as_f64(),
380        Value::String(s) => s.trim().parse::<f64>().ok(),
381        _ => None,
382    };
383    match f {
384        Some(x) if x.is_finite() => x.clamp(TEMP_MIN, TEMP_MAX),
385        _ => 1.0,
386    }
387}
388
389/// Temperature bucket key, e.g. `"choice:3-5"`.
390pub fn temp_bucket(qtype: usize, k: usize) -> String {
391    let size = if k <= 2 {
392        "2"
393    } else if k <= 5 {
394        "3-5"
395    } else if k <= 10 {
396        "6-10"
397    } else {
398        "11+"
399    };
400    let name = QType::from_index(qtype)
401        .map(|q| q.name())
402        .unwrap_or("choice");
403    format!("{}:{}", name, size)
404}
405
406/// One tokenized (state, question) pair before collation.
407#[derive(Debug, Clone)]
408pub struct Item {
409    pub ids: Vec<u32>,
410    pub markers: Vec<usize>,
411    pub qtype: usize,
412}
413
414/// A padded batch ready to be turned into tensors.
415#[derive(Debug, Clone)]
416pub struct CollatedBatch {
417    pub input_ids: Vec<Vec<u32>>,
418    pub attention_mask: Vec<Vec<u32>>,
419    pub marker_pos: Vec<Vec<usize>>,
420    pub marker_mask: Vec<Vec<bool>>,
421    pub qtype: Vec<usize>,
422}
423
424/// Collate groups of items into one padded batch. `groups` is a list of per-state item lists,
425/// flattened in order (mirrors `collate_items` in `laya/common.py`). Returns `None` if empty.
426pub fn collate_items(groups: &[Vec<Item>], _pad_id: u32) -> Option<CollatedBatch> {
427    let items: Vec<&Item> = groups.iter().flatten().collect();
428    if items.is_empty() {
429        return None;
430    }
431    let l = items.iter().map(|it| it.ids.len()).max().unwrap();
432    let kmax = items.iter().map(|it| it.markers.len()).max().unwrap();
433
434    let mut input_ids = Vec::with_capacity(items.len());
435    let mut attention_mask = Vec::with_capacity(items.len());
436    let mut marker_pos = Vec::with_capacity(items.len());
437    let mut marker_mask = Vec::with_capacity(items.len());
438    let mut qtype = Vec::with_capacity(items.len());
439
440    for it in &items {
441        let mut ids = it.ids.clone();
442        let mut att = vec![1u32; ids.len()];
443        ids.resize(l, _pad_id);
444        att.resize(l, 0);
445        input_ids.push(ids);
446        attention_mask.push(att);
447
448        let k = it.markers.len();
449        let mut mpos: Vec<usize> = it.markers.clone();
450        mpos.resize(kmax, 0);
451        let mut mmask: Vec<bool> = vec![true; k];
452        mmask.resize(kmax, false);
453        marker_pos.push(mpos);
454        marker_mask.push(mmask);
455
456        qtype.push(it.qtype);
457    }
458
459    Some(CollatedBatch {
460        input_ids,
461        attention_mask,
462        marker_pos,
463        marker_mask,
464        qtype,
465    })
466}
467
468#[cfg(test)]
469mod tests {
470    use super::*;
471    use serde_json::json;
472
473    /// A deterministic fake tokenizer for structural tests (one id per whitespace word).
474    struct Fake;
475    impl Tokenizer for Fake {
476        fn cls_id(&self) -> u32 {
477            1
478        }
479        fn sep_id(&self) -> u32 {
480            2
481        }
482        fn mask_id(&self) -> u32 {
483            3
484        }
485        fn pad_id(&self) -> u32 {
486            0
487        }
488        fn mask_token(&self) -> &str {
489            "[MASK]"
490        }
491        fn encode(&self, text: &str) -> Vec<u32> {
492            text.split_whitespace()
493                .map(|w| 100 + (w.len() as u32))
494                .collect()
495        }
496    }
497
498    #[test]
499    fn py_json_matches_python_spacing() {
500        // Python json.dumps({"a":1,"b":"x"}, ensure_ascii=False) -> '{"a": 1, "b": "x"}'
501        let v = json!({"a": 1, "b": "x"});
502        assert_eq!(py_json(&v), r#"{"a": 1, "b": "x"}"#);
503        assert_eq!(py_json(&json!([1, 2, 3])), "[1, 2, 3]");
504        assert_eq!(py_json(&json!("café")), "\"café\""); // non-ASCII stays raw
505        assert_eq!(py_json(&json!(true)), "true");
506        assert_eq!(py_json(&json!(null)), "null");
507    }
508
509    #[test]
510    fn serialize_state_passes_strings_through() {
511        assert_eq!(serialize_state(&json!("hello")), "hello");
512        assert_eq!(serialize_state(&json!({"k": "v"})), r#"{"k": "v"}"#);
513    }
514
515    #[test]
516    fn answer_confidence_is_max_prob() {
517        // max(p[:k]), clamped to [0, 1]; k < 1 is defined as fully confident.
518        assert_eq!(answer_confidence(&[0.2, 0.7, 0.1], 3), 0.7);
519        assert_eq!(answer_confidence(&[0.6, 0.4], 2), 0.6);
520        assert_eq!(answer_confidence(&[0.9], 1), 0.9);
521        assert_eq!(answer_confidence(&[], 0), 1.0);
522        // Only the first k entries count.
523        assert_eq!(answer_confidence(&[0.3, 0.3, 0.99], 2), 0.3);
524    }
525
526    #[test]
527    fn round4_half_to_even() {
528        assert_eq!(round4(0.12345), 0.1234); // 0.12345 -> 0.1234 (round half to even)
529        assert_eq!(round4(0.12355), 0.1236);
530        assert_eq!(round4(1.0), 1.0);
531    }
532
533    #[test]
534    fn render_options_per_type() {
535        let choice = InternalQ {
536            t: QType::Choice,
537            ins: String::new(),
538            crit: json!({"a": "first", "b": null, "c": ""}),
539            labels: None,
540        };
541        assert_eq!(render_options(&choice).unwrap(), vec!["a: first", "b", "c"]);
542
543        let score = InternalQ {
544            t: QType::Score,
545            ins: String::new(),
546            crit: json!(["low", "high"]),
547            labels: None,
548        };
549        assert_eq!(
550            render_options(&score).unwrap(),
551            vec!["level 0: low", "level 1: high"]
552        );
553
554        let noul = InternalQ {
555            t: QType::Noul,
556            ins: String::new(),
557            crit: json!({"true": "yes it does", "false": "no"}),
558            labels: None,
559        };
560        assert_eq!(
561            render_options(&noul).unwrap(),
562            vec!["false: no", "true: yes it does"]
563        );
564
565        let noul_default = InternalQ {
566            t: QType::Noul,
567            ins: String::new(),
568            crit: Value::Null,
569            labels: None,
570        };
571        assert_eq!(
572            render_options(&noul_default).unwrap(),
573            vec![
574                "false: no, the statement does not hold",
575                "true: yes, the statement holds"
576            ]
577        );
578    }
579
580    #[test]
581    fn temp_bucket_sizes() {
582        assert_eq!(temp_bucket(0, 2), "choice:2");
583        assert_eq!(temp_bucket(1, 4), "score:3-5");
584        assert_eq!(temp_bucket(2, 8), "noul:6-10");
585        assert_eq!(temp_bucket(0, 20), "choice:11+");
586    }
587
588    #[test]
589    fn build_sequence_structure() {
590        let q = InternalQ {
591            t: QType::Choice,
592            ins: "pick one".to_string(),
593            crit: json!({"a": null, "b": null}),
594            labels: None,
595        };
596        let (ids, markers) =
597            build_sequence(&Fake, &json!("some state text"), &q, 512, 192, None, false).unwrap();
598        assert_eq!(ids[0], 1); // [CLS]
599        assert_eq!(*ids.last().unwrap(), 2); // trailing [SEP]
600        assert_eq!(markers.len(), 2); // one marker per option
601                                      // each marker points at a [MASK] id
602        for &m in &markers {
603            assert_eq!(ids[m], 3);
604        }
605    }
606
607    #[test]
608    fn confidence_bounds() {
609        assert_eq!(confidence_from_probs(&[1.0], 1), 1.0); // k<2
610        let uniform = confidence_from_probs(&[0.5, 0.5], 2);
611        assert!(uniform.abs() < 1e-9); // max entropy -> confidence 0
612        let certain = confidence_from_probs(&[1.0, 0.0], 2);
613        assert!((certain - 1.0).abs() < 1e-9);
614    }
615}
laya/common.rs

laya/
common.rs
