| """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``. |
| """ |
| |
| |
| |
| |
| |
| |
| 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() |
|
|
