Agent(
    model_id_or_path: str = "convaiinnovations/laya",
    device: Optional[str] = None,
    token: Optional[str] = None,
    subfolder: Optional[str] = None,
    fast: bool = False,
    compile: bool = False,
    revision: Optional[str] = None,
    expected_sha256: Optional[Dict[str, str]] = None,
    lang_temperatures: Optional[Dict[str, Dict[str, Any]]] = None,
    hooks=None,
    on_predict_start=None,
    on_predict_end=None,
    hooks_raise: bool = True,
    hooks_concurrent: bool = True,
    hooks_timeout: Optional[float] = None,
    calibration: Optional[str] = None,
)

System 1 decision model runtime: fast, non-autoregressive, calibrated decisions.

dtype is the autocast target, not the precision of every call. On MPS a call autocasts only at or above mps_amp_min_rows rows, so dtype can say float16 while a call runs in float32. dtype_for(rows) returns the precision of a call with rows rows.

Load a Laya checkpoint.

revision optionally pins the Hub download to an explicit commit SHA/branch/tag; when omitted, huggingface_hub's normal default and existing offline cache are used. expected_sha256 ({path relative to the checkpoint dir: hexdigest}) verifies artifact integrity before any weight is parsed or executed; it is opt-in and applies to local directories too. A missing artifact raises FileNotFoundError and a digest mismatch raises ValueError; either error refuses the load.

fast=True swaps the encoder/head forward for the TileLang fast path (CUDA only, needs pip install laya[fast]); see Agent.accelerate.

compile=True runs the model under torch.compile and turns ModernBERT's encoder reference_compile on. torch.compile specializes per input shape and Laya sees a new one on almost every request, so those graphs usually cost more than they return; use it when the traffic is repetitive. fast=True takes precedence, because the TileLang path replaces the forward that would be compiled.

subfolder selects one checkpoint from a repo that bundles several, e.g. Agent("convaiinnovations/laya", subfolder="multilingual"). Only that subfolder is downloaded, so bundling does not cost every user the whole family.

calibration is an optional JSON path with temperature and temperature_by_options. It is applied after the checkpoint config, so a fitted map overrides shipped scalars without rewriting model.safetensors.

hooks / on_predict_start / on_predict_end observe or shape every prediction; see laya.hooks. hooks_raise=False warns and continues when a hook fails, hooks_concurrent=False serialises hooks that are not safe to run in parallel, and hooks_timeout bounds each hook call in seconds (None means no limit).

accelerate(use_graphs: bool = True, strict: bool = False)

Replace the model forward with the TileLang fast path (fused GEMM/GEGLU/LayerNorm/RoPE kernels, sliding-window flash attention, 16-bit resident weights, CUDA graphs per shape bucket).

The fast path runs in the agent's autocast dtype at the time of the call (bf16 or fp16), so it matches the stock forward it replaces within rounding (see benchmarks/parity_fast.py). After changing agent.dtype, call deaccelerate() then accelerate() to rebuild it. Returns True if enabled. With strict=False any failure (no CUDA, tilelang missing) leaves the stock path in place.

warmup(shapes=None) -> float

Run the forward on synthetic input of each shape now and return the seconds it took.

compile=True traces and compiles on the first request that needs a graph (tens of seconds on a GPU), and fast=True builds its kernels and CUDA graphs per shape bucket on first use. Calling this after loading, before serving, moves that cost out of the first requests. With the stock forward it is a few ordinary forward passes. shapes is a list of (rows, tokens, markers); tokens are capped at the agent's max_len. Nothing is returned to or recorded for any caller, and hooks do not run.

deaccelerate()

Restore the stock forward.

dtype_for(rows: int) -> torch.dtype

Precision that a forward pass with rows question rows runs in.

dtype is the autocast target, set once at load time. Whether a forward autocasts is decided per call: on MPS only at or above mps_amp_min_rows rows. This returns dtype when a forward with rows rows autocasts, and torch.float32 when it does not. A predict call runs one row per question.

