export_coreml.py

#!/usr/bin/env python3
"""Export the Kompress v2 PyTorch checkpoint to Apple CoreML (mlprogram).

Usage:
    # Run with default settings (inputs: merged.pt, output: coreml/kompress.mlpackage)
    python3 export_coreml.py

    # Run with custom parameters
    python3 export_coreml.py --checkpoint merged.pt --output-dir coreml --max-seq-len 1024
"""

import os
import argparse
import shutil
import torch
import torch.nn as nn
import numpy as np
from transformers import AutoModel, AutoTokenizer

# Standard ModernBERT-base model card
BASE_MODEL = "answerdotai/ModernBERT-base"

class HeadroomCompressorModel(nn.Module):
    """Dual-head ModernBERT: token classification + span importance CNN."""

    def __init__(self, model_name: str = BASE_MODEL):
        super().__init__()
        self.encoder = AutoModel.from_pretrained(model_name, attn_implementation="eager")
        hidden_size = self.encoder.config.hidden_size

        self.token_dropout = nn.Dropout(0.1)
        self.token_head = nn.Linear(hidden_size, 2)

        self.span_conv = nn.Sequential(
            nn.Conv1d(hidden_size, 256, kernel_size=5, padding=2),
            nn.GELU(),
            nn.Conv1d(256, 1, kernel_size=3, padding=1),
            nn.Sigmoid(),
        )

    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        hidden = self.encoder(input_ids, attention_mask=attention_mask).last_hidden_state
        token_probs = torch.softmax(self.token_head(hidden), dim=-1)[:, :, 1]
        span_scores = self.span_conv(hidden.transpose(1, 2)).squeeze(1)
        return token_probs * (0.5 + 0.5 * span_scores)

class CoreMLExportWrapper(nn.Module):
    """Wrapper that handles type conversions for Apple CoreML.
    
    CoreML accepts int32 inputs, but PyTorch and ModernBERT embedding layers
    expect int64 (long) tensors. This wrapper intercepts and casts inputs
    accordingly before feeding them to the model.
    """
    def __init__(self, inner):
        super().__init__()
        self.inner = inner

    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        return self.inner(input_ids.long(), attention_mask.long())

def export_to_coreml(
    checkpoint_path: str,
    output_path: str,
    max_seq_len: int = 2048,
    precision: str = "fp16",
    minimum_deployment_target: str = "iOS17"
):
    import coremltools as ct

    print("Step 1: Instantiating model and loading weights...")
    model = HeadroomCompressorModel()
    
    if not os.path.exists(checkpoint_path):
        raise FileNotFoundError(f"Checkpoint not found at: {checkpoint_path}")
        
    ckpt = torch.load(checkpoint_path, map_location="cpu")
    for key in ("encoder_state_dict", "token_head_state_dict", "span_conv_state_dict"):
        if key not in ckpt:
            raise ValueError(f"Checkpoint missing key: {key}. Are you sure this is a v2 merged checkpoint?")

    model.encoder.load_state_dict(ckpt["encoder_state_dict"], strict=True)
    model.token_head.load_state_dict(ckpt["token_head_state_dict"], strict=True)
    model.span_conv.load_state_dict(ckpt["span_conv_state_dict"], strict=True)
    model.eval()
    print("  Weights loaded successfully!")

    wrapper = CoreMLExportWrapper(model).eval()

    print("Step 2: JIT tracing the PyTorch model...")
    # Use sequence length of 512 as the default trace shape
    trace_len = min(512, max_seq_len)
    dummy_ids = torch.randint(0, 1000, (1, trace_len), dtype=torch.int32)
    dummy_mask = torch.ones((1, trace_len), dtype=torch.int32)

    with torch.no_grad():
        traced = torch.jit.trace(wrapper, (dummy_ids, dummy_mask), strict=False)
    print("  Model JIT-traced successfully!")

    print("Step 3: Converting to CoreML (mlprogram)...")
    # Define dynamic sequence length dimension (1 to max_seq_len)
    seq_dim = ct.RangeDim(lower_bound=1, upper_bound=max_seq_len, default=trace_len)
    
    inputs = [
        ct.TensorType(name="input_ids", shape=(1, seq_dim), dtype=np.int32),
        ct.TensorType(name="attention_mask", shape=(1, seq_dim), dtype=np.int32)
    ]
    
    outputs = [
        ct.TensorType(name="final_scores", dtype=np.float32)
    ]

    compute_precision = ct.precision.FLOAT16 if precision == "fp16" else ct.precision.FLOAT32
    target = getattr(ct.target, minimum_deployment_target)

    mlmodel = ct.convert(
        traced,
        inputs=inputs,
        outputs=outputs,
        convert_to="mlprogram",
        compute_precision=compute_precision,
        minimum_deployment_target=target
    )
    print("  Conversion completed successfully!")

    print(f"Step 4: Saving CoreML package to {output_path}...")
    if os.path.exists(output_path):
        shutil.rmtree(output_path)
    mlmodel.save(output_path)
    print("  Model saved successfully!")

    return wrapper, mlmodel

