| """Small real-fixture native/Core ML parity check for entity extraction.""" |
|
|
| 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.boundary.model import _group_scored_candidates |
| from gliner2.models.outputs import CandidateTensorBatch |
| from huggingface_hub import snapshot_download |
|
|
| from convert_extraction_names import FEATURE_NAMES, SCORE_INPUT_NAMES |
| from extraction_pool import select_candidates |
| from preprocessing import prepare_extraction |
|
|
| MODEL_ID = "fastino/gliner2.5-base-v1" |
| MODEL_REVISION = "1a8bc24e00dc7300b9017c81d63e3dcdabb26596" |
| FIXTURES = [ |
| ("Alice founded Acme in Toronto in 2020.", ["person", "organization", "location"]), |
| ("Apple acquired Beats for three billion dollars.", ["company", "product", "money"]), |
| ("Marie Curie was born in Warsaw and worked in Paris.", ["person", "city"]), |
| ] |
|
|
|
|
| def coreml_entity_result(native, features_model, scorer_model, text, schema, length, words, queries): |
| arrays, batch = prepare_extraction(native.processor, text, schema, length, words, queries) |
| features = features_model.predict(arrays) |
| tensor = {name: torch.from_numpy(np.asarray(features[name]).copy()) for name in FEATURE_NAMES} |
| text_mask = torch.from_numpy(arrays["text_mask"]).bool() |
| query_mask = torch.from_numpy(arrays["query_mask"]).bool() |
| head = native.boundary_head |
| pool = select_candidates( |
| tensor["pool_start_projection"], |
| tensor["pool_end_projection"], |
| tensor["boundary_mask"].bool(), |
| query_mask, |
| tensor["start_logits"], |
| tensor["end_logits"], |
| 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_values = ( |
| tensor["text_states"], |
| text_mask.float(), |
| tensor["query_states"], |
| query_mask.float(), |
| tensor["boundary_states"], |
| tensor["start_logits"], |
| tensor["end_logits"], |
| tensor["inside_prefix"], |
| tensor["inside_prefix_mean"], |
| pool.indices.int(), |
| pool.mask.float(), |
| pool.compat_logits, |
| ) |
| score_arrays = { |
| name: value.numpy().astype(np.int32 if name == "candidate_indices" else np.float32) |
| for name, value in zip(SCORE_INPUT_NAMES, score_values) |
| } |
| scores = scorer_model.predict(score_arrays) |
| candidate_batch = CandidateTensorBatch( |
| indices=pool.indices.unsqueeze(1).expand(1, queries, -1, 2), |
| proposal_logits=pool.proposal_logits.unsqueeze(1).expand(1, queries, -1), |
| pair_logits=torch.from_numpy(np.asarray(scores["pair_logits"]).copy()), |
| valid_mask=pool.mask.unsqueeze(1).expand(1, queries, -1), |
| query_mask=query_mask, |
| ) |
| _, metadata = native._build_schema_dicts_and_metadata([schema]) |
| with torch.no_grad(): |
| core = native._encode_core(batch) |
| native_output = native._extract_from_batch(batch, 0.5, metadata, True, True)[0] |
| probs = torch.sigmoid(candidate_batch.pair_logits / native.boundary_settings.pair_temperature) |
| grouped = _group_scored_candidates( |
| candidate_batch, |
| threshold=0.5, |
| probabilities=probs, |
| count_log_rates=tensor["count_log_rates"], |
| adaptive_threshold=native.boundary_settings.adaptive_threshold, |
| ) |
| extracted = native._decode_entities( |
| 0, |
| core, |
| core["ext_specs"][0], |
| grouped[0], |
| metadata[0], |
| torch.sigmoid(tensor["null_logits"])[0], |
| "flat", |
| core["word_offsets"][0], |
| batch.start_mappings[0], |
| batch.end_mappings[0], |
| text, |
| len(batch.start_mappings[0]), |
| True, |
| True, |
| ) |
| coreml_output = {"entities": [extracted]} if extracted else {} |
| return native_output, coreml_output, pool, candidate_batch |
|
|
|
|
| def stripped_spans(result): |
| output = {} |
| for name, entities in result.get("entities", [{}])[0].items(): |
| output[name] = [(item["text"], item["start"], item["end"]) for item in entities] |
| return output |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--model-dir", required=True) |
| parser.add_argument("--precision", choices=["fp16", "fp32"], default="fp32") |
| parser.add_argument("--length", type=int, default=128) |
| parser.add_argument("--max-words", type=int, default=64) |
| parser.add_argument("--max-queries", type=int, default=8) |
| args = parser.parse_args() |
| torch.set_num_threads(4) |
| source = snapshot_download(MODEL_ID, revision=MODEL_REVISION) |
| native = AutoExtractor.from_pretrained(str(source), map_location="cpu").eval() |
| suffix = f"{args.precision}_L{args.length}_W{args.max_words}_Q{args.max_queries}" |
| folder = Path(args.model_dir) |
| features = ct.models.MLModel( |
| str(folder / f"gliner2_base_extraction_features_{suffix}.mlpackage"), compute_units=ct.ComputeUnit.CPU_ONLY |
| ) |
| scorer = ct.models.MLModel( |
| str(folder / f"gliner2_base_extraction_scorer_{suffix}.mlpackage"), compute_units=ct.ComputeUnit.CPU_ONLY |
| ) |
| cases = [] |
| for text, labels in FIXTURES: |
| reference, actual, _, _ = coreml_entity_result( |
| native, |
| features, |
| scorer, |
| text, |
| Schema().entities(labels), |
| args.length, |
| args.max_words, |
| args.max_queries, |
| ) |
| expected_spans = stripped_spans(reference) |
| actual_spans = stripped_spans(actual) |
| cases.append( |
| { |
| "text": text, |
| "native_spans": expected_spans, |
| "coreml_spans": actual_spans, |
| "span_match": expected_spans == actual_spans, |
| } |
| ) |
| report = { |
| "source_model": MODEL_ID, |
| "source_revision": MODEL_REVISION, |
| "precision": args.precision, |
| "selected_manifest": "three fixed real-text entity fixtures", |
| "cases": cases, |
| "span_matches": sum(case["span_match"] for case in cases), |
| } |
| path = folder / f"verify-entities-{suffix}.json" |
| path.write_text(json.dumps(report, indent=2) + "\n") |
| print(json.dumps(report, indent=2)) |
| if report["span_matches"] != len(cases): |
| raise SystemExit(1) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
