convert-coreml.py

"""Convert and verify the pinned GLiNER2.5-Decide multi-head classification path."""
import argparse
import json
import math
from pathlib import Path

import coremltools as ct
import numpy as np
import torch
from gliner2 import AutoExtractor
from huggingface_hub import snapshot_download
from transformers.models.deberta_v2 import modeling_deberta_v2

from convert_names import MODEL_ID, MODEL_REVISION
from export_model import GLiNER2DecideExport, coreml_safe_attention_forward
from preprocessing import classification_schema, prepare_decision
from runtime import decode, package_name

EXAMPLES = [
    ("My subscription renewed on April 15 for ¥5,400 after the service was already down. "
     "Can I get that charge refunded?",
     {"intent": ["order_status", "refund_request", "cancel_subscription", "update_payment", "login_problem",
                 "shipping_delay", "bug_report", "speak_to_human"]}),
    ("Battery dies before lunch, but the keyboard and the screen are the best I have used on a laptop.",
     {"sentiment": ["positive", "negative", "mixed", "neutral"],
      "aspects": {"labels": ["battery", "keyboard", "screen", "camera", "price", "support"],
                  "multi_label": True, "cls_threshold": 0.4}}),
    ("From: compliance@group.example\nSubject: Protocol update — action required today\n\n"
     "Please confirm the new retention rule is applied before Friday's audit.",
     {"intent": ["fyi", "request", "approval", "complaint", "newsletter", "security_alert"],
      "urgency": ["low", "normal", "high", "critical"],
      "route": ["support", "billing", "legal", "security", "finance", "archive"]}),
]