def validate_model(wrapper, mlmodel, max_seq_len: int):
    print("\nStep 5: Running numeric validation and sequence length checks...")
    
    # Test a few different sequence lengths to verify dynamic shape behavior
    test_lengths = [64, 128, 512]
    if max_seq_len >= 1024:
        test_lengths.append(1024)
        
    for seq_len in test_lengths:
        print(f"\n- Testing sequence length: {seq_len}")
        test_ids = torch.randint(0, 1000, (1, seq_len), dtype=torch.int32)
        test_mask = torch.ones((1, seq_len), dtype=torch.int32)

        # PyTorch reference prediction
        with torch.no_grad():
            py_out = wrapper(test_ids, test_mask).numpy()

        # CoreML prediction
        cml_inputs = {
            "input_ids": test_ids.numpy(),
            "attention_mask": test_mask.numpy()
        }
        
        try:
            cml_out = mlmodel.predict(cml_inputs)["final_scores"]
            
            max_abs_diff = np.max(np.abs(py_out - cml_out))
            py_keep = py_out > 0.5
            cml_keep = cml_out > 0.5
            agreement = np.mean(py_keep == cml_keep)
            
            print(f"  PyTorch range: [{py_out.min():.4f}, {py_out.max():.4f}]")
            print(f"  CoreML range : [{cml_out.min():.4f}, {cml_out.max():.4f}]")
            print(f"  Max absolute difference: {max_abs_diff:.6f}")
            print(f"  Decision agreement (threshold 0.5): {agreement * 100:.2f}%")
        except Exception as e:
            print(f"  ⚠️ CoreML execution warning/error: {e}")
            print("  (If this is an ANE compiler failure on macOS, it will fall back to CPU/GPU at runtime.)")

def main():
    parser = argparse.ArgumentParser(description="Convert Kompress v2 base model to CoreML.")
    parser.add_argument("--checkpoint", default="merged.pt", help="Path to input PyTorch merged.pt checkpoint.")
    parser.add_argument("--output-dir", default="coreml", help="Directory where output .mlpackage will be saved.")
    parser.add_argument("--max-seq-len", type=int, default=2048, help="Maximum sequence length dimension boundary.")
    parser.add_argument("--precision", choices=["fp16", "fp32"], default="fp16", help="Target weight and execution precision.")
    parser.add_argument("--deployment-target", default="iOS17", help="Minimum deployment target (iOS17, iOS16, macOS14, etc.)")
    
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)
    output_path = os.path.join(args.output_dir, "kompress.mlpackage")

    try:
        wrapper, mlmodel = export_to_coreml(
            checkpoint_path=args.checkpoint,
            output_path=output_path,
            max_seq_len=args.max_seq_len,
            precision=args.precision,
            minimum_deployment_target=args.deployment_target
        )
        validate_model(wrapper, mlmodel, args.max_seq_len)
        print("\n🎉 Export and validation completed successfully!")
    except Exception as e:
        print(f"\n❌ Export failed: {e}")
        return 1
        
    return 0

if __name__ == "__main__":
    import sys
    sys.exit(main())