#!/usr/bin/env python3
"""ModernBERT -> Core ML, at a fixed sequence length.

    python convert.py --length 128 [--model ModernBERT-base]

The length is fixed on purpose. A flexible `RangeDim` traced at one token
converts without error and is wrong at every other length — see the README.

Pin the toolchain. transformers 5.x emits `new_ones` in ModernBERT's attention
path and coremltools cannot convert it; the failure also hits BERT and
DistilBERT, so it is the toolchain and not this model.

    torch==2.7.0   transformers==4.48.3   coremltools==9.0
"""
import argparse
import numpy as np
import torch
import coremltools as ct
from transformers import AutoModelForMaskedLM

parser = argparse.ArgumentParser()
parser.add_argument("--model", default="ModernBERT-base")
parser.add_argument("--length", type=int, default=128)
args = parser.parse_args()


class Wrapped(torch.nn.Module):
    """The mask stays `ones_like` inside the graph.

    Taking `attention_mask` as a second input makes ModernBERT build it with
    `new_ones`, which coremltools has no conversion for. The cost is that
    padding is attended as if it were text — pick a length close to your
    sentences, or convert several.
    """

    def __init__(self, name):
        super().__init__()
        self.model = AutoModelForMaskedLM.from_pretrained(f"answerdotai/{name}")

    def forward(self, input_ids):
        return self.model(
            input_ids=input_ids, attention_mask=torch.ones_like(input_ids)
        ).logits


model = Wrapped(args.model).eval()
example = (torch.ones((1, args.length), dtype=torch.int32),)
traced = torch.jit.trace(model, example, strict=False)

mlmodel = ct.convert(
    traced,
    inputs=[ct.TensorType(name="input_ids", shape=(1, args.length), dtype=np.int32)],
    outputs=[ct.TensorType(name="logits")],
    minimum_deployment_target=ct.target.macOS15,
)
mlmodel.short_description = (
    f"{args.model}, masked-LM head, fixed sequence length {args.length}"
)
out = f"{args.model}-{args.length}.mlpackage"
mlmodel.save(out)
print("wrote", out)