| |
| """Check a converted model against the PyTorch original. |
| |
| python verify.py ModernBERT-base-128.mlpackage --length 128 |
| |
| Prints the correlation of the logits at a masked position. Anything below |
| 0.99 means the conversion is wrong, whatever its top predictions look like. |
| """ |
| import argparse |
| import numpy as np |
| import torch |
| import coremltools as ct |
| from transformers import AutoTokenizer, AutoModelForMaskedLM |
|
|
| parser = argparse.ArgumentParser() |
| parser.add_argument("package") |
| parser.add_argument("--length", type=int, default=128) |
| parser.add_argument("--model", default="answerdotai/ModernBERT-base") |
| args = parser.parse_args() |
|
|
| tok = AutoTokenizer.from_pretrained(args.model) |
| ref = AutoModelForMaskedLM.from_pretrained(args.model).eval() |
| mlmodel = ct.models.MLModel(args.package) |
|
|
| SENTENCES = [ |
| "The capital of Ireland is {m}.", |
| "You should see a parrot {m} at the top right of your screen.", |
| "The PR looks good but there is another failure mode {m} we might cover this.", |
| ] |
|
|
| for sentence in SENTENCES: |
| text = sentence.format(m=tok.mask_token) |
| encoded = tok(text, return_tensors="pt") |
| length = encoded.input_ids.shape[1] |
| ids = np.full((1, args.length), tok.pad_token_id, np.int32) |
| ids[0, :length] = encoded.input_ids.numpy() |
| padded = torch.tensor(ids, dtype=torch.long) |
| with torch.no_grad(): |
| expected = ref( |
| input_ids=padded, attention_mask=torch.ones_like(padded) |
| ).logits[0].numpy() |
| got = mlmodel.predict({"input_ids": ids})["logits"][0] |
| at = int(np.where(encoded.input_ids[0].numpy() == tok.mask_token_id)[0][0]) |
| print(f"\n {text[:56]}") |
| print(f" torch {[tok.decode([int(i)]) for i in np.argsort(-expected[at])[:4]]}") |
| print(f" coreml {[tok.decode([int(i)]) for i in np.argsort(-got[at])[:4]]}") |
| print(f" r = {np.corrcoef(expected[at], got[at])[0, 1]:.5f}") |
|
|
