scripts/validate_bina_coreml_transcripts.py

#!/usr/bin/env python3
"""End-to-end greedy transcript parity for the fixed 512px Bina CoreML runtime."""

from __future__ import annotations

import argparse
import hashlib
import importlib.util
import json
import time
from pathlib import Path

import coremltools as ct
import numpy as np
import torch
from PIL import Image
from qwen_vl_utils import process_vision_info
from transformers import AutoProcessor


EOS_TOKEN_ID = 2
PAD_TOKEN_ID = 0


def load_runtime_module(path: Path):
    spec = importlib.util.spec_from_file_location("bina_coreml_runtime_export", path)
    if spec is None or spec.loader is None:
        raise RuntimeError(f"Could not import runtime module from {path}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def letterbox(image: Image.Image, size: int) -> Image.Image:
    source = image.convert("RGB")
    source.thumbnail((size, size), Image.Resampling.LANCZOS)
    canvas = Image.new("RGB", (size, size), "white")
    canvas.paste(source, ((size - source.width) // 2, (size - source.height) // 2))
    return canvas


def build_inputs(processor, image: Image.Image, prompt: str):
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": prompt},
            ],
        }
    ]
    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    image_inputs, video_inputs = process_vision_info(messages)
    return processor(
        text=[text],
        images=image_inputs,
        videos=video_inputs,
        return_tensors="pt",
    )


def compute_units(name: str):
    return {
        "cpu_only": ct.ComputeUnit.CPU_ONLY,
        "cpu_and_gpu": ct.ComputeUnit.CPU_AND_GPU,
        "all": ct.ComputeUnit.ALL,
    }[name]


def prefix_match(left: list[int], right: list[int]) -> tuple[int, dict | None]:
    for index, (actual, expected) in enumerate(zip(left, right)):
        if actual != expected:
            return index, {"index": index, "coreml": actual, "reference": expected}
    matched = min(len(left), len(right))
    if len(left) != len(right):
        return matched, {
            "index": matched,
            "coreml": left[matched] if matched < len(left) else None,
            "reference": right[matched] if matched < len(right) else None,
        }
    return matched, None


