Skip to main content

predict/
predict.rs

1//! Smoke test / demo: download the English checkpoint, run a preset, print the answers.
2//!
3//!   cargo run --release --example predict
4//!
5//! Set LAYA_MODEL to a local checkpoint dir or another repo id; LAYA_SUBFOLDER for a bundle
6//! subfolder (e.g. "multilingual").
7
8use laya::agent::{Agent, LoadOptions};
9use laya::{triage_questions, State};
10use std::time::Instant;
11
12fn main() -> Result<(), Box<dyn std::error::Error>> {
13    let model =
14        std::env::var("LAYA_MODEL").unwrap_or_else(|_| "convaiinnovations/laya".to_string());
15    let subfolder = std::env::var("LAYA_SUBFOLDER").ok();
16    let device = std::env::var("LAYA_DEVICE").ok().filter(|s| !s.is_empty());
17
18    let t = Instant::now();
19    let agent = Agent::load(
20        &model,
21        LoadOptions {
22            subfolder,
23            device: device.clone(),
24            ..Default::default()
25        },
26    )?;
27    println!(
28        "loaded {model} on {} in {:?}",
29        device.as_deref().unwrap_or("cpu"),
30        t.elapsed()
31    );
32
33    let state: State = serde_json::json!({
34        "message": "I was charged twice this month and I'm furious. Refund me now or I'm leaving."
35    });
36    let questions = triage_questions();
37
38    let t = Instant::now();
39    let result = agent.system_one(&state, &questions)?;
40    println!("forward in {:?}\n", t.elapsed());
41    println!("{}", serde_json::to_string_pretty(&result.to_json())?);
42
43    // Sanity: every probability distribution should sum to ~1.
44    for (qid, ans) in &result.answers {
45        if let Some(probs) = ans.get("probabilities").and_then(|p| p.as_object()) {
46            let sum: f64 = probs.values().filter_map(|v| v.as_f64()).sum();
47            assert!((sum - 1.0).abs() < 1e-2, "{qid} probabilities sum to {sum}");
48        }
49    }
50    println!("\nOK: probability distributions sum to 1.");
51    Ok(())
52}