#!/usr/bin/env python3
"""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}")