convert-relation-coreml.py

"""Export the trained GLiNER2.5 multilingual sparse relation scorer to Core ML."""

import argparse
import json
from pathlib import Path

import coremltools as ct
import numpy as np
import torch
from gliner2 import AutoExtractor, Schema
from gliner2.models.base import QueryLayout
from gliner2.training.trainer import ExtractorCollator
from huggingface_hub import snapshot_download

from extraction_export import ExtractionRelationExport

MODEL_ID = "fastino/gliner2.5-multi-v1"
MODEL_REVISION = "a221b77a8baf4a613b8f8652661d41fa10a5641e"
INPUT_NAMES = (
    "text_states",
    "text_length",
    "relation_states",
    "batch_index",
    "relation_index",
    "head_start",
    "head_end",
    "tail_start",
    "tail_end",
    "pair_mask",
)


def pad(value, size: int, fill=0):
    if value.shape[0] > size:
        raise ValueError(f"Relation fixture exceeds capacity {size}")
    output = value.new_full((size, *value.shape[1:]), fill)
    output[: value.shape[0]] = value
    return output


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", default="build/extraction")
    parser.add_argument("--precision", choices=["fp16", "fp32"], default="fp32")
    parser.add_argument("--max-words", type=int, default=64)
    parser.add_argument("--max-relations", type=int, default=4)
    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",
        ],
    )
    native = AutoExtractor.from_pretrained(source, map_location="cpu").eval()
    text = "Alice founded Acme in Toronto."
    schema = Schema().relations(["founded"])
    batch = ExtractorCollator(native.processor, is_training=False, max_len=None, architecture="boundary")(
        [(text, schema.build())]
    )
    with torch.no_grad():
        core = native._encode_core(batch)
        output = native.boundary_head(core["text_states"], core["text_mask"], core["query_states"], core["query_mask"])
        sample = native._single_sample_candidates(output.candidates, 0)
        relation_specs = core["rel_specs"][0]
        pairs = native.relation_pair_generator.generate_batched(
            sample, [QueryLayout(queries=())], [[entry["spec"] for entry in relation_specs]], compact=False
        )
        relation_states = torch.stack([entry["query_state"] for entry in relation_specs]).unsqueeze(0)
        native_scores = native.relation_scorer(core["text_states"], relation_states, sample, pairs)
    pair_cap = args.max_relations * native.boundary_settings.relation_pair_cap
    text_states = torch.zeros(1, args.max_words, core["text_states"].shape[-1])
    text_states[:, : core["text_states"].shape[1]] = core["text_states"]
    relation_padded = torch.zeros(1, args.max_relations, relation_states.shape[-1])
    relation_padded[:, : relation_states.shape[1]] = relation_states
    arguments = (
        text_states,
        torch.tensor([core["text_states"].shape[1]], dtype=torch.int32),
        relation_padded,
        pad(pairs.batch_index.int(), pair_cap),
        pad(pairs.relation_index.int(), pair_cap),
        pad(pairs.head_start.int(), pair_cap),
        pad(pairs.head_end.int(), pair_cap),
        pad(pairs.tail_start.int(), pair_cap),
        pad(pairs.tail_end.int(), pair_cap),
        pad(pairs.pair_mask.float(), pair_cap),
    )
    wrapper = ExtractionRelationExport(native).eval()
    with torch.no_grad():
        reference = wrapper(*arguments)
        wrapper_error = float((reference[: len(pairs)] - native_scores).abs().max())
        traced = torch.jit.trace(wrapper, arguments, check_trace=False)
    if wrapper_error > 1e-4:
        raise RuntimeError(f"Relation wrapper differs from native: {wrapper_error}")
    precision = ct.precision.FLOAT16 if args.precision == "fp16" else ct.precision.FLOAT32
    converted = ct.convert(
        traced,
        convert_to="mlprogram",
        minimum_deployment_target=ct.target.iOS17,
        compute_precision=precision,
        compute_units=ct.ComputeUnit.CPU_ONLY,
        inputs=[
            ct.TensorType(
                name=name,
                shape=tuple(value.shape),
                dtype=np.float32 if name in ("text_states", "relation_states", "pair_mask") else np.int32,
            )
            for name, value in zip(INPUT_NAMES, arguments)
        ],
        outputs=[ct.TensorType(name="relation_logits", dtype=np.float32)],
    )
    converted.short_description = "GLiNER2.5 multilingual trained sparse relation scoring head"
    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,
            "stage": "trained relation scorer",
            "word_capacity": str(args.max_words),
            "relation_capacity": str(args.max_relations),
            "pair_capacity": str(pair_cap),
        }
    )
    out = Path(args.output_dir)
    out.mkdir(parents=True, exist_ok=True)
    suffix = f"{args.precision}_W{args.max_words}_R{args.max_relations}_P{pair_cap}"
    package = out / f"gliner2_multi_relation_{suffix}.mlpackage"
    converted.save(str(package))
    model = ct.models.MLModel(str(package), compute_units=ct.ComputeUnit.CPU_ONLY)
    prediction = model.predict(
        {
            name: value.numpy().astype(
                np.float32 if name in ("text_states", "relation_states", "pair_mask") else np.int32
            )
            for name, value in zip(INPUT_NAMES, arguments)
        }
    )["relation_logits"]
    runtime_error = float(np.max(np.abs(prediction[: len(pairs)] - reference.numpy()[: len(pairs)])))
    if not np.isfinite(runtime_error):
        raise RuntimeError("Relation scorer produced non-finite values")
    report = {
        "source_model": MODEL_ID,
        "source_revision": MODEL_REVISION,
        "precision": args.precision,
        "fixture": text,
        "valid_pairs": int(pairs.pair_mask.sum()),
        "wrapper_max_absolute_error": wrapper_error,
        "coreml_max_absolute_error": runtime_error,
        "package": str(package),
        "package_bytes": sum(file.stat().st_size for file in package.rglob("*") if file.is_file()),
        "coremltools": ct.__version__,
        "torch": torch.__version__,
    }
    (out / f"relation-{suffix}.json").write_text(json.dumps(report, indent=2) + "\n")
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()