laya-onnx

Laya, exported to ONNX

ONNX exports of the three Laya checkpoints, published for ruby-laya. Laya is a non-autoregressive System 1 decision engine by Convai Innovations: typed decisions over any state in a single forward pass, with calibrated probabilities.

These are conversions, not new models. All credit for the weights belongs to Convai Innovations.

Source revision: convaiinnovations/laya at commit 1c5edc17a7acd8701df6fc341c0d179f1c62c982, exported with laya 0.3.7 and PyTorch 2.14.0.

What each folder holds

model.onnx             the decision model: encoder, typed decision head and action head
rl_agent_config.json   copied from the source checkpoint (token budgets, temperatures)
onnx_config.json       export provenance and the verified deviation from PyTorch
tokenizer/             copied from the source checkpoint

Graph

Inputs are input_ids, attention_mask (int64 [batch, seq]), marker_pos, marker_mask (int64 and bool [batch, markers]) and qtype (int64 [batch]). Outputs are logits [batch, markers], act_logits [batch, 2] and last_hidden_state [batch, seq, dim], which the embedding shortlist uses. ONNX Runtime prunes the graph to the outputs you request, so asking for the first two costs nothing extra.

A question with a single option is padded to two markers with the second masked off, and a batch shorter than eight tokens is padded up to it. Both are what the traced graph expects, and neither changes an answer.

Weights are stored as float16 and cast to float32 in the graph. The published checkpoints are float16 on disk and upstream widens them at load, so the round trip is exact and the download stays the size of the original safetensors.

Verification

Every export is checked against upstream laya on CPU: same sequence construction, same temperature buckets, same rounding. Each probability, score, confidence and action probability upstream reports is reproduced, and raw logits agree within 2.5e-5. Reproduce with tools/export_onnx.py:

uv run tools/export_onnx.py <checkpoint-dir> <output-dir>

Use

gem install ruby-laya

require "laya"
agent = Laya.load("convaiinnovations/laya")   # resolves to english/ here

Any ONNX Runtime can load these directly; the gem is a convenience, not a requirement.

License

Apache 2.0, inherited from the source checkpoints. Laya was developed by Convai Innovations.