predict_batch(
    states: List[Union[str, dict, list]],
    questions: Dict[str, Dict[str, Any]],
    batch_size: Optional[int] = None,
    lang: Optional[str] = None,
    hooks=None,
    on_predict_start=None,
    on_predict_end=None,
    hooks_raise: Optional[bool] = None,
    hooks_timeout: Optional[float] = None,
    max_len: Optional[int] = None,
    head_max_len: Optional[int] = None,
    sort_by_length: bool = False,
    min_confidence: Optional[float] = None,
) -> List[Dict[str, Any]]

Evaluate the same questions over many states, packing them into shared forward passes.

This is the throughput path. system_one/predict handle one state per forward pass; on a GPU that leaves most of the batch dimension idle. predict_batch collates several states' question rows into one tensor, so a call that would take N sequential forward passes takes one (or ceil(len(states) / batch_size)), which is several times faster per decision on GPU.

predict_long(
    state: Union[str, dict, list],
    questions: Dict[str, Dict[str, Any]],
    window: Optional[int] = None,
    stride: Optional[int] = None,
    aggregate: str = "auto",
    batch_size: Optional[int] = None,
    lang: Optional[str] = None,
    hooks=None,
    on_predict_start=None,
    on_predict_end=None,
    hooks_raise: Optional[bool] = None,
    hooks_timeout: Optional[float] = None,
) -> Dict[str, Any]

Evaluate questions over a state longer than the context window, scanning it in overlapping windows and aggregating per question.

system_one/predict truncate a state that exceeds max_len to a single window (the first, or for a conversation list the last), silently dropping the rest. predict_long tokenizes the state once, splits it into overlapping token windows, scores every window in shared forward passes (via predict_batch), and combines the per-window answers:

• noul -> P(true) is the max over windows (the statement holds if any window supports it)
• choice-> the answer from the single most-confident window, so a localized signal isn't out-voted by the many neutral windows a long document is mostly made of (averaging drowns it -- the neutral majority dominates)
• score -> the level from the most-confident window, likewise

The returned probability/confidence is the deciding window's, not a calibrated number for the whole document: a noul max over many windows drifts up with the window count even with no signal, and choice can land on a confidently-neutral window when nothing in the document is decisive. Each answer therefore carries answer["window"] — the deciding window's index, token_start/token_end into the tokenized state, and the window count — so a caller can inspect the span the answer came from rather than trust the raw number.

A state that already fits one window is passed straight to system_one (identical output).

The hooks wrap the inference that answers the state, which for a document needing several windows is the one shared predict_batch over them: on_predict_start fires once, and ctx.states holds the decoded window texts in scan order -- not the caller's state, which was tokenized to produce them. Three outcomes follow from what the chain leaves behind:

• ctx.skip([result]) answers the document: the payload comes back with no window attribution and usage["windows"] at 0, because nothing was scored
• a scan left as this method built it: every window is scored, each answer carries answer["window"], and usage["windows"] is the window count
• a rewritten scan (ctx.states replaced, in any way): the answers are aggregated over the states that were scored, but no answer carries answer["window"] -- the offsets above describe this method's windows, not the text the model read

Returns a single result dict, the same shape as system_one, with usage["windows"] added. The key is always present and counts the windows the model scored to produce the answer: 1 for a state that fit one window, N for a document scanned in N overlapping windows (or the N a start hook rewrote them to), and 0 when a start hook answered the document, or left no states to score, before any window was read -- on either path, so a cached answer never reads as a window the model read.

Across several windows the truncation keys are combined like every other usage field: truncated, state_tokens and state_tokens_dropped are summed (so truncated is the number of windows that were cut, and the token counts include the overlap), and truncated_questions is the last window's list. The two can disagree: when only an earlier window was cut, truncated is above 0 and truncated_questions is empty. A window is cut when it is larger than the room a question's head leaves, from a window above the default or a start hook that narrows max_len / head_max_len. Test usage["truncated"] > 0 here, not is True.

system_one(
    state: Union[str, dict, list],
    questions: Dict[str, Dict[str, Any]],
    lang: Optional[str] = None,
    hooks=None,
    on_predict_start=None,
    on_predict_end=None,
    hooks_raise: Optional[bool] = None,
    hooks_timeout: Optional[float] = None,
    max_len: Optional[int] = None,
    head_max_len: Optional[int] = None,
    min_confidence: Optional[float] = None,
) -> Dict[str, Any]

