feat: add ONNX export script for production inference

- Convert PyTorch model to ONNX format
- Apply ONNX optimizations for BERT models
- Verify inference matches PyTorch outputs
- Benchmark PyTorch vs ONNX latency

ml/export/convert_to_onnx.py

@@ -0,0 +1,254 @@
+"""Convert trained PyTorch model to ONNX format for fast inference."""
+
+import json
+import time
+from pathlib import Path
+
+import numpy as np
+import onnx
+import onnxruntime as ort
+import torch
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
+
+def convert_to_onnx(
+    model_dir: str = "ml/artifacts/complexity-classifier",
+    output_path: str | None = None,
+    opset_version: int = 14,
+    optimize: bool = True,
+) -> str:
+    """
+    Convert a trained HuggingFace model to ONNX format.
+
+    Args:
+        model_dir: Directory containing trained model
+        output_path: Output path for ONNX model (defaults to model_dir/model.onnx)
+        opset_version: ONNX opset version
+        optimize: Whether to apply ONNX optimizations
+
+    Returns:
+        Path to the saved ONNX model
+    """
+    model_dir = Path(model_dir)
+    output_path = Path(output_path or model_dir / "model.onnx")
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    print(f"Converting model to ONNX")
+    print(f"  Model dir: {model_dir}")
+    print(f"  Output: {output_path}")
+    print(f"  Opset: {opset_version}")
+
+    # Load model and tokenizer
+    print("\nLoading model...")
+    tokenizer = AutoTokenizer.from_pretrained(model_dir)
+    model = AutoModelForSequenceClassification.from_pretrained(model_dir)
+    model.eval()
+
+    # Create dummy input for tracing
+    dummy_text = "This is a sample text for tracing the model."
+    dummy_inputs = tokenizer(
+        dummy_text,
+        padding="max_length",
+        truncation=True,
+        max_length=128,
+        return_tensors="pt",
+    )
+
+    # Define input names and dynamic axes
+    input_names = ["input_ids", "attention_mask"]
+    output_names = ["logits"]
+
+    dynamic_axes = {
+        "input_ids": {0: "batch_size", 1: "sequence"},
+        "attention_mask": {0: "batch_size", 1: "sequence"},
+        "logits": {0: "batch_size"},
+    }
+
+    # Export to ONNX
+    print("\nExporting to ONNX...")
+    torch.onnx.export(
+        model,
+        (dummy_inputs["input_ids"], dummy_inputs["attention_mask"]),
+        str(output_path),
+        input_names=input_names,
+        output_names=output_names,
+        dynamic_axes=dynamic_axes,
+        opset_version=opset_version,
+        do_constant_folding=True,
+    )
+
+    print(f"Model exported to: {output_path}")
+
+    # Validate the model
+    print("\nValidating ONNX model...")
+    onnx_model = onnx.load(str(output_path))
+    onnx.checker.check_model(onnx_model)
+    print("ONNX model validation passed!")
+
+    # Apply optimizations if requested
+    if optimize:
+        print("\nApplying ONNX optimizations...")
+        from onnxruntime.transformers import optimizer
+
+        optimized_path = output_path.with_suffix(".optimized.onnx")
+        optimized_model = optimizer.optimize_model(
+            str(output_path),
+            model_type="bert",
+            num_heads=12,
+            hidden_size=768,
+        )
+        optimized_model.save_model_to_file(str(optimized_path))
+        print(f"Optimized model saved to: {optimized_path}")
+
+        # Use optimized model
+        output_path = optimized_path
+
+    # Verify inference
+    print("\nVerifying inference...")
+    _verify_onnx_inference(model, tokenizer, output_path)
+
+    # Benchmark
+    print("\nBenchmarking...")
+    pytorch_time, onnx_time = _benchmark_inference(model, tokenizer, output_path)
+
+    # Save conversion info
+    info = {
+        "original_model": str(model_dir),
+        "onnx_path": str(output_path),
+        "opset_version": opset_version,
+        "optimized": optimize,
+        "benchmark": {
+            "pytorch_ms": pytorch_time,
+            "onnx_ms": onnx_time,
+            "speedup": pytorch_time / onnx_time if onnx_time > 0 else 0,
+        },
+    }
+
+    info_path = output_path.with_suffix(".json")
+    with open(info_path, "w") as f:
+        json.dump(info, f, indent=2)
+
+    print("\n" + "=" * 50)
+    print("Conversion complete!")
+    print("=" * 50)
+    print(f"\nONNX model: {output_path}")
+    print(f"PyTorch latency: {pytorch_time:.2f}ms")
+    print(f"ONNX latency: {onnx_time:.2f}ms")
+    print(f"Speedup: {pytorch_time / onnx_time:.2f}x")
+
+    return str(output_path)
+
+
+def _verify_onnx_inference(model, tokenizer, onnx_path: Path) -> None:
+    """Verify ONNX model produces same outputs as PyTorch."""
+    # Test inputs
+    test_texts = [
+        "Hello, how are you?",
+        "Write a Python function to calculate the factorial of a number recursively.",
+    ]
+
+    for text in test_texts:
+        inputs = tokenizer(
+            text,
+            padding="max_length",
+            truncation=True,
+            max_length=128,
+            return_tensors="pt",
+        )
+
+        # PyTorch inference
+        with torch.no_grad():
+            pytorch_outputs = model(**inputs)
+            pytorch_logits = pytorch_outputs.logits.numpy()
+
+        # ONNX inference
+        session = ort.InferenceSession(str(onnx_path))
+        onnx_inputs = {
+            "input_ids": inputs["input_ids"].numpy(),
+            "attention_mask": inputs["attention_mask"].numpy(),
+        }
+        onnx_outputs = session.run(None, onnx_inputs)
+        onnx_logits = onnx_outputs[0]
+
+        # Compare
+        np.testing.assert_allclose(pytorch_logits, onnx_logits, rtol=1e-3, atol=1e-4)
+
+    print("  Inference verification passed!")
+
+
+def _benchmark_inference(
+    model, tokenizer, onnx_path: Path, num_runs: int = 100
+) -> tuple[float, float]:
+    """Benchmark PyTorch vs ONNX inference latency."""
+    test_text = "What is the capital of France?"
+    inputs = tokenizer(
+        test_text,
+        padding="max_length",
+        truncation=True,
+        max_length=128,
+        return_tensors="pt",
+    )
+
+    # Warmup
+    with torch.no_grad():
+        _ = model(**inputs)
+
+    session = ort.InferenceSession(str(onnx_path))
+    onnx_inputs = {
+        "input_ids": inputs["input_ids"].numpy(),
+        "attention_mask": inputs["attention_mask"].numpy(),
+    }
+    _ = session.run(None, onnx_inputs)
+
+    # Benchmark PyTorch
+    start = time.perf_counter()
+    for _ in range(num_runs):
+        with torch.no_grad():
+            _ = model(**inputs)
+    pytorch_time = (time.perf_counter() - start) / num_runs * 1000  # ms
+
+    # Benchmark ONNX
+    start = time.perf_counter()
+    for _ in range(num_runs):
+        _ = session.run(None, onnx_inputs)
+    onnx_time = (time.perf_counter() - start) / num_runs * 1000  # ms
+
+    return pytorch_time, onnx_time
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Convert model to ONNX")
+    parser.add_argument(
+        "--model-dir",
+        type=str,
+        default="ml/artifacts/complexity-classifier",
+        help="Model directory",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default=None,
+        help="Output path for ONNX model",
+    )
+    parser.add_argument(
+        "--opset",
+        type=int,
+        default=14,
+        help="ONNX opset version",
+    )
+    parser.add_argument(
+        "--no-optimize",
+        action="store_true",
+        help="Skip ONNX optimizations",
+    )
+
+    args = parser.parse_args()
+
+    convert_to_onnx(
+        model_dir=args.model_dir,
+        output_path=args.output,
+        opset_version=args.opset,
+        optimize=not args.no_optimize,
+    )