def native_logits(native, text, tasks):
    """Per-head native logits via the upstream span collator, encoder and shared classifier."""
    from gliner2.training.trainer import ExtractorCollator

    collator = ExtractorCollator(native.processor, is_training=False, max_len=None, architecture="span")
    batch = collator([(text, classification_schema(tasks).build())])
    _, schema_embs = native._encode_batch(batch)
    return [native.classifier(torch.stack(embs[1:])).squeeze(-1) for embs in schema_embs[0]]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", default="build")
    parser.add_argument("--length", type=int, default=128)
    parser.add_argument("--max-heads", type=int, default=4)
    parser.add_argument("--max-options", type=int, default=8)
    parser.add_argument("--precision", choices=["fp16", "fp32"], default="fp16")
    args = parser.parse_args()
    torch.set_num_threads(4)
    source = snapshot_download(
        MODEL_ID, revision=MODEL_REVISION,
        allow_patterns=[
            "config.json", "encoder_config/*", "model.safetensors", "tokenizer.json", "tokenizer_config.json",
            "special_tokens_map.json",
        ],
    )
    native = AutoExtractor.from_pretrained(source, map_location="cpu").eval()
    wrapper = GLiNER2DecideExport(native).eval()
    bucket = (args.length, args.max_heads, args.max_options)
    text, tasks = EXAMPLES[2]
    arrays = prepare_decision(native.processor, text, tasks, *bucket)
    tensors = tuple(torch.from_numpy(value) for value in arrays.values())

    def wrapper_error():
        with torch.no_grad():
            expected = native_logits(native, text, tasks)
            actual = wrapper(*tensors)[0][0]
        return max(float((row - actual[h, : len(row)]).abs().max()) for h, row in enumerate(expected))

    error = wrapper_error()
    if error > 1e-4:
        raise RuntimeError(f"Wrapper/native logit mismatch: {error}")
    # The upstream scale is a constant for a fixed DeBERTa attention head width.
    # Its traced int32 sqrt is rejected by Core ML; freeze the identical float32
    # value while tracing, and restore the upstream implementation immediately.
    original_scale = modeling_deberta_v2.scaled_size_sqrt
    original_rpos = modeling_deberta_v2.build_rpos
    original_attention = modeling_deberta_v2.DisentangledSelfAttention.forward

    def static_scale(query_layer, scale_factor):
        value = math.sqrt(float(query_layer.shape[-1] * scale_factor))
        return torch.tensor(value, dtype=torch.float32, device=query_layer.device)

    modeling_deberta_v2.scaled_size_sqrt = static_scale
    # The encoder only uses self-attention: query and key sequence lengths are
    # identical, so the scripted build_rpos returns relative_pos unchanged.
    modeling_deberta_v2.build_rpos = lambda query, key, relative_pos, buckets, max_pos: relative_pos
    modeling_deberta_v2.DisentangledSelfAttention.forward = coreml_safe_attention_forward
    try:
        frozen_error = wrapper_error()
        if frozen_error > 1e-4:
            raise RuntimeError(f"Frozen attention scale changed native logits: {frozen_error}")
        with torch.no_grad():
            traced = torch.jit.trace(wrapper, tensors, check_trace=False)
    finally:
        modeling_deberta_v2.scaled_size_sqrt = original_scale
        modeling_deberta_v2.build_rpos = original_rpos
        modeling_deberta_v2.DisentangledSelfAttention.forward = original_attention
    grid = (1, args.max_heads, args.max_options)
    converted = ct.convert(
        traced, convert_to="mlprogram", minimum_deployment_target=ct.target.iOS17,
        compute_precision=ct.precision.FLOAT16 if args.precision == "fp16" else ct.precision.FLOAT32,
        compute_units=ct.ComputeUnit.CPU_ONLY,
        inputs=[
            ct.TensorType(name="input_ids", shape=(1, args.length), dtype=np.int32),
            ct.TensorType(name="attention_mask", shape=(1, args.length), dtype=np.int32),
            ct.TensorType(name="marker_indices", shape=grid, dtype=np.int32),
            ct.TensorType(name="marker_mask", shape=grid, dtype=np.float32),
        ],
        outputs=[ct.TensorType(name="logits", dtype=np.float32), ct.TensorType(name="probabilities", dtype=np.float32)],
    )
    converted.short_description = "GLiNER2.5-Decide multi-head schema classification path"
    converted.author = "Fastino (original); Fluid Inference (Core ML conversion)"
    converted.license = "Apache-2.0"
    converted.user_defined_metadata.update({
        "source_model": MODEL_ID, "source_revision": MODEL_REVISION,
        "scope": "classification only; span and count heads not exported",
        "length": str(args.length), "max_heads": str(args.max_heads), "max_options": str(args.max_options),
    })
    out = Path(args.output_dir)
    out.mkdir(parents=True, exist_ok=True)
    package = out / package_name(args.precision, *bucket)
    converted.save(str(package))
    runtime = ct.models.MLModel(str(package), compute_units=ct.ComputeUnit.ALL)
    cases = []
    for text, tasks in EXAMPLES:
        native_output = native.classify_text(text, tasks, include_confidence=True)
        logits = np.asarray(runtime.predict(prepare_decision(native.processor, text, tasks, *bucket))["logits"])[0]
        coreml_output = decode(tasks, logits)
        if json.dumps(labels_only(native_output)) != json.dumps(labels_only(coreml_output)):
            raise RuntimeError(f"Core ML/native mismatch: {coreml_output} != {native_output}")
        cases.append({"text": text, "native": native_output, "coreml": coreml_output,
                      "max_confidence_error": confidence_error(native_output, coreml_output)})
    report = {
        "source_model": MODEL_ID, "source_revision": MODEL_REVISION, "package": str(package),
        "package_bytes": sum(f.stat().st_size for f in package.rglob("*") if f.is_file()),
        "native_total_parameters": sum(p.numel() for p in native.parameters()),
        "exported_parameters": sum(p.numel() for p in wrapper.parameters()),
        "wrapper_max_logit_error": error, "coremltools": ct.__version__, "torch": torch.__version__,
        "cases": cases,
    }
    (out / f"conversion-{package.stem}.json").write_text(json.dumps(report, indent=2, ensure_ascii=False) + "\n")
    print(json.dumps(report, indent=2, ensure_ascii=False))


def _entries(value):
    return value if isinstance(value, list) else [value]


def labels_only(result: dict) -> dict:
    return {task: sorted(entry["label"] for entry in _entries(value)) for task, value in result.items()}


def confidence_error(native: dict, coreml: dict) -> float:
    errors = []
    for task, value in native.items():
        predicted = {entry["label"]: entry["confidence"] for entry in _entries(coreml[task])}
        errors += [abs(entry["confidence"] - predicted[entry["label"]]) for entry in _entries(value)]
    return max(errors)


if __name__ == "__main__":
    main()