def coreml_generate(
    rt,
    model,
    sample,
    vision_ml,
    prefill_ml,
    decode_ml,
    max_cache_length: int,
    max_new_tokens: int,
):
    started = time.perf_counter()
    pixel_values = sample["pixel_values"].detach().cpu().numpy().astype(np.float32)
    vision_outputs = vision_ml.predict({"pixel_values": pixel_values})
    image_embeds = torch.from_numpy(vision_outputs["image_embeds"])

    _, prefill_example = rt.build_prefill_example_from_image_embeds(model, sample, image_embeds)
    if int(prefill_example[0].shape[1]) >= max_cache_length:
        raise ValueError(
            f"prompt length {prefill_example[0].shape[1]} exceeds cache {max_cache_length}"
        )
    prefill_feed = {
        spec.name: tensor.detach().cpu().to(torch.float32).numpy()
        for spec, tensor in zip(rt.prefill_input_specs(prefill_example), prefill_example, strict=True)
    }
    prefill_outputs = prefill_ml.predict(prefill_feed)
    current_token = int(prefill_outputs["logits"][:, -1, :].argmax(axis=-1)[0])

    full_keys = [
        torch.from_numpy(prefill_outputs[f"full_key_{i}"])
        for i in range(rt.SuryaCoreMLPrefillFlat.full_layers)
    ]
    full_values = [
        torch.from_numpy(prefill_outputs[f"full_value_{i}"])
        for i in range(rt.SuryaCoreMLPrefillFlat.full_layers)
    ]
    conv_states = [
        torch.from_numpy(prefill_outputs[f"conv_state_{i}"])
        for i in range(rt.SuryaCoreMLPrefillFlat.linear_layers)
    ]
    recurrent_states = [
        torch.from_numpy(prefill_outputs[f"recurrent_state_{i}"])
        for i in range(rt.SuryaCoreMLPrefillFlat.linear_layers)
    ]

    generated = [current_token]
    cache_len = int(sample["input_ids"].shape[1])
    rope_delta = rt.sample_rope_delta(model, sample)
    stop_reason = "eos" if current_token == EOS_TOKEN_ID else "max_new_tokens"

    for _ in range(max_new_tokens - 1):
        if current_token == EOS_TOKEN_ID:
            stop_reason = "eos"
            break
        if cache_len >= max_cache_length:
            stop_reason = "cache_limit"
            break
        token = torch.tensor([[current_token]], dtype=torch.long)
        with torch.no_grad():
            inputs_embeds = model.model.language_model.embed_tokens(token)
            position_ids = torch.full((3, 1, 1), cache_len + rope_delta, dtype=torch.long)
            cos, sin = model.model.language_model.rotary_emb(inputs_embeds, position_ids)
        attention_mask = rt.decode_attention_mask(cache_len, max_cache_length)
        example = (inputs_embeds, cos, sin, attention_mask) + tuple(
            full_keys + full_values + conv_states + recurrent_states
        )
        feed = {
            spec.name: tensor.detach().cpu().to(torch.float32).numpy()
            for spec, tensor in zip(rt.decode_input_specs(example), example, strict=True)
        }
        outputs = decode_ml.predict(feed)
        next_token = int(outputs["logits"][:, -1, :].argmax(axis=-1)[0])

        for index in range(rt.SuryaCoreMLDecodeStepFlat.full_layers):
            full_keys[index][:, :, cache_len : cache_len + 1, :] = torch.from_numpy(
                outputs[f"new_full_key_{index}"]
            )
            full_values[index][:, :, cache_len : cache_len + 1, :] = torch.from_numpy(
                outputs[f"new_full_value_{index}"]
            )
        conv_states = [
            torch.from_numpy(outputs[f"new_conv_state_{i}"])
            for i in range(rt.SuryaCoreMLDecodeStepFlat.linear_layers)
        ]
        recurrent_states = [
            torch.from_numpy(outputs[f"new_recurrent_state_{i}"])
            for i in range(rt.SuryaCoreMLDecodeStepFlat.linear_layers)
        ]
        cache_len += 1
        current_token = next_token
        generated.append(next_token)

    return {
        "tokens": generated,
        "stop_reason": stop_reason,
        "prompt_tokens": int(sample["input_ids"].shape[1]),
        "mrope_position_delta": rope_delta,
        "seconds": time.perf_counter() - started,
    }


