convert-extraction-coreml.py

"""Export the trained GLiNER2.5 multilingual boundary extraction stages to Core ML."""

import argparse
import json
import shutil
from pathlib import Path

import coremltools as ct
import numpy as np
import torch
from gliner2 import AutoExtractor, Schema
from huggingface_hub import snapshot_download

from convert_extraction_names import FEATURE_NAMES, SCORE_INPUT_NAMES
from extraction_export import ExtractionFeaturesExport, ExtractionScoreExport, coreml_trace_patches
from extraction_pool import select_candidates
from preprocessing import prepare_extraction

MODEL_ID = "fastino/gliner2.5-multi-v1"
MODEL_REVISION = "a221b77a8baf4a613b8f8652661d41fa10a5641e"
FIXTURE_TEXT = "Alice founded Acme in Toronto in 2020."
FIXTURE_SCHEMA = Schema().entities(["person", "organization", "location"])


def package_bytes(path: Path) -> int:
    return sum(file.stat().st_size for file in path.rglob("*") if file.is_file())


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", default="build/extraction")
    parser.add_argument("--length", type=int, default=128, help="Subword capacity")
    parser.add_argument("--max-words", type=int, default=64)
    parser.add_argument("--max-queries", type=int, default=8)
    parser.add_argument("--precision", choices=["fp16", "fp32"], default="fp32")
    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()
    arrays, batch = prepare_extraction(
        native.processor, FIXTURE_TEXT, FIXTURE_SCHEMA, args.length, args.max_words, args.max_queries
    )
    tensors = tuple(torch.from_numpy(value) for value in arrays.values())
    features_wrapper = ExtractionFeaturesExport(native).eval()
    with torch.no_grad(), coreml_trace_patches():
        features_reference = features_wrapper(*tensors)
        traced_features = torch.jit.trace(features_wrapper, tensors, check_trace=False)
    with torch.no_grad():
        native_core = native._encode_core(batch)
        valid_words = native_core["text_states"].shape[1]
        valid_queries = native_core["query_states"].shape[1]
        routing_error = max(
            float((features_reference[0][:, :valid_words] - native_core["text_states"]).abs().max()),
            float((features_reference[1][:, :valid_queries] - native_core["query_states"]).abs().max()),
        )
    if routing_error > 1e-4:
        raise RuntimeError(f"Traced routing differs from native: {routing_error}")

    precision = ct.precision.FLOAT16 if args.precision == "fp16" else ct.precision.FLOAT32
    input_names = tuple(arrays)
    features_model = ct.convert(
        traced_features,
        convert_to="mlprogram",
        minimum_deployment_target=ct.target.iOS17,
        compute_precision=precision,
        compute_units=ct.ComputeUnit.CPU_ONLY,
        inputs=[ct.TensorType(name=name, shape=arrays[name].shape, dtype=arrays[name].dtype) for name in input_names],
        outputs=[ct.TensorType(name=name, dtype=np.float32) for name in FEATURE_NAMES],
    )
    features_model.short_description = "GLiNER2.5 multilingual trained boundary extraction features"
    features_model.author = "Fastino (original); Fluid Inference (Core ML conversion)"
    features_model.license = "Apache-2.0"
    features_model.user_defined_metadata.update(
        {
            "source_model": MODEL_ID,
            "source_revision": MODEL_REVISION,
            "stage": "extraction features and trained boundary heads",
            "subword_capacity": str(args.length),
            "word_capacity": str(args.max_words),
            "query_capacity": str(args.max_queries),
        }
    )
    out = Path(args.output_dir)
    out.mkdir(parents=True, exist_ok=True)
    shutil.copy2(Path(source) / "config.json", out / "config.json")
    tokenizer_dir = out / "tokenizer"
    tokenizer_dir.mkdir(exist_ok=True)
    for name in ("tokenizer.json", "tokenizer_config.json"):
        shutil.copy2(Path(source) / name, tokenizer_dir / name)
    suffix = f"{args.precision}_L{args.length}_W{args.max_words}_Q{args.max_queries}"
    features_path = out / f"gliner2_multi_extraction_features_{suffix}.mlpackage"
    if features_path.exists():
        shutil.rmtree(features_path)
    features_model.save(str(features_path))
    print(f"Saved {features_path}", flush=True)

    head = native.boundary_head
    pooled = select_candidates(
        features_reference[8],
        features_reference[9],
        features_reference[3].bool(),
        tensors[5].bool(),
        features_reference[4],
        features_reference[5],
        boundary_top_k=head.shared_pool_builder.pool_boundary_top_k,
        pool_size=head.shared_pool_builder.pool_size,
        min_pool_per_query=head.shared_pool_builder.min_pool_per_query,
    )
    score_tensors = (
        features_reference[0],
        tensors[3],
        features_reference[1],
        tensors[5],
        features_reference[2],
        features_reference[4],
        features_reference[5],
        features_reference[6],
        features_reference[7],
        pooled.indices.int(),
        pooled.mask.float(),
        pooled.compat_logits,
    )
    scorer_wrapper = ExtractionScoreExport(native).eval()
    with torch.no_grad():
        scores_reference = scorer_wrapper(*score_tensors)
        traced_scores = torch.jit.trace(scorer_wrapper, score_tensors, check_trace=False)
    scorer_model = ct.convert(
        traced_scores,
        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.int32 if name == "candidate_indices" else np.float32
            )
            for name, value in zip(SCORE_INPUT_NAMES, score_tensors)
        ],
        outputs=[
            ct.TensorType(name="pair_logits", dtype=np.float32),
            ct.TensorType(name="candidate_states", dtype=np.float32),
        ],
    )
    scorer_model.short_description = "GLiNER2.5 multilingual trained shared-pool extraction scorer"
    scorer_model.author = "Fastino (original); Fluid Inference (Core ML conversion)"
    scorer_model.license = "Apache-2.0"
    scorer_model.user_defined_metadata.update(
        {
            "source_model": MODEL_ID,
            "source_revision": MODEL_REVISION,
            "stage": "trained extraction candidate scorer",
            "candidate_capacity": str(head.shared_pool_builder.pool_size),
        }
    )
    scorer_path = out / f"gliner2_multi_extraction_scorer_{suffix}.mlpackage"
    if scorer_path.exists():
        shutil.rmtree(scorer_path)
    scorer_model.save(str(scorer_path))
    print(f"Saved {scorer_path}", flush=True)

    # The first runtime check uses the same selected real fixture as the trace.
    runtime_features = ct.models.MLModel(str(features_path), compute_units=ct.ComputeUnit.CPU_ONLY)
    predicted_features = runtime_features.predict(arrays)
    errors = {
        name: float(np.max(np.abs(np.asarray(predicted_features[name]) - reference.detach().numpy())))
        for name, reference in zip(FEATURE_NAMES, features_reference)
    }
    if any(not np.isfinite(value) for value in errors.values()):
        raise RuntimeError("Extraction features contain non-finite values")
    runtime_scorer = ct.models.MLModel(str(scorer_path), compute_units=ct.ComputeUnit.CPU_ONLY)
    score_arrays = {
        name: value.detach().numpy().astype(np.int32 if name == "candidate_indices" else np.float32)
        for name, value in zip(SCORE_INPUT_NAMES, score_tensors)
    }
    predicted_scores = runtime_scorer.predict(score_arrays)
    errors["pair_logits"] = float(np.max(np.abs(predicted_scores["pair_logits"] - scores_reference[0].numpy())))
    errors["candidate_states"] = float(
        np.max(np.abs(predicted_scores["candidate_states"] - scores_reference[1].numpy()))
    )
    if any(not np.isfinite(value) for value in errors.values()):
        raise RuntimeError("Extraction scorer contains non-finite values")
    report = {
        "source_model": MODEL_ID,
        "source_revision": MODEL_REVISION,
        "precision": args.precision,
        "fixture": FIXTURE_TEXT,
        "shape": {
            "subwords": args.length,
            "words": args.max_words,
            "queries": args.max_queries,
            "candidates": head.shared_pool_builder.pool_size,
        },
        "routing_max_absolute_error": routing_error,
        "runtime_max_absolute_errors": errors,
        "packages": {
            "features": {"path": str(features_path), "bytes": package_bytes(features_path)},
            "scorer": {"path": str(scorer_path), "bytes": package_bytes(scorer_path)},
        },
        "coremltools": ct.__version__,
        "torch": torch.__version__,
    }
    (out / f"conversion-{suffix}.json").write_text(json.dumps(report, indent=2) + "\n")
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()