"""Build the dynamic FP32 OpenThai-SystemOne ONNX Runtime graph.""" from __future__ import annotations import argparse import shutil from pathlib import Path import numpy as np import onnx_ir as ir import torch from huggingface_hub import snapshot_download from transformers import AutoConfig, AutoModel from onnxruntime_genai.models.builders.qwen import Qwen35TextModel from .artifact import write_manifest from .constants import ( ABSTAIN_SLOT, MAX_POSITION_EMBEDDINGS, N_SLOTS, REQUIRED_ORT_GENAI_VERSION, SOURCE_MODEL, SOURCE_REVISION, ) class OpenThaiDynamicBuilder(Qwen35TextModel): """Qwen3.5 builder with internal state and the OpenThai decision head.""" def make_node(self, op_type, inputs, outputs, *, name, domain="", **kwargs): """Normalize ONNX integer attributes before ONNX IR serialization. ONNX represents flags such as ``keepdims`` as INT attributes. ORT GenAI 0.16 passes Python booleans for a few of them. ONNX IR 0.1.3 infers the correct INT attribute type but otherwise retains the bool, which protobuf serialization rejects because ``bool`` is not accepted for ``AttributeProto.i``. """ # ORT GenAI 0.16 emits an ORT-private default-domain # SimplifiedLayerNormalization op. ONNX Runtime can load that op, but # the ONNX checker correctly rejects it because it has no ONNX schema. # ONNX opset 23 standardizes the same RMSNorm operation, so lower the # standalone form to RMSNormalization. The fused skip form remains the # com.microsoft contrib op. if op_type == "SimplifiedLayerNormalization" and not domain: op_type = "RMSNormalization" self.graph.opset_imports[""] = 23 normalized = { key: int(value) if isinstance(value, (bool, np.bool_)) else value for key, value in kwargs.items() } return super().make_node( op_type, inputs, outputs, name=name, domain=domain, **normalized, ) def load_weights(self, input_path): return AutoModel.from_pretrained( input_path, trust_remote_code=True, local_files_only=True, dtype=torch.float32, ).eval() def make_inputs_and_outputs(self): self.input_shapes["position_ids"] = [3, "batch_size", "sequence_length"] public_inputs = { "input_ids": (ir.DataType.INT64, ["batch_size", "sequence_length"]), "attention_mask": (ir.DataType.INT64, ["batch_size", "sequence_length"]), "answer_positions": (ir.DataType.INT64, ["batch_size", "num_questions"]), "option_counts": (ir.DataType.INT64, ["batch_size", "num_questions"]), "qtypes": (ir.DataType.INT64, ["batch_size", "num_questions"]), } for name, (dtype, shape) in public_inputs.items(): self.model.graph.inputs.append(self.make_value(name, dtype=dtype, shape=shape)) output_shape = ["batch_size", "num_questions", N_SLOTS] self.model.graph.outputs.append(self.make_value("logits", ir.DataType.FLOAT, output_shape)) self.model.graph.outputs.append(self.make_value("probabilities", ir.DataType.FLOAT, output_shape)) def _dynamic_batch_dimension(self) -> str: shape_name = "/openthai/internal/input_shape" self.make_shape(shape_name, "input_ids", shape=[2]) gather_name = "/openthai/internal/batch_dimension" self.make_gather( gather_name, [f"{shape_name}/output_0", "/model/constants/INT64/0"], dtype=ir.DataType.INT64, shape=[], axis=0, ) unsqueeze_name = "/openthai/internal/batch_dimension_unsqueezed" self.make_unsqueeze( unsqueeze_name, [f"{gather_name}/output_0", "/model/constants/INT64/[0]"], dtype=ir.DataType.INT64, shape=[1], ) return f"{unsqueeze_name}/output_0" def _zero_tensor(self, name: str, shape_input: str, dtype: ir.DataType, shape): if dtype != ir.DataType.FLOAT: raise TypeError(f"internal request state must be FP32, got {dtype}") zero = ir.tensor(np.asarray([0], dtype=np.float32)) node_name = f"/openthai/internal/zero/{name.replace('.', '_')}" self.make_node( "ConstantOfShape", inputs=[shape_input], outputs=[name], name=node_name, value=zero, ) self.make_value(name, dtype=dtype, shape=shape) def _make_internal_request_state(self): batch = self._dynamic_batch_dimension() shape_specs = { "kv": [self.num_kv_heads, 0, self.head_size], "conv": [self.linear_conv_dim, self.linear_conv_kernel_dim - 1], "recurrent": [ self.linear_num_value_heads, self.linear_key_head_dim, self.linear_value_head_dim, ], } shape_paths: dict[str, str] = {} for label, tail in shape_specs.items(): concat_name = f"/openthai/internal/{label}_state_shape" self.make_concat( concat_name, [batch, f"/model/constants/INT64/{tail}"], dtype=ir.DataType.INT64, shape=[1 + len(tail)], axis=0, ) shape_paths[label] = f"{concat_name}/output_0" for cache_kind in ("past_key_values.key", "past_key_values.value"): for cache_name in self.input_names[cache_kind].values(): self._zero_tensor( cache_name, shape_paths["kv"], self.io_dtype, ["batch_size", self.num_kv_heads, 0, self.head_size], ) for cache_name in self.input_names["past.conv"].values(): self._zero_tensor( cache_name, shape_paths["conv"], self.io_dtype, ["batch_size", self.linear_conv_dim, self.linear_conv_kernel_dim - 1], ) for cache_name in self.input_names["past.recurrent"].values(): self._zero_tensor( cache_name, shape_paths["recurrent"], self.input_types["past.recurrent"], [ "batch_size", self.linear_num_value_heads, self.linear_key_head_dim, self.linear_value_head_dim, ], ) def _make_dynamic_text_positions(self): cumsum_name = "/openthai/positions/CumSum" self.make_node( "CumSum", inputs=["attention_mask", "/model/constants/INT64/1"], outputs=[f"{cumsum_name}/output_0"], name=cumsum_name, ) self.make_value( f"{cumsum_name}/output_0", ir.DataType.INT64, ["batch_size", "sequence_length"], ) sub_name = "/openthai/positions/SubOne" self.make_sub( sub_name, [f"{cumsum_name}/output_0", "/model/constants/INT64/1"], ir.DataType.INT64, ["batch_size", "sequence_length"], ) padded_name = "/openthai/positions/IsPadding" self.make_equal( padded_name, ["attention_mask", "/model/constants/INT64/0"], ["batch_size", "sequence_length"], ) where_name = "/openthai/positions/ReplacePadding" self.make_node( "Where", inputs=[ f"{padded_name}/output_0", "/model/constants/INT64/0", f"{sub_name}/output_0", ], outputs=[f"{where_name}/output_0"], name=where_name, ) self.make_value( f"{where_name}/output_0", ir.DataType.INT64, ["batch_size", "sequence_length"], ) unsqueeze_name = "/openthai/positions/Unsqueeze" self.make_unsqueeze( unsqueeze_name, [f"{where_name}/output_0", "/model/constants/INT64/[0]"], ir.DataType.INT64, [1, "batch_size", "sequence_length"], ) tile_name = "/openthai/positions/TileMRoPE" self.make_node( "Tile", inputs=[f"{unsqueeze_name}/output_0", "/model/constants/INT64/[3, 1, 1]"], outputs=[self.input_names["position_ids"]], name=tile_name, ) self.make_value( self.input_names["position_ids"], ir.DataType.INT64, [3, "batch_size", "sequence_length"], ) def make_preprocessing_nodes(self): self._make_internal_request_state() self._make_dynamic_text_positions() super().make_preprocessing_nodes() def make_postprocessing_nodes(self): """Gather answers, apply the calibrated head, mask slots, and softmax.""" shape = ["batch_size", "num_questions", N_SLOTS] index_shape_name = "/openthai/head/AnswerIndexShape" self.make_shape(index_shape_name, "answer_positions", shape=[2]) target_shape_name = "/openthai/head/ExpandedIndexShape" self.make_concat( target_shape_name, [f"{index_shape_name}/output_0", f"/model/constants/INT64/[{self.hidden_size}]"], ir.DataType.INT64, [3], ) clamped_name = "/openthai/head/ClampAnswerPositions" self.make_node( "Max", inputs=["answer_positions", "/model/constants/INT64/0"], outputs=[f"{clamped_name}/output_0"], name=clamped_name, ) self.make_value( f"{clamped_name}/output_0", ir.DataType.INT64, ["batch_size", "num_questions"], ) unsqueeze_name = "/openthai/head/UnsqueezeAnswerPositions" self.make_unsqueeze( unsqueeze_name, [f"{clamped_name}/output_0", "/model/constants/INT64/[-1]"], ir.DataType.INT64, ["batch_size", "num_questions", 1], ) expand_name = "/openthai/head/ExpandAnswerPositions" self.make_expand( expand_name, [f"{unsqueeze_name}/output_0", f"{target_shape_name}/output_0"], ir.DataType.INT64, ["batch_size", "num_questions", self.hidden_size], ) gather_name = "/openthai/head/GatherAnswerStates" self.make_node( "GatherElements", inputs=["hidden_states", f"{expand_name}/output_0"], outputs=[f"{gather_name}/output_0"], name=gather_name, axis=1, ) self.make_value( f"{gather_name}/output_0", ir.DataType.FLOAT, ["batch_size", "num_questions", self.hidden_size], ) weight_name = "openthai.slot_head.weight.transposed" bias_name = "openthai.slot_head.bias" temperature_name = "openthai.temperature" self.make_initializer( self.weights.slot_head.weight.detach().float().t().contiguous(), weight_name, to=ir.DataType.FLOAT, ) self.make_initializer( self.weights.slot_head.bias.detach().float(), bias_name, to=ir.DataType.FLOAT, ) self.make_initializer( self.weights.log_temperature.detach().float().exp(), temperature_name, to=ir.DataType.FLOAT, ) matmul_name = "/openthai/head/SlotMatMul" self.make_node( "MatMul", inputs=[f"{gather_name}/output_0", weight_name], outputs=[f"{matmul_name}/output_0"], name=matmul_name, ) self.make_value(f"{matmul_name}/output_0", ir.DataType.FLOAT, shape) add_name = "/openthai/head/SlotBias" self.make_add( add_name, [f"{matmul_name}/output_0", bias_name], ir.DataType.FLOAT, shape, ) qtype_name = "/openthai/head/ClampQuestionType" self.make_node( "Max", inputs=["qtypes", "/model/constants/INT64/0"], outputs=[f"{qtype_name}/output_0"], name=qtype_name, ) self.make_value( f"{qtype_name}/output_0", ir.DataType.INT64, ["batch_size", "num_questions"], ) temp_gather_name = "/openthai/head/GatherTemperature" self.make_gather( temp_gather_name, [temperature_name, f"{qtype_name}/output_0"], ir.DataType.FLOAT, ["batch_size", "num_questions"], axis=0, ) temp_unsqueeze_name = "/openthai/head/UnsqueezeTemperature" self.make_unsqueeze( temp_unsqueeze_name, [f"{temp_gather_name}/output_0", "/model/constants/INT64/[-1]"], ir.DataType.FLOAT, ["batch_size", "num_questions", 1], ) scaled_name = "/openthai/head/ApplyTemperature" self.make_div( scaled_name, [f"{add_name}/output_0", f"{temp_unsqueeze_name}/output_0"], ir.DataType.FLOAT, shape, ) count_name = "/openthai/head/UnsqueezeOptionCounts" self.make_unsqueeze( count_name, ["option_counts", "/model/constants/INT64/[-1]"], ir.DataType.INT64, ["batch_size", "num_questions", 1], ) range_name = "/openthai/head/SlotRange" self.make_range( range_name, [ "/model/constants/INT64/0", f"/model/constants/INT64/{N_SLOTS}", "/model/constants/INT64/1", ], ir.DataType.INT64, [N_SLOTS], ) less_name = "/openthai/head/ValidOptions" self.make_node( "Less", inputs=[f"{range_name}/output_0", f"{count_name}/output_0"], outputs=[f"{less_name}/output_0"], name=less_name, ) self.make_value(f"{less_name}/output_0", ir.DataType.BOOL, shape) abstain_name = "/openthai/head/IsAbstain" self.make_equal( abstain_name, [f"{range_name}/output_0", f"/model/constants/INT64/{ABSTAIN_SLOT}"], [N_SLOTS], ) valid_name = "/openthai/head/IncludeAbstain" self.make_node( "Or", inputs=[f"{less_name}/output_0", f"{abstain_name}/output_0"], outputs=[f"{valid_name}/output_0"], name=valid_name, ) self.make_value(f"{valid_name}/output_0", ir.DataType.BOOL, shape) empty_name = "/openthai/head/EmptyQuestion" self.make_equal( empty_name, ["option_counts", "/model/constants/INT64/0"], ["batch_size", "num_questions"], ) empty_unsqueeze_name = "/openthai/head/UnsqueezeEmptyQuestion" self.make_unsqueeze( empty_unsqueeze_name, [f"{empty_name}/output_0", "/model/constants/INT64/[-1]"], ir.DataType.BOOL, ["batch_size", "num_questions", 1], ) slot_zero_name = "/openthai/head/IsSlotZero" self.make_equal( slot_zero_name, [f"{range_name}/output_0", "/model/constants/INT64/0"], [N_SLOTS], ) padding_slot_name = "/openthai/head/PaddingSlot" self.make_node( "And", inputs=[f"{empty_unsqueeze_name}/output_0", f"{slot_zero_name}/output_0"], outputs=[f"{padding_slot_name}/output_0"], name=padding_slot_name, ) self.make_value(f"{padding_slot_name}/output_0", ir.DataType.BOOL, shape) final_valid_name = "/openthai/head/FinalValidMask" self.make_node( "Or", inputs=[f"{valid_name}/output_0", f"{padding_slot_name}/output_0"], outputs=[f"{final_valid_name}/output_0"], name=final_valid_name, ) self.make_value(f"{final_valid_name}/output_0", ir.DataType.BOOL, shape) negative_infinity_name = "openthai.negative_infinity" self.make_initializer( np.asarray(float("-inf"), dtype=np.float32), negative_infinity_name, to=ir.DataType.FLOAT, ) self.make_node( "Where", inputs=[ f"{final_valid_name}/output_0", f"{scaled_name}/output_0", negative_infinity_name, ], outputs=["logits"], name="/openthai/head/MaskInvalidSlots", ) self.make_value("logits", ir.DataType.FLOAT, shape) self.make_node( "Softmax", inputs=["logits"], outputs=["probabilities"], name="/openthai/head/Softmax", axis=-1, ) self.make_value("probabilities", ir.DataType.FLOAT, shape) def _copy_release_files(snapshot: Path, destination: Path): tokenizer_names = { "tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "chat_template.jinja", } for name in tokenizer_names: source = snapshot / name if source.is_file(): shutil.copy2(source, destination / name) project_root = Path(__file__).resolve().parents[2] shutil.copy2(project_root / "LICENSE", destination / "LICENSE") shutil.copy2(project_root / "MODEL_CARD.md", destination / "README.md") shutil.copy2(project_root / "MODEL_CARD.md", destination / "MODEL_CARD.md") shutil.copy2(project_root / "ATTRIBUTION.md", destination / "ATTRIBUTION.md") shutil.copy2(project_root / "pyproject.toml", destination / "pyproject.toml") shutil.copy2(project_root / "uv.lock", destination / "uv.lock") shutil.copytree( project_root / "src" / "openthai_systemone_onnx", destination / "src" / "openthai_systemone_onnx", dirs_exist_ok=True, ignore=shutil.ignore_patterns("__pycache__", "*.pyc"), ) shutil.copytree( project_root / "examples", destination / "examples", dirs_exist_ok=True, ignore=shutil.ignore_patterns("__pycache__", "*.pyc"), ) def export_dynamic_model(output_dir: Path, cache_dir: Path, source: str, revision: str) -> Path: if revision != SOURCE_REVISION: raise ValueError( f"This validated recipe is pinned to source revision {SOURCE_REVISION}; got {revision}" ) snapshot = Path( snapshot_download( repo_id=source, revision=revision, cache_dir=cache_dir, allow_patterns=[ "*.json", "*.py", "*.jinja", "*.safetensors", "LICENSE*", "README.md", ], ) ) config = AutoConfig.from_pretrained(snapshot, trust_remote_code=True, local_files_only=True) if config.text_config.max_position_embeddings != MAX_POSITION_EMBEDDINGS: raise ValueError("the pinned backbone no longer reports max_position_embeddings=262144") output_dir.mkdir(parents=True, exist_ok=True) onnx_dir = output_dir / "onnx" onnx_dir.mkdir(parents=True, exist_ok=True) builder_cache = cache_dir / "builder" builder_cache.mkdir(parents=True, exist_ok=True) options = { "exclude_lm_head": True, "exclude_mtp": True, "filename": "model.onnx", "hf_remote": True, "linear_attn_op": "linear_attention", "state_window": 0, "use_paged_attention": False, } builder = OpenThaiDynamicBuilder( config, ir.DataType.FLOAT, ir.DataType.FLOAT, "cpu", str(builder_cache), options, ) builder.model.metadata_props.update( { "source_model": source, "source_revision": revision, "max_position_embeddings": str(MAX_POSITION_EMBEDDINGS), "onnxruntime_genai_builder": REQUIRED_ORT_GENAI_VERSION, "public_state": "internal-zero-state-per-request", } ) builder.make_model(str(snapshot)) builder.save_model(str(onnx_dir)) _copy_release_files(snapshot, output_dir) return write_manifest(output_dir) def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--output", type=Path, default=Path("artifacts/OpenThai-SystemOne-ONNX")) parser.add_argument("--cache-dir", type=Path, default=Path(".cache/huggingface")) parser.add_argument("--source", default=SOURCE_MODEL) parser.add_argument("--revision", default=SOURCE_REVISION) return parser def main(argv: list[str] | None = None): args = build_parser().parse_args(argv) manifest = export_dynamic_model(args.output.resolve(), args.cache_dir.resolve(), args.source, args.revision) print(f"Export complete: {manifest}") if __name__ == "__main__": main()