From 1cf7843cd2214303ab0e23740be35d603451ec34 Mon Sep 17 00:00:00 2001 From: sp25-bai-047-wq Date: Wed, 12 Aug 2026 10:14:45 +0500 Subject: [PATCH] Add inference/onnx_exporter.py --- inference/onnx_exporter.py | 61 +++++++++++++++++++------------------- 1 file changed, 31 insertions(+), 30 deletions(-) diff --git a/inference/onnx_exporter.py b/inference/onnx_exporter.py index 458353b..cc0c285 100644 --- a/inference/onnx_exporter.py +++ b/inference/onnx_exporter.py @@ -1,6 +1,12 @@ """ ONNX Model Exporter for CalculusSolver PyTorch weights. -Converts best.pt / model.pkl into best.onnx for lightweight CPU inference. +Converts best.pt into best.onnx for lightweight CPU inference. + +Rewritten to target model/simple_transformer.py's SimpleCalculusModel +(the architecture actually trained since PR #27's pivot), not the old +tree-structured model/architecture.py::CalculusModel. That class expects +src_positions/parent_child_pairs which were never produced by training +and are not part of this model's forward() signature. """ import os @@ -12,27 +18,26 @@ def export_to_onnx( checkpoint_path: str = "checkpoints/final/best.pt", output_path: str = "checkpoints/final/best.onnx", - vocab_path: str = "tokenizer/vocab.json" + vocab_path: str = "tokenizer/vocab.json", + hidden_dim: int = 256, # must match docs/TRAINING_RESULTS.md's "Hidden Dim" for the checkpoint being exported + max_len: int = 32, # must match config.json / TRAINING_RESULTS.md for the checkpoint being exported ) -> str: - """Export PyTorch CalculusModel state dict to ONNX format.""" + """Export PyTorch SimpleCalculusModel state dict to ONNX format.""" if not os.path.exists(checkpoint_path): raise FileNotFoundError(f"PyTorch checkpoint not found: {checkpoint_path}") - from model.architecture import CalculusModel + from model.simple_transformer import SimpleCalculusModel from inference.beam_search import load_vocab vocab_map = load_vocab(vocab_path) vocab_size = len(vocab_map["token_to_id"]) - rule_labels = [k.replace("RULE:", "") for k in vocab_map.get("rule_tokens", {}).keys()] + pad_id = vocab_map["token_to_id"].get("[PAD]", 0) - model = CalculusModel( + model = SimpleCalculusModel( vocab_size=vocab_size, - rule_labels=rule_labels, - hidden_dim=512, - num_heads=8, - num_layers=8, - ffn_dim=2048, - dropout=0.0 + hidden_dim=hidden_dim, + pad_id=pad_id, + max_len=max_len, ) checkpoint = torch.load(checkpoint_path, map_location="cpu") @@ -40,37 +45,33 @@ def export_to_onnx( model.load_state_dict(state_dict) model.eval() - # Dummy inputs for tracing + # Dummy inputs matching forward(self, src_seq, tgt_in_seq) -- no positions, + # no parent_child_pairs; this model never took them. batch_size = 1 - seq_len = 256 - dummy_src = torch.randint(0, vocab_size, (batch_size, seq_len), dtype=torch.long) - dummy_positions = torch.zeros((batch_size, seq_len, 3), dtype=torch.float32) - dummy_pairs = torch.zeros((batch_size, seq_len, seq_len), dtype=torch.float32) + dummy_src = torch.randint(1, vocab_size, (batch_size, max_len), dtype=torch.long) + dummy_tgt_in = torch.randint(1, vocab_size, (batch_size, max_len), dtype=torch.long) os.makedirs(os.path.dirname(output_path), exist_ok=True) torch.onnx.export( model, - (dummy_src, dummy_positions, dummy_pairs), + (dummy_src, dummy_tgt_in), output_path, - input_names=["src_tokens", "positions", "parent_child_pairs"], - output_names=["logits", "rule_logits"], + input_names=["src_seq", "tgt_in_seq"], + output_names=["logits"], dynamic_axes={ - "src_tokens": {0: "batch_size", 1: "seq_len"}, - "positions": {0: "batch_size", 1: "seq_len"}, - "parent_child_pairs": {0: "batch_size", 1: "seq_len", 2: "seq_len"}, - "logits": {0: "batch_size", 1: "seq_len"} + "src_seq": {0: "batch_size", 1: "seq_len"}, + "tgt_in_seq": {0: "batch_size", 1: "tgt_len"}, + "logits": {0: "batch_size", 1: "tgt_len"}, }, - opset_version=14 + opset_version=14, ) - print(f"[ONNX Export] Model successfully exported to: {output_path}") + size_mb = os.path.getsize(output_path) / (1024 * 1024) + print(f"[ONNX Export] Model successfully exported to: {output_path} ({size_mb:.1f} MB)") return output_path if __name__ == "__main__": ckpt = sys.argv[1] if len(sys.argv) > 1 else "checkpoints/final/best.pt" out = sys.argv[2] if len(sys.argv) > 2 else "checkpoints/final/best.onnx" - try: - export_to_onnx(ckpt, out) - except Exception as exc: - print(f"Export skipped/failed: {exc}") + export_to_onnx(ckpt, out) \ No newline at end of file