#!/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())