Struct Agent 

pub struct Agent { /* private fields */ }

System-1 decision model runtime.

Implementations

impl Agent

pub fn load(model_id_or_path: &str, opts: LoadOptions) -> Result<Agent>

Load a Laya checkpoint from a local directory, a Hugging Face repo id, or a registry checkpoint name or alias (the same ones Router resolves: "typed-decisions", "ml", …).

A bare name or alias resolves to the same (repo, subfolder) the Router would pick, instead of being handed to the Hub as a repo id. A caller who spelled out a subfolder is loading that subfolder of whatever repo they named, so the name is only resolved when it is the whole argument, and only when it does not already look like a local path or name a directory that holds a checkpoint (upstream load name resolution + _is_local_checkpoint_arg).

Examples found in repository

examples/predict.rs (lines 19-26)

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let model =
        std::env::var("LAYA_MODEL").unwrap_or_else(|_| "convaiinnovations/laya".to_string());
    let subfolder = std::env::var("LAYA_SUBFOLDER").ok();
    let device = std::env::var("LAYA_DEVICE").ok().filter(|s| !s.is_empty());

    let t = Instant::now();
    let agent = Agent::load(
        &model,
        LoadOptions {
            subfolder,
            device: device.clone(),
            ..Default::default()
        },
    )?;
    println!(
        "loaded {model} on {} in {:?}",
        device.as_deref().unwrap_or("cpu"),
        t.elapsed()
    );

    let state: State = serde_json::json!({
        "message": "I was charged twice this month and I'm furious. Refund me now or I'm leaving."
    });
    let questions = triage_questions();

    let t = Instant::now();
    let result = agent.system_one(&state, &questions)?;
    println!("forward in {:?}\n", t.elapsed());
    println!("{}", serde_json::to_string_pretty(&result.to_json())?);

    // Sanity: every probability distribution should sum to ~1.
    for (qid, ans) in &result.answers {
        if let Some(probs) = ans.get("probabilities").and_then(|p| p.as_object()) {
            let sum: f64 = probs.values().filter_map(|v| v.as_f64()).sum();
            assert!((sum - 1.0).abs() < 1e-2, "{qid} probabilities sum to {sum}");
        }
    }
    println!("\nOK: probability distributions sum to 1.");
    Ok(())
}

examples/shortlist.rs (lines 11-17)

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let model =
        std::env::var("LAYA_MODEL").unwrap_or_else(|_| "convaiinnovations/laya".to_string());
    let device = std::env::var("LAYA_DEVICE").ok().filter(|s| !s.is_empty());
    let agent = Agent::load(
        &model,
        LoadOptions {
            device,
            ..Default::default()
        },
    )?;

    // A large label set (25 intents); shortlist to the 5 most relevant before scoring.
    let criteria = json!({
        "refund": null, "cancel_subscription": null, "reset_password": null, "update_billing": null,
        "report_bug": null, "feature_request": null, "track_shipment": null, "return_item": null,
        "change_address": null, "upgrade_plan": null, "downgrade_plan": null, "dispute_charge": null,
        "account_locked": null, "two_factor_help": null, "api_error": null, "integration_help": null,
        "pricing_question": null, "demo_request": null, "gdpr_request": null, "delete_account": null,
        "invoice_copy": null, "tax_question": null, "partnership": null, "press_inquiry": null,
        "other": null
    });
    let state =
        json!({ "message": "I was double charged and want my money back for the duplicate." });

    let embed = embed_fn_from_agent(&agent, 512, 32);
    let top = shortlist_choice(&state, &criteria, &embed, 5, None)?;
    println!("top-5 shortlisted labels: {top:?}");

    let dim = agent.embed(&["hello world".to_string()], 512, 32)?[0].len();
    println!("encoder embedding dim: {dim}");
    assert!(top.len() == 5, "expected 5 labels");
    println!("OK");
    Ok(())
}

pub fn system_one( &self, state: &State, questions: &Questions, ) -> Result<SystemOneResult>

