Upload convert_to_onnx.py with huggingface_hub

convert_to_onnx.py

@@ -0,0 +1,47 @@
+import torch
+from pathlib import Path
+from src.s01_config import ModelConfig
+from src.s02_tokenizer import MusicTokenizer
+from src.s04_model import MusicTransformer
+
+def export_onnx():
+    # 1. Initialize tokenizer and model configuration
+    tokenizer = MusicTokenizer()
+    model_config = ModelConfig(vocab_size=tokenizer.vocab_size)
+    model = MusicTransformer.from_config(model_config)
+
+    # 2. Load model weights from the PyTorch checkpoint
+    checkpoint_path = Path("checkpoints/best.pt")
+    if not checkpoint_path.exists():
+        print(f"Error: Checkpoint not found at {checkpoint_path}")
+        return
+
+    print("Loading PyTorch model weights...")
+    checkpoint = torch.load(checkpoint_path, map_location="cpu")
+    model.load_state_dict(checkpoint["model_state_dict"])
+    model.eval()
+
+    # 3. Prepare dummy inputs (batch_size=1, sequence_length=10)
+    dummy_input_ids = torch.zeros((1, 10), dtype=torch.long)
+
+    # 4. Export to ONNX format with dynamic axes
+    onnx_output_path = Path("checkpoints/model.onnx")
+    print(f"Exporting model to ONNX format: {onnx_output_path}...")
+    
+    torch.onnx.export(
+        model,
+        (dummy_input_ids,),
+        str(onnx_output_path),
+        input_names=["input_ids"],
+        output_names=["logits"],
+        dynamic_axes={
+            "input_ids": {0: "batch_size", 1: "sequence_length"},
+            "logits": {0: "batch_size", 1: "sequence_length"}
+        },
+        opset_version=15,
+        do_constant_folding=True
+    )
+    print("ONNX export completed successfully!")
+
+if __name__ == "__main__":
+    export_onnx()