m1rhan/laya-typed-decisions-ONNX

Laya (typed-decisions) — ONNX

ONNX export of `convaiinnovations/laya`, typed-decisions checkpoint — a ModernBERT-large encoder (421M params) with a small custom decision head, tuned specifically for typed choice / score / noul decisions (0.766 acc on that task family). All credit for the model itself goes to the original authors; this repo only adds ONNX conversion and quantization for in-browser (transformers.js / onnxruntime-web / WebGPU) use.

Files

Plain fp32 and int8 are not included:

• —fp32 export is ~1.69GB — verified byte-exact-enough (max abs diff 2.6e-6 to 7.1e-6 across all 6 fixtures) but too large to be a useful browser artifact.
• —int8 (onnxruntime.quantization.quantize_dynamic) fails during its own internal shape-inference pass on this graph (ShapeInferenceError: Inferred shape and existing shape differ in dimension 0: (1028) vs (256)), reproduced identically across two onnxruntime versions (1.23.2, 1.30.0) and two machines. Not a config issue — a real tooling incompatibility with this graph.
• —q4f16 (fp16-converted graph + 4-bit weights, for faster WebGPU compute) was attempted three ways — blocking the specific failing node, blocking the Where/Gemm op types, blocking MatMulNBits/DynamicQuantizeLinear — and hit the same graph type-consistency error each time (Type (tensor(float16)) ... does not match expected type (tensor(float)) on the model's mixed-type Where output node). Looks like a real bug in onnxconverter_common's handling of this dynamo-exported graph shape, not a one-line fix. Not included.

Usage (transformers.js / onnxruntime-web)

The model takes 5 inputs and returns 2 outputs:

• —Inputs: input_ids (int64), attention_mask (int64), marker_pos (int64), marker_mask (bool), qtype (int64, 0=choice/1=score/2=noul)
• —Outputs: logits (float32, one score per marker), act_logits (float32, 2-way escalate/answer gate)

Sequence construction (state + typed instructions + markers) and decode (temperature calibration → softmax → argmax/expected-value) follow the original SDK's build_sequence / Agent.system_one logic — see `convaiinnovations/laya` for the reference implementation.

Export details

• —Exported with torch.onnx.export(..., dynamo=True, dynamic_shapes=<torch.export.Dim-based>, opset_version=18) — the legacy TorchScript exporter (dynamo=False) silently bakes in the traced sequence length despite declaring dynamic_axes, which is a real trap if only verified at one shape. This export is verified across varying real sequence lengths (50–133 tokens) and marker counts (2–14).
• —Quantized with onnxruntime.quantization.matmul_nbits_quantizer.MatMulNBitsQuantizer.