Evaluate typed questions across state in a single, parallel forward pass.

To score many states at once, see predict_batch, which shares forward passes across them.

decide(
    state: Union[str, dict, list],
    schema: Any = None,
    questions: Optional[Dict[str, Any]] = None,
    return_details: bool = False,
    min_confidence: Optional[float] = None,
    predict_kwargs,
) -> Any

Answer state against a schema (JSON schema or pydantic model) and return typed values.

See laya.structured. Pass exactly one of schema or questions; extra keyword arguments are forwarded to predict / system_one.

decide_batch(
    states: List[Union[str, dict, list]],
    schema: Any = None,
    questions: Optional[Dict[str, Any]] = None,
    return_details: bool = False,
    min_confidence: Optional[float] = None,
    predict_kwargs,
) -> List[Any]

Answer many states against one schema (JSON schema or pydantic model) in one batched call.

The throughput form of :meth:decide: the schema is planned once and its questions run over every state through :meth:predict_batch (shared forward passes, results in input order), then each state's answers are projected as decide does. Extra keyword arguments (batch_size=, lang=, hooks=, ...) are forwarded to predict_batch. See laya.structured.

fit_temperatures(records, compute_ece: bool = False, seed: int = 0) -> Dict[str, Any]

Fit per-bucket temperatures from CPU records and store them on this agent.

records are (qtype, logits, target, k). Build them with laya.calibrate.records_from_labeled when you have labeled forwards; this method does not download weights or write model.safetensors. seed only affects the held-out ECE split when compute_ece is true. The checkpoint cfg is left as loaded.

save_calibration(path: str) -> None

Write temperatures and the checkpoint they were fitted for. Does not write weights.

load_calibration(path: str) -> None

Read a JSON map written by save_calibration onto this agent.

A file with no version is treated as version 1 and still loads. A newer file whose recorded checkpoint does not match this agent warns and still loads. Values that are not numbers, or that sit outside [TEMP_MIN, TEMP_MAX], are clamped with clamp_temperature the same way checkpoint load is.

load(
    model_id_or_path: str = "convaiinnovations/laya",
    device: Optional[str] = None,
    token: Optional[str] = None,
    subfolder: Optional[str] = None,
    fast: bool = False,
    compile: bool = False,
    revision: Optional[str] = None,
    expected_sha256: Optional[Dict[str, str]] = None,
    lang_temperatures: Optional[Dict[str, Dict[str, Any]]] = None,
    hooks=None,
    on_predict_start=None,
    on_predict_end=None,
    hooks_raise: bool = True,
    hooks_concurrent: bool = True,
    hooks_timeout: Optional[float] = None,
    calibration: Optional[str] = None,
) -> Agent

Load a Laya agent.

subfolder picks one checkpoint out of a repo that bundles several:

laya.load("convaiinnovations/laya")                           # English (repo root)
laya.load("convaiinnovations/laya", subfolder="multilingual")
laya.load("convaiinnovations/laya", fast=True)                # TileLang GPU fast path
laya.load("convaiinnovations/laya", compile=True)              # torch.compile the model

revision/expected_sha256 pin and verify the downloaded artifacts; see Agent. hooks / on_predict_start / on_predict_end observe or shape every prediction; see laya.hooks. calibration is the same optional JSON path accepted by Agent.

ONNXAgent(
    model_id_or_path: str,
    onnx_path: str = "laya.onnx",
    token: Optional[str] = None,
    subfolder: Optional[str] = None,
    revision: Optional[str] = None,
    expected_sha256: Optional[Dict[str, str]] = None,
    hooks=None,
    on_predict_start=None,
    on_predict_end=None,
    hooks_raise: bool = True,
    hooks_concurrent: bool = True,
    hooks_timeout: Optional[float] = None,
    lang_temperatures: Optional[Dict[str, Dict[str, Any]]] = None,
    calibration: Optional[str] = None,
)

System 1 decision model runtime via ONNX: fast CPU-optimized decisions.

Load a Laya agent backed by ONNX Runtime.

load_calibration(path: str) -> None

Read a JSON map written by save_calibration onto this agent.

