File size: 7,123 Bytes
0a423df | 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 | """Convert and verify the pinned GLiNER2.5-multi classification decision 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 export_model import GLiNER2ClassificationExport, coreml_safe_attention_forward
from preprocessing import native_batch, prepare_classification
MODEL_ID = "fastino/gliner2.5-multi-v1"
MODEL_REVISION = "a221b77a8baf4a613b8f8652661d41fa10a5641e"
EXAMPLES = [
("The rocket launched successfully.", "topic", ["science", "sports", "politics"]),
("The team won the football championship.", "topic", ["science", "sports", "politics"]),
("The budget was approved by parliament.", "topic", ["science", "sports", "politics"]),
]
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--output-dir", default="build")
parser.add_argument("--length", type=int, default=128)
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"
],
)
native = AutoExtractor.from_pretrained(source, map_location="cpu").eval()
wrapper = GLiNER2ClassificationExport(native).eval()
text, task, labels = EXAMPLES[0]
arrays = prepare_classification(native, text, task, labels, args.length, args.max_options)
tensors = tuple(torch.from_numpy(value) for value in arrays.values())
with torch.no_grad():
batch = native_batch(native, text, task, labels, args.length)
core = native._encode_core(batch)
expected = native.classifier(core["cls_specs"][0][0]["choice_states"]).squeeze(-1)
actual = wrapper(*tensors)[0][0, : len(labels)]
wrapper_error = float((expected - actual).abs().max())
if wrapper_error > 1e-4:
raise RuntimeError(f"Wrapper/native logit mismatch: {wrapper_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.
# Freeze that branch to avoid a Core ML conditional with mismatched ranks.
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:
with torch.no_grad():
frozen = wrapper(*tensors)[0][0, : len(labels)]
frozen_error = float((expected - frozen).abs().max())
if frozen_error > 1e-4:
raise RuntimeError(f"Frozen attention scale changed native logits: {frozen_error}")
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
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=(1, args.max_options), dtype=np.int32),
ct.TensorType(name="marker_mask", shape=(1, args.max_options), dtype=np.float32),
],
outputs=[ct.TensorType(name="logits", dtype=np.float32), ct.TensorType(name="probabilities", dtype=np.float32)],
)
converted.short_description = "GLiNER2.5-multi native 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; entity/relation/record extraction heads not exported",
"length": str(args.length), "max_options": str(args.max_options),
})
out = Path(args.output_dir)
out.mkdir(parents=True, exist_ok=True)
package = out / f"gliner2_multi_classification_{args.precision}_L{args.length}_K{args.max_options}.mlpackage"
converted.save(str(package))
runtime = ct.models.MLModel(str(package), compute_units=ct.ComputeUnit.ALL)
cases = []
for text, task, labels in EXAMPLES:
arrays = prepare_classification(native, text, task, labels, args.length, args.max_options)
native_output = native.classify_text(text, {task: labels}, include_confidence=True, max_len=args.length)[task]
prediction = runtime.predict(arrays)
scores = np.asarray(prediction["probabilities"])[0, : len(labels)]
choice = labels[int(scores.argmax())]
if choice != native_output["label"]:
raise RuntimeError(f"Core ML/native choice mismatch: {choice} != {native_output['label']}")
cases.append({
"text": text, "native_label": native_output["label"], "coreml_label": choice,
"native_confidence": native_output["confidence"], "coreml_confidence": float(scores.max()),
"absolute_confidence_error": abs(float(scores.max()) - native_output["confidence"]),
})
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": wrapper_error, "coremltools": ct.__version__,
"torch": torch.__version__, "cases": cases,
}
(out / "conversion.json").write_text(json.dumps(report, indent=2) + "\n")
print(json.dumps(report, indent=2))
if __name__ == "__main__":
main()
|