Evaluate typed questions over one state in a single forward pass.

Examples found in repository

examples/predict.rs (line 39)

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let model =
        std::env::var("LAYA_MODEL").unwrap_or_else(|_| "convaiinnovations/laya".to_string());
    let subfolder = std::env::var("LAYA_SUBFOLDER").ok();
    let device = std::env::var("LAYA_DEVICE").ok().filter(|s| !s.is_empty());

    let t = Instant::now();
    let agent = Agent::load(
        &model,
        LoadOptions {
            subfolder,
            device: device.clone(),
            ..Default::default()
        },
    )?;
    println!(
        "loaded {model} on {} in {:?}",
        device.as_deref().unwrap_or("cpu"),
        t.elapsed()
    );

    let state: State = serde_json::json!({
        "message": "I was charged twice this month and I'm furious. Refund me now or I'm leaving."
    });
    let questions = triage_questions();

    let t = Instant::now();
    let result = agent.system_one(&state, &questions)?;
    println!("forward in {:?}\n", t.elapsed());
    println!("{}", serde_json::to_string_pretty(&result.to_json())?);

    // Sanity: every probability distribution should sum to ~1.
    for (qid, ans) in &result.answers {
        if let Some(probs) = ans.get("probabilities").and_then(|p| p.as_object()) {
            let sum: f64 = probs.values().filter_map(|v| v.as_f64()).sum();
            assert!((sum - 1.0).abs() < 1e-2, "{qid} probabilities sum to {sum}");
        }
    }
    println!("\nOK: probability distributions sum to 1.");
    Ok(())
}

pub fn predict( &self, state: &State, questions: &Questions, ) -> Result<SystemOneResult>

Alias for Agent::system_one.

pub fn predict_batch( &self, states: &[State], questions: &Questions, batch_size: Option<usize>, ) -> Result<Vec<SystemOneResult>>

Evaluate the same questions over many states, sharing forward passes.

pub fn embed( &self, texts: &[String], max_length: usize, batch_size: usize, ) -> Result<Vec<Vec<f32>>>

Mean-pooled encoder embeddings for a list of texts ([n, hidden_size]).

Tokenizes with special tokens, truncates to max_length, runs the checkpoint encoder in batches of batch_size, and mean-pools over non-padding tokens. Backs crate::shortlist::embed_fn_from_agent; a dedicated bi-encoder will usually shortlist better, but this needs only the loaded checkpoint.

Examples found in repository

examples/shortlist.rs (line 36)

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let model =
        std::env::var("LAYA_MODEL").unwrap_or_else(|_| "convaiinnovations/laya".to_string());
    let device = std::env::var("LAYA_DEVICE").ok().filter(|s| !s.is_empty());
    let agent = Agent::load(
        &model,
        LoadOptions {
            device,
            ..Default::default()
        },
    )?;

    // A large label set (25 intents); shortlist to the 5 most relevant before scoring.
    let criteria = json!({
        "refund": null, "cancel_subscription": null, "reset_password": null, "update_billing": null,
        "report_bug": null, "feature_request": null, "track_shipment": null, "return_item": null,
        "change_address": null, "upgrade_plan": null, "downgrade_plan": null, "dispute_charge": null,
        "account_locked": null, "two_factor_help": null, "api_error": null, "integration_help": null,
        "pricing_question": null, "demo_request": null, "gdpr_request": null, "delete_account": null,
        "invoice_copy": null, "tax_question": null, "partnership": null, "press_inquiry": null,
        "other": null
    });
    let state =
        json!({ "message": "I was double charged and want my money back for the duplicate." });

    let embed = embed_fn_from_agent(&agent, 512, 32);
    let top = shortlist_choice(&state, &criteria, &embed, 5, None)?;
    println!("top-5 shortlisted labels: {top:?}");

    let dim = agent.embed(&["hello world".to_string()], 512, 32)?[0].len();
    println!("encoder embedding dim: {dim}");
    assert!(top.len() == 5, "expected 5 labels");
    println!("OK");
    Ok(())
}