system_one(
    state: Union[str, dict, list],
    questions: Dict[str, Dict[str, Any]],
    lang: Optional[str] = None,
    hooks=None,
    on_predict_start=None,
    on_predict_end=None,
    hooks_raise: Optional[bool] = None,
    hooks_timeout: Optional[float] = None,
    max_len: Optional[int] = None,
    head_max_len: Optional[int] = None,
    min_confidence: Optional[float] = None,
) -> Dict[str, Any]

Evaluate typed questions across state in one ONNX Runtime session run.

lang selects a per-language temperature override (see lang_temperatures), matching the PyTorch Agent.system_one signature so either backend is a drop-in for the other.

Defined in terms of predict_batch, exactly as the PyTorch Agent.system_one is, so the single-state and batched paths cannot drift apart.

To score many states at once, see predict_batch, which shares session runs across them.

predict_batch(
    states: List[Union[str, dict, list]],
    questions: Dict[str, Dict[str, Any]],
    batch_size: Optional[int] = None,
    lang: Optional[str] = None,
    hooks=None,
    on_predict_start=None,
    on_predict_end=None,
    hooks_raise: Optional[bool] = None,
    hooks_timeout: Optional[float] = None,
    max_len: Optional[int] = None,
    head_max_len: Optional[int] = None,
    sort_by_length: bool = False,
    min_confidence: Optional[float] = None,
) -> List[Dict[str, Any]]

Evaluate the same questions over many states, sharing ONNX Runtime session runs.

The throughput path, mirroring laya.agent.Agent.predict_batch: system_one collates one state's question rows per session run, so N states cost N runs. predict_batch collates several states' rows into one run -- or ceil(len(states) / batch_size) of them -- which is where ONNX Runtime's own parallelism pays off on CPU.

predict_long(
    state: Union[str, dict, list],
    questions: Dict[str, Dict[str, Any]],
    window: Optional[int] = None,
    stride: Optional[int] = None,
    aggregate: str = "auto",
    batch_size: Optional[int] = None,
    lang: Optional[str] = None,
    hooks=None,
    on_predict_start=None,
    on_predict_end=None,
    hooks_raise: Optional[bool] = None,
    hooks_timeout: Optional[float] = None,
) -> Dict[str, Any]

Evaluate questions over a state longer than the context window, scanning it in overlapping windows and aggregating per question.

The ONNX port of laya.agent.Agent.predict_long, with the same aggregation rules: system_one truncates a state that exceeds max_len to a single window, silently dropping the rest. predict_long tokenizes the state once, splits it into overlapping token windows, scores every window through predict_batch -- so the windows share ONNX Runtime session runs rather than costing one each -- and combines the per-window answers:

• noul -> P(true) is the max over windows (the statement holds if any window supports it)
• choice-> the answer from the single most-confident window, so a localized signal isn't out-voted by the many neutral windows a long document is mostly made of
• score -> the level from the most-confident window, likewise

The returned probability/confidence is the deciding window's, not a calibrated number for the whole document, for the same reasons the PyTorch docstring gives. Each answer carries answer["window"] -- the deciding window's index, token_start/token_end into the tokenized state, and the window count.

A state that already fits one window is passed straight to system_one (identical output).

Returns a single result dict, the same shape as system_one, with usage["windows"] added. Across several windows the truncation keys are summed or carried the same way as in Agent.predict_long: truncated is a window count and truncated_questions is the last window's list, so truncated can be above 0 while the list is empty.

decide(
    state: Union[str, dict, list],
    schema: Any = None,
    questions: Optional[Dict[str, Dict[str, Any]]] = None,
    return_details: bool = False,
    min_confidence: Optional[float] = None,
    predict_kwargs,
) -> Any

Answer state against a schema (JSON schema or pydantic model) and return typed values.

See laya.structured. Pass exactly one of schema or questions; extra keyword arguments are forwarded to predict / system_one.

decide_batch(
    states: List[Union[str, dict, list]],
    schema: Any = None,
    questions: Optional[Dict[str, Dict[str, Any]]] = None,
    return_details: bool = False,
    min_confidence: Optional[float] = None,
    predict_kwargs,
) -> List[Any]

Answer many states against one schema through predict_batch; see laya.structured.

python scripts/export_onnx.py --model convaiinnovations/laya --output laya.onnx --quantize