| |
| """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) |
|
|
