Skip to main content

shortlist/
shortlist.rs

1//! Demo: encoder-embedding shortlist for a high-cardinality choice question.
2//!   cargo run --release --example shortlist
3use laya::agent::{Agent, LoadOptions};
4use laya::shortlist::{embed_fn_from_agent, shortlist_choice};
5use serde_json::json;
6
7fn main() -> Result<(), Box<dyn std::error::Error>> {
8    let model =
9        std::env::var("LAYA_MODEL").unwrap_or_else(|_| "convaiinnovations/laya".to_string());
10    let device = std::env::var("LAYA_DEVICE").ok().filter(|s| !s.is_empty());
11    let agent = Agent::load(
12        &model,
13        LoadOptions {
14            device,
15            ..Default::default()
16        },
17    )?;
18
19    // A large label set (25 intents); shortlist to the 5 most relevant before scoring.
20    let criteria = json!({
21        "refund": null, "cancel_subscription": null, "reset_password": null, "update_billing": null,
22        "report_bug": null, "feature_request": null, "track_shipment": null, "return_item": null,
23        "change_address": null, "upgrade_plan": null, "downgrade_plan": null, "dispute_charge": null,
24        "account_locked": null, "two_factor_help": null, "api_error": null, "integration_help": null,
25        "pricing_question": null, "demo_request": null, "gdpr_request": null, "delete_account": null,
26        "invoice_copy": null, "tax_question": null, "partnership": null, "press_inquiry": null,
27        "other": null
28    });
29    let state =
30        json!({ "message": "I was double charged and want my money back for the duplicate." });
31
32    let embed = embed_fn_from_agent(&agent, 512, 32);
33    let top = shortlist_choice(&state, &criteria, &embed, 5, None)?;
34    println!("top-5 shortlisted labels: {top:?}");
35
36    let dim = agent.embed(&["hello world".to_string()], 512, 32)?[0].len();
37    println!("encoder embedding dim: {dim}");
38    assert!(top.len() == 5, "expected 5 labels");
39    println!("OK");
40    Ok(())
41}