def reference_generate(model, sample, max_new_tokens: int):
    started = time.perf_counter()
    model.model.rope_deltas = None
    with torch.inference_mode():
        output = model.generate(
            **sample,
            do_sample=False,
            max_new_tokens=max_new_tokens,
            eos_token_id=EOS_TOKEN_ID,
            pad_token_id=PAD_TOKEN_ID,
            use_cache=True,
        )
    prompt_length = int(sample["input_ids"].shape[1])
    tokens = output[0, prompt_length:].detach().cpu().tolist()
    return {
        "tokens": [int(token) for token in tokens],
        "stop_reason": "eos" if tokens and tokens[-1] == EOS_TOKEN_ID else "max_new_tokens",
        "seconds": time.perf_counter() - started,
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-id", required=True)
    parser.add_argument("--runtime-script", type=Path, required=True)
    parser.add_argument("--vision-package", type=Path, required=True)
    parser.add_argument("--prefill-package", type=Path, required=True)
    parser.add_argument("--decode-package", type=Path, required=True)
    parser.add_argument("--image", type=Path, action="append", required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--canvas-size", type=int, default=512)
    parser.add_argument("--prompt-tokens", type=int, default=416)
    parser.add_argument("--max-cache-length", type=int, default=1024)
    parser.add_argument("--max-new-tokens", type=int, default=512)
    parser.add_argument(
        "--reference-dtype",
        choices=["bfloat16", "float32"],
        default="bfloat16",
    )
    parser.add_argument(
        "--compute-units",
        choices=["cpu_only", "cpu_and_gpu", "all"],
        default="cpu_only",
    )
    args = parser.parse_args()

    rt = load_runtime_module(args.runtime_script.expanduser().resolve())
    dtype = torch.bfloat16 if args.reference_dtype == "bfloat16" else torch.float32
    processor = AutoProcessor.from_pretrained(args.model_id, trust_remote_code=True)
    model = rt.load_model(args.model_id, dtype)
    units = compute_units(args.compute_units)
    vision_ml = ct.models.MLModel(str(args.vision_package.expanduser().resolve()), compute_units=units)
    prefill_ml = ct.models.MLModel(str(args.prefill_package.expanduser().resolve()), compute_units=units)
    decode_ml = ct.models.MLModel(str(args.decode_package.expanduser().resolve()), compute_units=units)

    rows = []
    for image_path in args.image:
        image_path = image_path.expanduser().resolve()
        original_bytes = image_path.read_bytes()
        with Image.open(image_path) as opened:
            source_size = list(opened.size)
            fixed = letterbox(opened, args.canvas_size)
        sample = build_inputs(processor, fixed, rt.PROMPT)
        if int(sample["input_ids"].shape[1]) != args.prompt_tokens:
            raise ValueError(
                f"{image_path.name}: expected fixed prompt length {args.prompt_tokens}, "
                f"got {sample['input_ids'].shape[1]}"
            )

        reference = reference_generate(model, sample, args.max_new_tokens)
        coreml = coreml_generate(
            rt,
            model,
            sample,
            vision_ml,
            prefill_ml,
            decode_ml,
            args.max_cache_length,
            args.max_new_tokens,
        )
        matched, mismatch = prefix_match(coreml["tokens"], reference["tokens"])
        coreml_text = processor.tokenizer.decode(
            coreml["tokens"], skip_special_tokens=True, clean_up_tokenization_spaces=False
        )
        reference_text = processor.tokenizer.decode(
            reference["tokens"], skip_special_tokens=True, clean_up_tokenization_spaces=False
        )
        row = {
            "image": image_path.name,
            "source_sha256": hashlib.sha256(original_bytes).hexdigest(),
            "source_size": source_size,
            "canvas_size": [args.canvas_size, args.canvas_size],
            "prompt_tokens": coreml["prompt_tokens"],
            "mrope_position_delta": coreml["mrope_position_delta"],
            "coreml_tokens": coreml["tokens"],
            "reference_tokens": reference["tokens"],
            "coreml_stop_reason": coreml["stop_reason"],
            "reference_stop_reason": reference["stop_reason"],
            "matched_prefix_tokens": matched,
            "token_exact": mismatch is None,
            "mismatch": mismatch,
            "coreml_text": coreml_text,
            "reference_text": reference_text,
            "transcript_exact": coreml_text == reference_text,
            "coreml_seconds": coreml["seconds"],
            "reference_seconds": reference["seconds"],
        }
        print(json.dumps(row, ensure_ascii=False), flush=True)
        rows.append(row)

    receipt = {
        "model_id": args.model_id,
        "source_precision": "bfloat16",
        "reference_execution_dtype": args.reference_dtype,
        "coreml_compute_precision": {
            "vision": "float32",
            "prefill": "float16",
            "decode": "float16",
        },
        "compute_units": args.compute_units,
        "canvas_size": args.canvas_size,
        "max_cache_length": args.max_cache_length,
        "eos_token_id": EOS_TOKEN_ID,
        "tested": len(rows),
        "token_exact": sum(row["token_exact"] for row in rows),
        "transcript_exact": sum(row["transcript_exact"] for row in rows),
        "all_eos": all(
            row["coreml_stop_reason"] == "eos" and row["reference_stop_reason"] == "eos"
            for row in rows
        ),
        "rows": rows,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(receipt, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(receipt, ensure_ascii=False, indent=2), flush=True)
    if (
        receipt["token_exact"] != receipt["tested"]
        or receipt["transcript_exact"] != receipt["tested"]
        or not receipt["all_eos"]
    ):
        raise SystemExit(1)


if __name__ == "__main__":
    main()