File size: 6,714 Bytes
fcf4209 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | """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()
|