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/// Normalized Shannon-entropy confidence: `1 - H(p)/log(k)`, clamped to `[0, 1]`.
348/// `p` is expected to already hold the `k` option probabilities.
349pub fn confidence_from_probs(p: &[f64], k: usize) -> f64 {
350    if k < 2 {
351        return 1.0;
352    }
353    let p = &p[..k.min(p.len())];
354    let ent: f64 = -p.iter().map(|&v| v * v.clamp(1e-12, 1.0).ln()).sum::<f64>();
355    (1.0 - ent / (k as f64).ln()).clamp(0.0, 1.0)
356}
357
358pub const TEMP_MIN: f64 = 0.5;
359pub const TEMP_MAX: f64 = 5.0;
360
361/// A usable temperature: `t` confined to `[TEMP_MIN, TEMP_MAX]`, falling back to 1.0 if not finite.
362pub fn clamp_temperature(t: &Value) -> f64 {
363    let f = match t {
364        Value::Number(n) => n.as_f64(),
365        Value::String(s) => s.trim().parse::<f64>().ok(),
366        _ => None,
367    };
368    match f {
369        Some(x) if x.is_finite() => x.clamp(TEMP_MIN, TEMP_MAX),
370        _ => 1.0,
371    }
372}
373
374/// Temperature bucket key, e.g. `"choice:3-5"`.
375pub fn temp_bucket(qtype: usize, k: usize) -> String {
376    let size = if k <= 2 {
377        "2"
378    } else if k <= 5 {
379        "3-5"
380    } else if k <= 10 {
381        "6-10"
382    } else {
383        "11+"
384    };
385    let name = QType::from_index(qtype)
386        .map(|q| q.name())
387        .unwrap_or("choice");
388    format!("{}:{}", name, size)
389}
390
391/// One tokenized (state, question) pair before collation.
392#[derive(Debug, Clone)]
393pub struct Item {
394    pub ids: Vec<u32>,
395    pub markers: Vec<usize>,
396    pub qtype: usize,
397}
398
399/// A padded batch ready to be turned into tensors.
400#[derive(Debug, Clone)]
401pub struct CollatedBatch {
402    pub input_ids: Vec<Vec<u32>>,
403    pub attention_mask: Vec<Vec<u32>>,
404    pub marker_pos: Vec<Vec<usize>>,
405    pub marker_mask: Vec<Vec<bool>>,
406    pub qtype: Vec<usize>,
407}
408
409/// Collate groups of items into one padded batch. `groups` is a list of per-state item lists,
410/// flattened in order (mirrors `collate_items` in `laya/common.py`). Returns `None` if empty.
411pub fn collate_items(groups: &[Vec<Item>], _pad_id: u32) -> Option<CollatedBatch> {
412    let items: Vec<&Item> = groups.iter().flatten().collect();
413    if items.is_empty() {
414        return None;
415    }
416    let l = items.iter().map(|it| it.ids.len()).max().unwrap();
417    let kmax = items.iter().map(|it| it.markers.len()).max().unwrap();
418
419    let mut input_ids = Vec::with_capacity(items.len());
420    let mut attention_mask = Vec::with_capacity(items.len());
421    let mut marker_pos = Vec::with_capacity(items.len());
422    let mut marker_mask = Vec::with_capacity(items.len());
423    let mut qtype = Vec::with_capacity(items.len());
424
425    for it in &items {
426        let mut ids = it.ids.clone();
427        let mut att = vec![1u32; ids.len()];
428        ids.resize(l, _pad_id);
429        att.resize(l, 0);
430        input_ids.push(ids);
431        attention_mask.push(att);
432
433        let k = it.markers.len();
434        let mut mpos: Vec<usize> = it.markers.clone();
435        mpos.resize(kmax, 0);
436        let mut mmask: Vec<bool> = vec![true; k];
437        mmask.resize(kmax, false);
438        marker_pos.push(mpos);
439        marker_mask.push(mmask);
440
441        qtype.push(it.qtype);
442    }
443
444    Some(CollatedBatch {
445        input_ids,
446        attention_mask,
447        marker_pos,
448        marker_mask,
449        qtype,
450    })
451}
452
453#[cfg(test)]
454mod tests {
455    use super::*;
456    use serde_json::json;
457
458    /// A deterministic fake tokenizer for structural tests (one id per whitespace word).
459    struct Fake;
460    impl Tokenizer for Fake {
461        fn cls_id(&self) -> u32 {
462            1
463        }
464        fn sep_id(&self) -> u32 {
465            2
466        }
467        fn mask_id(&self) -> u32 {
468            3
469        }
470        fn pad_id(&self) -> u32 {
471            0
472        }
473        fn mask_token(&self) -> &str {
474            "[MASK]"
475        }
476        fn encode(&self, text: &str) -> Vec<u32> {
477            text.split_whitespace()
478                .map(|w| 100 + (w.len() as u32))
479                .collect()
480        }
481    }
482
483    #[test]
484    fn py_json_matches_python_spacing() {
485        // Python json.dumps({"a":1,"b":"x"}, ensure_ascii=False) -> '{"a": 1, "b": "x"}'
486        let v = json!({"a": 1, "b": "x"});
487        assert_eq!(py_json(&v), r#"{"a": 1, "b": "x"}"#);
488        assert_eq!(py_json(&json!([1, 2, 3])), "[1, 2, 3]");
489        assert_eq!(py_json(&json!("café")), "\"café\""); // non-ASCII stays raw
490        assert_eq!(py_json(&json!(true)), "true");
491        assert_eq!(py_json(&json!(null)), "null");
492    }
493
494    #[test]
495    fn serialize_state_passes_strings_through() {
496        assert_eq!(serialize_state(&json!("hello")), "hello");
497        assert_eq!(serialize_state(&json!({"k": "v"})), r#"{"k": "v"}"#);
498    }
499
500    #[test]
501    fn round4_half_to_even() {
502        assert_eq!(round4(0.12345), 0.1234); // 0.12345 -> 0.1234 (round half to even)
503        assert_eq!(round4(0.12355), 0.1236);
504        assert_eq!(round4(1.0), 1.0);
505    }
506
507    #[test]
508    fn render_options_per_type() {
509        let choice = InternalQ {
510            t: QType::Choice,
511            ins: String::new(),
512            crit: json!({"a": "first", "b": null, "c": ""}),
513            labels: None,
514        };
515        assert_eq!(render_options(&choice).unwrap(), vec!["a: first", "b", "c"]);
516
517        let score = InternalQ {
518            t: QType::Score,
519            ins: String::new(),
520            crit: json!(["low", "high"]),
521            labels: None,
522        };
523        assert_eq!(
524            render_options(&score).unwrap(),
525            vec!["level 0: low", "level 1: high"]
526        );
527
528        let noul = InternalQ {
529            t: QType::Noul,
530            ins: String::new(),
531            crit: json!({"true": "yes it does", "false": "no"}),
532            labels: None,
533        };
534        assert_eq!(
535            render_options(&noul).unwrap(),
536            vec!["false: no", "true: yes it does"]
537        );
538
539        let noul_default = InternalQ {
540            t: QType::Noul,
541            ins: String::new(),
542            crit: Value::Null,
543            labels: None,
544        };
545        assert_eq!(
546            render_options(&noul_default).unwrap(),
547            vec![
548                "false: no, the statement does not hold",
549                "true: yes, the statement holds"
550            ]
551        );
552    }
553
554    #[test]
555    fn temp_bucket_sizes() {
556        assert_eq!(temp_bucket(0, 2), "choice:2");
557        assert_eq!(temp_bucket(1, 4), "score:3-5");
558        assert_eq!(temp_bucket(2, 8), "noul:6-10");
559        assert_eq!(temp_bucket(0, 20), "choice:11+");
560    }
561
562    #[test]
563    fn build_sequence_structure() {
564        let q = InternalQ {
565            t: QType::Choice,
566            ins: "pick one".to_string(),
567            crit: json!({"a": null, "b": null}),
568            labels: None,
569        };
570        let (ids, markers) =
571            build_sequence(&Fake, &json!("some state text"), &q, 512, 192, None, false).unwrap();
572        assert_eq!(ids[0], 1); // [CLS]
573        assert_eq!(*ids.last().unwrap(), 2); // trailing [SEP]
574        assert_eq!(markers.len(), 2); // one marker per option
575                                      // each marker points at a [MASK] id
576        for &m in &markers {
577            assert_eq!(ids[m], 3);
578        }
579    }
580
581    #[test]
582    fn confidence_bounds() {
583        assert_eq!(confidence_from_probs(&[1.0], 1), 1.0); // k<2
584        let uniform = confidence_from_probs(&[0.5, 0.5], 2);
585        assert!(uniform.abs() < 1e-9); // max entropy -> confidence 0
586        let certain = confidence_from_probs(&[1.0, 0.0], 2);
587        assert!((certain - 1.0).abs() < 1e-9);
588    }
589}
laya/common.rs

laya/
common.rs
