Skip to main content

laya/
agent.rs

1//! High-level inference runtime. Ported from `laya/agent.py` (load, validate/normalize
2//! questions, tokenize, forward, decode typed answers) — minus the TileLang fast path.
3
4use std::path::{Path, PathBuf};
5
6use candle_core::Device;
7use indexmap::IndexMap;
8use serde_json::{json, Map, Value};
9
10use crate::common::{
11    answer_confidence, build_sequence, clamp_temperature, collate_items, confidence_from_probs,
12    render_options, round4, softmax, temp_bucket, InternalQ, Item, QType,
13};
14use crate::error::{LayaError, Result};
15use crate::model::DecisionModel;
16use crate::tokenizer::LayaTokenizer;
17use crate::{Questions, State};
18
19/// Options for [`Agent::load`].
20#[derive(Debug, Clone, Default)]
21pub struct LoadOptions {
22    /// Compute device: `None`/`"cpu"` for CPU (the only backend built by default).
23    pub device: Option<String>,
24    /// Hugging Face token (falls back to `HF_TOKEN`).
25    pub token: Option<String>,
26    /// Subfolder within a bundled repo (e.g. `"multilingual"`).
27    pub subfolder: Option<String>,
28}
29
30/// The result of a single-state evaluation (Jev-shaped).
31#[derive(Debug, Clone)]
32pub struct SystemOneResult {
33    pub model: String,
34    pub answers: IndexMap<String, Value>,
35    pub input_tokens: usize,
36    pub output_tokens: usize,
37    /// Set by [`crate::Router::predict`]; the routing decision as JSON.
38    pub routing: Option<Value>,
39}
40
41impl SystemOneResult {
42    /// Serialise to the Jev wire shape: `{model, answers, usage, routing?}`.
43    pub fn to_json(&self) -> Value {
44        let answers: Map<String, Value> = self
45            .answers
46            .iter()
47            .map(|(k, v)| (k.clone(), v.clone()))
48            .collect();
49        let mut out = json!({
50            "model": self.model,
51            "answers": Value::Object(answers),
52            "usage": {"input_tokens": self.input_tokens, "output_tokens": self.output_tokens},
53        });
54        if let Some(r) = &self.routing {
55            out.as_object_mut()
56                .unwrap()
57                .insert("routing".to_string(), r.clone());
58        }
59        out
60    }
61}
62
63struct CheckpointFiles {
64    config: PathBuf,
65    weights: PathBuf,
66    encoder_config: PathBuf,
67    tokenizer: PathBuf,
68}
69
70/// System-1 decision model runtime.
71pub struct Agent {
72    model: DecisionModel,
73    tok: LayaTokenizer,
74    max_len: usize,
75    head_max_len: usize,
76    temperature: [f64; 3],
77    temperature_by_options: std::collections::HashMap<String, f64>,
78}
79
80impl Agent {
81    /// Load a Laya checkpoint from a local directory or a Hugging Face repo id.
82    pub fn load(model_id_or_path: &str, opts: LoadOptions) -> Result<Agent> {
83        let files = resolve_files(model_id_or_path, &opts)?;
84
85        let cfg: Value = serde_json::from_slice(&std::fs::read(&files.config)?)?;
86        let encoder_config: Value = serde_json::from_slice(&std::fs::read(&files.encoder_config)?)?;
87
88        let max_len = cfg.get("max_len").and_then(|v| v.as_u64()).unwrap_or(512) as usize;
89        let head_max_len = cfg
90            .get("head_max_len")
91            .and_then(|v| v.as_u64())
92            .unwrap_or(192) as usize;
93        let head_layers = cfg.get("head_layers").and_then(|v| v.as_u64()).unwrap_or(2) as usize;
94        let n_act = cfg
95            .get("act_costs")
96            .and_then(|v| v.as_object())
97            .map(|m| m.len())
98            .unwrap_or(0)
99            + 1;
100
101        let temperature = load_temperature(&cfg);
102        let temperature_by_options = load_temperature_by_options(&cfg);
103
104        let device = resolve_device(opts.device.as_deref());
105        let model =
106            DecisionModel::load(&files.weights, &encoder_config, head_layers, n_act, device)?;
107        let tok = LayaTokenizer::from_file(&files.tokenizer)?;
108
109        Ok(Agent {
110            model,
111            tok,
112            max_len,
113            head_max_len,
114            temperature,
115            temperature_by_options,
116        })
117    }
118
119    /// Evaluate typed questions over one state in a single forward pass.
120    pub fn system_one(&self, state: &State, questions: &Questions) -> Result<SystemOneResult> {
121        Ok(self
122            .predict_batch(std::slice::from_ref(state), questions, None)?
123            .pop()
124            .expect("one state → one result"))
125    }
126
127    /// Alias for [`Agent::system_one`].
128    pub fn predict(&self, state: &State, questions: &Questions) -> Result<SystemOneResult> {
129        self.system_one(state, questions)
130    }
131
132    /// Evaluate the same questions over many states, sharing forward passes.
133    pub fn predict_batch(
134        &self,
135        states: &[State],
136        questions: &Questions,
137        batch_size: Option<usize>,
138    ) -> Result<Vec<SystemOneResult>> {
139        if states.is_empty() {
140            return Ok(Vec::new());
141        }
142        let ids: Vec<String> = questions.keys().cloned().collect();
143        if ids.is_empty() {
144            return Ok(states
145                .iter()
146                .map(|_| SystemOneResult {
147                    model: "laya-rl-agent".to_string(),
148                    answers: IndexMap::new(),
149                    input_tokens: 0,
150                    output_tokens: 0,
151                    routing: None,
152                })
153                .collect());
154        }
155
156        // Validate + normalize each question once (state-independent).
157        let mut internal: Vec<(String, InternalQ)> = Vec::with_capacity(ids.len());
158        for qid in &ids {
159            let qdef = &questions[qid];
160            check_question(qid, qdef)?;
161            internal.push((qid.clone(), to_internal(qdef)?));
162        }
163
164        let chunk = batch_size.filter(|&b| b > 0).unwrap_or(states.len());
165        let mut results = Vec::with_capacity(states.len());
166
167        for part in states.chunks(chunk) {
168            let mut per_state_items: Vec<Vec<Item>> = Vec::with_capacity(part.len());
169            for st in part {
170                per_state_items.push(self.encode_state(st, &internal)?);
171            }
172            let batch = collate_items(&per_state_items, self.tok_pad())
173                .ok_or_else(|| LayaError::Model("empty batch".to_string()))?;
174            let (logits, act) = self.model.forward(&batch)?;
175
176            let mut row = 0usize;
177            for items in &per_state_items {
178                let nrows = items.len();
179                let n_tokens: usize = (row..row + nrows)
180                    .map(|r| {
181                        batch.attention_mask[r]
182                            .iter()
183                            .map(|&x| x as usize)
184                            .sum::<usize>()
185                    })
186                    .sum();
187                let answers = self.decode(&logits, &act, items, &internal, row);
188                results.push(SystemOneResult {
189                    model: "laya-rl-agent".to_string(),
190                    answers,
191                    input_tokens: n_tokens,
192                    output_tokens: 0,
193                    routing: None,
194                });
195                row += nrows;
196            }
197        }
198        Ok(results)
199    }
200
201    fn tok_pad(&self) -> u32 {
202        use crate::common::Tokenizer;
203        self.tok.pad_id()
204    }
205
206    /// Mean-pooled encoder embeddings for a list of texts (`[n, hidden_size]`).
207    ///
208    /// Tokenizes with special tokens, truncates to `max_length`, runs the checkpoint encoder in
209    /// batches of `batch_size`, and mean-pools over non-padding tokens. Backs
210    /// [`crate::shortlist::embed_fn_from_agent`]; a dedicated bi-encoder will usually shortlist
211    /// better, but this needs only the loaded checkpoint.
212    pub fn embed(
213        &self,
214        texts: &[String],
215        max_length: usize,
216        batch_size: usize,
217    ) -> Result<Vec<Vec<f32>>> {
218        if texts.is_empty() {
219            return Ok(Vec::new());
220        }
221        let pad = self.tok_pad();
222        let mut out = Vec::with_capacity(texts.len());
223        for chunk in texts.chunks(batch_size.max(1)) {
224            let rows: Vec<Vec<u32>> = chunk
225                .iter()
226                .map(|t| self.tok.encode_with_special(t, max_length))
227                .collect();
228            let l = rows.iter().map(|r| r.len()).max().unwrap_or(0).max(1);
229            let mut ids = Vec::with_capacity(rows.len());
230            let mut att = Vec::with_capacity(rows.len());
231            for r in &rows {
232                let mut row = r.clone();
233                let mut a = vec![1u32; row.len()];
234                row.resize(l, pad);
235                a.resize(l, 0);
236                ids.push(row);
237                att.push(a);
238            }
239            out.extend(self.model.embed(&ids, &att)?);
240        }
241        Ok(out)
242    }
243
244    fn encode_state(&self, state: &State, internal: &[(String, InternalQ)]) -> Result<Vec<Item>> {
245        let truncate_left = state.is_array();
246        let mut items = Vec::with_capacity(internal.len());
247        for (qid, q) in internal {
248            let (seq, markers) = build_sequence(
249                &self.tok,
250                state,
251                q,
252                self.max_len,
253                self.head_max_len,
254                None,
255                truncate_left,
256            )?;
257            if markers.len() != render_options(q)?.len() {
258                return Err(LayaError::InvalidQuestion(format!(
259                    "question {:?} options exceed head_max_len={}",
260                    qid, self.head_max_len
261                )));
262            }
263            items.push(Item {
264                ids: seq,
265                markers,
266                qtype: q.t.index(),
267            });
268        }
269        Ok(items)
270    }
271
272    fn decode(
273        &self,
274        logits: &[Vec<f32>],
275        act: &[Vec<f32>],
276        items: &[Item],
277        internal: &[(String, InternalQ)],
278        offset: usize,
279    ) -> IndexMap<String, Value> {
280        let mut answers = IndexMap::new();
281        for (j, (qid, q)) in internal.iter().enumerate() {
282            let r = offset + j;
283            let k = items[j].markers.len();
284            let qt = q.t.index();
285            let t_scale = self
286                .temperature_by_options
287                .get(&temp_bucket(qt, k))
288                .copied()
289                .unwrap_or(self.temperature[qt]);
290
291            let z: Vec<f64> = logits[r][..k].iter().map(|&v| v as f64 / t_scale).collect();
292            let p = softmax(&z);
293            let conf = round4(confidence_from_probs(&p, k));
294            // `confidence` means one thing for `noul` (max(p)) and another for `choice`/`score`
295            // (normalized entropy), and only the first is the quantity temperature scaling fits and
296            // ECE measures. Report both: `answer_confidence` is the calibrated one on every type, so
297            // a caller can gate across question types on a single number.
298            let ans_conf = round4(answer_confidence(&p, k));
299
300            let act_p = softmax(&act[r].iter().map(|&v| v as f64).collect::<Vec<_>>());
301            let act_probability = round4(act_p.first().copied().unwrap_or(0.0));
302            let action = json!({ "act_probability": act_probability });
303
304            let answer = match q.t {
305                QType::Choice => {
306                    let keys: Vec<String> = q
307                        .crit
308                        .as_object()
309                        .map(|m| m.keys().cloned().collect())
310                        .unwrap_or_default();
311                    let best = argmax(&p);
312                    let probs: Map<String, Value> = keys
313                        .iter()
314                        .zip(p.iter())
315                        .map(|(kk, &v)| (kk.clone(), json!(round4(v))))
316                        .collect();
317                    json!({
318                        "type": "choice",
319                        "choice": keys.get(best).cloned().unwrap_or_default(),
320                        "probabilities": Value::Object(probs),
321                        "confidence": conf,
322                        "answer_confidence": ans_conf,
323                        "action": action,
324                    })
325                }
326                QType::Score => {
327                    let exp: f64 = p.iter().enumerate().map(|(i, &v)| i as f64 * v).sum();
328                    let legend: Map<String, Value> = q
329                        .crit
330                        .as_array()
331                        .map(|a| {
332                            a.iter()
333                                .enumerate()
334                                .map(|(i, c)| (i.to_string(), c.clone()))
335                                .collect()
336                        })
337                        .unwrap_or_default();
338                    let probs: Map<String, Value> = p
339                        .iter()
340                        .enumerate()
341                        .map(|(i, &v)| (i.to_string(), json!(round4(v))))
342                        .collect();
343                    json!({
344                        "type": "score",
345                        "score": round4(exp),
346                        "legend": Value::Object(legend),
347                        "probabilities": Value::Object(probs),
348                        "confidence": conf,
349                        "answer_confidence": ans_conf,
350                        "action": action,
351                    })
352                }
353                QType::Noul => {
354                    let p_true = p.get(1).copied().unwrap_or(0.0);
355                    json!({
356                        "type": "noul",
357                        "noul": round4(p_true),
358                        "confidence": round4(p_true.max(1.0 - p_true)),
359                        // identical here: over two options max(p_true, 1 - p_true) is max(p)
360                        "answer_confidence": ans_conf,
361                        "action": action,
362                    })
363                }
364            };
365            answers.insert(qid.clone(), answer);
366        }
367        answers
368    }
369}
370
371fn argmax(p: &[f64]) -> usize {
372    let mut best = 0usize;
373    for i in 1..p.len() {
374        if p[i] > p[best] {
375            best = i;
376        }
377    }
378    best
379}
380
381fn load_temperature(cfg: &Value) -> [f64; 3] {
382    let arr = cfg.get("temperature").and_then(|v| v.as_array());
383    let mut out = [1.0; 3];
384    if let Some(a) = arr {
385        for (i, slot) in out.iter_mut().enumerate() {
386            if let Some(v) = a.get(i) {
387                *slot = clamp_temperature(v);
388            }
389        }
390    }
391    out
392}
393
394fn load_temperature_by_options(cfg: &Value) -> std::collections::HashMap<String, f64> {
395    cfg.get("temperature_by_options")
396        .and_then(|v| v.as_object())
397        .map(|m| {
398            m.iter()
399                .map(|(k, v)| (k.clone(), clamp_temperature(v)))
400                .collect()
401        })
402        .unwrap_or_default()
403}
404
405fn resolve_device(device: Option<&str>) -> Device {
406    match device.map(|d| d.trim().to_lowercase()).as_deref() {
407        Some("metal") => {
408            #[cfg(feature = "metal")]
409            {
410                match Device::new_metal(0) {
411                    Ok(d) => d,
412                    Err(e) => {
413                        eprintln!(
414                            "laya: Metal requested but unavailable ({e}); falling back to CPU."
415                        );
416                        Device::Cpu
417                    }
418                }
419            }
420            #[cfg(not(feature = "metal"))]
421            {
422                eprintln!(
423                    "laya: device=\"metal\" requested but the crate was built without the `metal` \
424                     feature; using CPU. Rebuild with --features metal."
425                );
426                Device::Cpu
427            }
428        }
429        // "cuda" and other accelerators are follow-ups; everything else runs on CPU.
430        _ => Device::Cpu,
431    }
432}
433
434/// Validate a question definition, matching `Agent._check_question`.
435fn check_question(qid: &str, qdef: &Value) -> Result<()> {
436    let obj = match qdef {
437        Value::Object(m) => m,
438        other => {
439            let kind = match other {
440                Value::Array(_) => "list",
441                Value::Null => "NoneType",
442                Value::String(_) => "str",
443                Value::Bool(_) => "bool",
444                Value::Number(_) => "number",
445                Value::Object(_) => unreachable!(),
446            };
447            return Err(LayaError::InvalidQuestion(format!(
448                "question {:?}: definition must be a dict, got {}",
449                qid, kind
450            )));
451        }
452    };
453    let t = obj.get("type").and_then(|v| v.as_str());
454    let t = match t.and_then(QType::from_name) {
455        Some(t) => t,
456        None => {
457            return Err(LayaError::InvalidQuestion(format!(
458                "question {:?}: unknown type {:?}; use one of [\"choice\", \"noul\", \"score\"]",
459                qid,
460                obj.get("type").unwrap_or(&Value::Null)
461            )))
462        }
463    };
464    if !obj.contains_key("instructions") {
465        return Err(LayaError::InvalidQuestion(format!(
466            "question {:?}: no 'instructions'; add the text the model should answer",
467            qid
468        )));
469    }
470    let crit = obj.get("criteria");
471    match t {
472        QType::Choice => {
473            let ok = matches!(crit, Some(Value::Object(_)) | Some(Value::Array(_)));
474            if !ok {
475                return Err(LayaError::InvalidQuestion(format!(
476                    "question {:?}: a choice question takes 'criteria' as a dict of label -> description, or a list of labels",
477                    qid
478                )));
479            }
480            let empty = match crit {
481                Some(Value::Object(m)) => m.is_empty(),
482                Some(Value::Array(a)) => a.is_empty(),
483                _ => true,
484            };
485            if empty {
486                return Err(LayaError::InvalidQuestion(format!(
487                    "question {:?}: a choice question needs at least one criterion",
488                    qid
489                )));
490            }
491            // A label in the list form is used as the answer key and rendered as option text, so a
492            // nested list or object label has no meaning here. In the dict form keys are always
493            // strings, so only the list form can carry one. Reject it, naming the question and the
494            // label, rather than letting it be JSON-stringified into a key (upstream #425). The dict
495            // form (JSON object) cannot hit this: its keys are strings by construction.
496            if let Some(Value::Array(a)) = crit {
497                for (i, label) in a.iter().enumerate() {
498                    let kind = match label {
499                        Value::Array(_) => "list",
500                        Value::Object(_) => "object",
501                        _ => continue,
502                    };
503                    return Err(LayaError::InvalidQuestion(format!(
504                        "question {:?}: choice label {} is a {}; a label is rendered as option text \
505                         and used as the answer key, so it must be a scalar (a string, number or \
506                         null), got {}",
507                        qid,
508                        i,
509                        kind,
510                        crate::common::py_json(label)
511                    )));
512                }
513            }
514        }
515        QType::Score => {
516            let arr = match crit {
517                Some(Value::Array(a)) => a,
518                _ => {
519                    return Err(LayaError::InvalidQuestion(format!(
520                        "question {:?}: a score question takes 'criteria' as a list of level descriptions, index 0 first",
521                        qid
522                    )))
523                }
524            };
525            if arr.is_empty() {
526                return Err(LayaError::InvalidQuestion(format!(
527                    "question {:?}: a score question needs at least one level",
528                    qid
529                )));
530            }
531            // A null level would silently drop that description; reject it and name the index so
532            // every level gets described, index 0 first (upstream: reject a null score level).
533            if let Some(idx) = arr.iter().position(|v| v.is_null()) {
534                return Err(LayaError::InvalidQuestion(format!(
535                    "question {:?}: score level {} is null; give every level a description, index 0 first",
536                    qid, idx
537                )));
538            }
539        }
540        QType::Noul => {
541            if let Some(c) = crit {
542                if !c.is_null() && !c.is_object() {
543                    return Err(LayaError::InvalidQuestion(format!(
544                        "question {:?}: a noul question takes 'criteria' as a dict with optional 'true'/'false' descriptions, or omits it",
545                        qid
546                    )));
547                }
548                // `render_options` reads only `false`/`true` out by name, so a dict keyed any other
549                // way is not a noul description at all — it used to be silently dropped and replaced
550                // with the defaults (#156). A noul is a boolean question, so those are its only keys.
551                if let Some(map) = c.as_object() {
552                    let mut keys: Vec<String> = map.keys().map(|k| k.to_lowercase()).collect();
553                    if keys.iter().any(|k| k != "true" && k != "false") {
554                        keys.sort();
555                        let shown: Vec<Value> = keys.into_iter().map(Value::String).collect();
556                        return Err(LayaError::InvalidQuestion(format!(
557                            "question {:?}: a noul question takes 'criteria' keyed only 'true'/'false' (either or both, and omitted is fine), got {}. Those keys are the option texts the model reads; any other key was silently dropped and replaced with the defaults. If you want the answer worded differently, keep 'criteria' keyed 'true'/'false' and set 'labels' instead.",
558                            qid,
559                            Value::Array(shown)
560                        )));
561                    }
562                }
563            }
564        }
565    }
566    if obj.contains_key("labels") {
567        if t != QType::Noul {
568            return Err(LayaError::InvalidQuestion(format!(
569                "question {:?}: 'labels' is only supported for noul questions",
570                qid
571            )));
572        }
573        crate::common::resolve_noul_labels(obj.get("labels"))
574            .map_err(|e| LayaError::InvalidQuestion(format!("question {:?}: {}", qid, e)))?;
575    }
576    Ok(())
577}
578
579/// Normalize a validated question into the internal form.
580fn to_internal(qdef: &Value) -> Result<InternalQ> {
581    let obj = qdef.as_object().expect("validated as object");
582    let t = QType::from_name(obj.get("type").and_then(|v| v.as_str()).unwrap_or("")).unwrap();
583    let mut crit = obj.get("criteria").cloned().unwrap_or(Value::Null);
584    match t {
585        QType::Choice => {
586            if let Value::Array(a) = &crit {
587                let mut m = Map::new();
588                for c in a {
589                    let key = c
590                        .as_str()
591                        .map(|s| s.to_string())
592                        .unwrap_or_else(|| crate::common::py_json(c));
593                    m.insert(key, Value::Null);
594                }
595                crit = Value::Object(m);
596            }
597        }
598        QType::Noul => {
599            if let Value::Object(m) = &crit {
600                let lowered: Map<String, Value> = m
601                    .iter()
602                    .map(|(k, v)| (k.to_lowercase(), v.clone()))
603                    .collect();
604                crit = Value::Object(lowered);
605            }
606        }
607        QType::Score => {}
608    }
609    let ins = match obj.get("instructions") {
610        Some(Value::String(s)) => s.clone(),
611        Some(other) => crate::common::py_json(other),
612        None => String::new(),
613    };
614    let labels = obj.get("labels").cloned();
615    Ok(InternalQ {
616        t,
617        ins,
618        crit,
619        labels,
620    })
621}
622
623/// Resolve the four checkpoint files from a local dir or a Hugging Face repo id.
624fn resolve_files(model_id_or_path: &str, opts: &LoadOptions) -> Result<CheckpointFiles> {
625    let base = Path::new(model_id_or_path);
626    if base.exists() {
627        let dir = match &opts.subfolder {
628            Some(sub) => base.join(sub),
629            None => base.to_path_buf(),
630        };
631        let config = require(dir.join("rl_agent_config.json"))?;
632        let weights = require(dir.join("model.safetensors"))?;
633        let encoder_config = if dir.join("encoder/config.json").exists() {
634            dir.join("encoder/config.json")
635        } else {
636            require(dir.join("config.json"))?
637        };
638        let tokenizer = if dir.join("tokenizer/tokenizer.json").exists() {
639            dir.join("tokenizer/tokenizer.json")
640        } else {
641            require(dir.join("tokenizer.json"))?
642        };
643        return Ok(CheckpointFiles {
644            config,
645            weights,
646            encoder_config,
647            tokenizer,
648        });
649    }
650    if model_id_or_path.starts_with('/')
651        || model_id_or_path.starts_with("./")
652        || model_id_or_path.starts_with("../")
653    {
654        return Err(LayaError::NotFound(format!(
655            "Local model path not found: {:?}.",
656            model_id_or_path
657        )));
658    }
659    download_files(model_id_or_path, opts)
660}
661
662fn require(p: PathBuf) -> Result<PathBuf> {
663    if p.exists() {
664        Ok(p)
665    } else {
666        Err(LayaError::NotFound(format!(
667            "missing checkpoint file: {}",
668            p.display()
669        )))
670    }
671}
672
673fn download_files(repo: &str, opts: &LoadOptions) -> Result<CheckpointFiles> {
674    use hf_hub::api::sync::ApiBuilder;
675    // An empty token (explicit or from an empty `HF_TOKEN`) is treated as no token, so we never
676    // send an empty `Bearer` header to the hub (parity with upstream `token or HF_TOKEN or None`).
677    let token = opts
678        .token
679        .clone()
680        .or_else(|| std::env::var("HF_TOKEN").ok())
681        .filter(|t| !t.is_empty());
682    let api = ApiBuilder::new()
683        .with_token(token)
684        .build()
685        .map_err(|e| LayaError::Download(e.to_string()))?;
686    let repo_api = api.model(repo.to_string());
687    let prefix = match &opts.subfolder {
688        Some(sub) => format!("{}/", sub),
689        None => String::new(),
690    };
691    let get = |name: &str| -> Result<PathBuf> {
692        repo_api
693            .get(&format!("{}{}", prefix, name))
694            .map_err(|e| LayaError::Download(format!("{}{}: {}", prefix, name, e)))
695    };
696    let config = get("rl_agent_config.json")?;
697    let weights = get("model.safetensors")?;
698    let encoder_config = get("encoder/config.json")?;
699    let tokenizer = get("tokenizer/tokenizer.json").or_else(|_| get("tokenizer.json"))?;
700    Ok(CheckpointFiles {
701        config,
702        weights,
703        encoder_config,
704        tokenizer,
705    })
706}
707
708#[cfg(test)]
709mod tests {
710    use super::*;
711
712    #[test]
713    fn rejects_null_score_level() {
714        // A null level would silently drop that description; it must be rejected, naming the index.
715        let q = json!({"type": "score", "instructions": "rate severity",
716                       "criteria": ["low", null, "high"]});
717        let msg = check_question("severity", &q).unwrap_err().to_string();
718        assert!(msg.contains("score level 1 is null"), "got: {msg}");
719    }
720
721    #[test]
722    fn accepts_fully_described_score_levels() {
723        let q = json!({"type": "score", "instructions": "rate", "criteria": ["low", "high"]});
724        assert!(check_question("s", &q).is_ok());
725    }
726
727    #[test]
728    fn rejects_nested_choice_label() {
729        // A list/object label in the list form is used as the answer key: reject it as a named
730        // caller error, naming the question and the label index (upstream #425).
731        let q = json!({"type": "choice", "instructions": "route",
732                       "criteria": [["billing"], "tech"]});
733        let msg = check_question("dept", &q).unwrap_err().to_string();
734        assert!(msg.contains("choice label 0 is a list"), "got: {msg}");
735
736        let q = json!({"type": "choice", "instructions": "route",
737                       "criteria": ["billing", {"tech": "x"}]});
738        let msg = check_question("dept", &q).unwrap_err().to_string();
739        assert!(msg.contains("choice label 1 is a object"), "got: {msg}");
740    }
741
742    #[test]
743    fn accepts_scalar_choice_labels() {
744        // Strings, numbers and null are all valid answer keys.
745        let q = json!({"type": "choice", "instructions": "route",
746                       "criteria": ["billing", 2, null]});
747        assert!(check_question("dept", &q).is_ok());
748        let q = json!({"type": "choice", "instructions": "route",
749                       "criteria": {"billing": "money", "tech": "bugs"}});
750        assert!(check_question("dept", &q).is_ok());
751    }
752}