"""GLiNER2.5 multilingual extraction runtime using only Core ML trained weights.""" import gc import json import threading from contextvars import ContextVar from pathlib import Path import coremltools as ct import numpy as np import torch from gliner2.configuration import BoundaryHeadSettings from gliner2.models.boundary.engine import BoundaryExtractor from gliner2.models.boundary.records import RecordGroupOutput from gliner2.models.boundary.relations import ( RelationProposalSettings, RelationTypeSpec, TypedRelationPairGenerator, ) from gliner2.models.outputs import CandidateTensorBatch, ExtractorOutput from convert_extraction_names import FEATURE_NAMES, SCORE_INPUT_NAMES from extraction_pool import select_candidates from preprocessing import load_processor, prepare_extraction MODEL_PREFIX = "gliner2_multi" def to_tensor(value): return torch.from_numpy(np.asarray(value).copy()) def padded(value, size): if value.shape[0] > size: raise ValueError(f"Request exceeds Core ML bucket capacity {size}") result = value.new_zeros((size, *value.shape[1:])) result[: value.shape[0]] = value return result class CoreMLBoundaryHead(torch.nn.Module): """Provide native decoder tensors from the converted extraction stages.""" def __init__(self, context: ContextVar, explicit_model, max_queries: int, max_spans: int): super().__init__() self.context = context self.explicit_model = explicit_model self.max_queries = max_queries self.max_spans = max_spans def forward(self, text, text_mask, query, query_mask, return_candidates=True): state = self.context.get() features = state["features"] full = state["candidates"] count = query.shape[1] selected = CandidateTensorBatch( indices=full.indices[:, :count], proposal_logits=full.proposal_logits[:, :count], pair_logits=full.pair_logits[:, :count], valid_mask=full.valid_mask[:, :count], query_mask=full.query_mask[:, :count], candidate_states=full.candidate_states[:, :count], ) return ExtractorOutput( candidates=selected if return_candidates else None, start_logits=features["start_logits"][:, :count], end_logits=features["end_logits"][:, :count], null_logits=features["null_logits"][:, :count], count_log_rates=features["count_log_rates"][:, :count], batch_size=1, ) def score_explicit_spans(self, text, text_mask, query, query_mask, indices, valid_mask=None): state = self.context.get() features = state["features"] full_queries = features["query_states"][0] active_queries = int(state["arrays"]["query_mask"].sum()) query_count = query.shape[1] span_count = indices.shape[2] if query_count > self.max_queries or span_count > self.max_spans: raise ValueError("Explicit-span request exceeds Core ML bucket capacity") selected = [] for row in query[0]: equal = torch.isclose(full_queries[:active_queries], row, atol=1e-6, rtol=0).all(-1) matches = equal.nonzero(as_tuple=False).flatten() if matches.numel() != 1: raise ValueError("Explicit-span query cannot be mapped to the encoded schema") selected.append(int(matches[0])) selection = torch.tensor(selected, dtype=torch.long) query_states = padded(query[0], self.max_queries).unsqueeze(0) query_mask_padded = padded(query_mask[0].float(), self.max_queries).unsqueeze(0) def selected_feature(name): source = features[name][0].index_select(0, selection) return padded(source, self.max_queries).unsqueeze(0) span_indices = torch.zeros(1, self.max_queries, self.max_spans, 2, dtype=torch.int32) span_mask = torch.zeros(1, self.max_queries, self.max_spans, dtype=torch.float32) span_indices[:, :query_count, :span_count] = indices.int() span_mask[:, :query_count, :span_count] = valid_mask.float() if valid_mask is not None else 1.0 values = ( text, text_mask.float(), query_states, query_mask_padded, features["boundary_states"], selected_feature("start_logits"), selected_feature("end_logits"), selected_feature("inside_prefix"), selected_feature("inside_prefix_mean"), span_indices, span_mask, ) names = ( "text_states", "text_mask", "query_states", "query_mask", "boundary_states", "start_logits", "end_logits", "inside_prefix", "inside_prefix_mean", "span_indices", "span_mask", ) output = self.explicit_model.predict( { name: value.numpy().astype(np.int32 if name == "span_indices" else np.float32) for name, value in zip(names, values) } )["span_logits"] return to_tensor(output)[:, :query_count, :span_count] class CoreMLRelationScorer(torch.nn.Module): """Call the trained Core ML relation graph after native pair selection.""" def __init__(self, context: ContextVar, model, max_words: int, max_relations: int, pair_cap: int): super().__init__() self.context = context self.model = model self.max_words = max_words self.max_relations = max_relations self.pair_cap = pair_cap def forward(self, text, relation, candidates, pairs): count = len(pairs) if text.shape[1] != self.max_words or relation.shape[1] > self.max_relations or count > self.pair_cap: raise ValueError("Relation request exceeds Core ML bucket capacity") relation_states = torch.zeros(1, self.max_relations, relation.shape[-1], dtype=relation.dtype) relation_states[:, : relation.shape[1]] = relation state = self.context.get() text_length = int(state["arrays"]["text_mask"].sum()) fields = ( text, torch.tensor([text_length], dtype=torch.int32), relation_states, padded(pairs.batch_index.int(), self.pair_cap), padded(pairs.relation_index.int(), self.pair_cap), padded(pairs.head_start.int(), self.pair_cap), padded(pairs.head_end.int(), self.pair_cap), padded(pairs.tail_start.int(), self.pair_cap), padded(pairs.tail_end.int(), self.pair_cap), padded(pairs.pair_mask.float(), self.pair_cap), ) names = ( "text_states", "text_length", "relation_states", "batch_index", "relation_index", "head_start", "head_end", "tail_start", "tail_end", "pair_mask", ) output = self.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(names, fields) } )["relation_logits"] return to_tensor(output)[:count] class CoreMLRecordHead(torch.nn.Module): """Keep GLiNER2's instance ordering while running learned layers in Core ML.""" def __init__(self, assignment_model, anchorless_model, max_fields=8, max_candidates=192, max_instances=1536): super().__init__() self.assignment_model = assignment_model self.anchorless_model = anchorless_model self.max_fields = max_fields self.max_candidates = max_candidates self.max_instances = max_instances def forward_group(self, spec, query_states, candidates, sample_index): field_specs = list(spec.fields) field_query_ids = [field.query_id for field in field_specs] if len(field_specs) > self.max_fields: raise ValueError("Record has more fields than the Core ML bucket") query_count = query_states.shape[0] if any(query_id < 0 or query_id >= query_count for query_id in field_query_ids): raise ValueError("Record field query is outside the encoded schema") field_states, field_spans, field_logits, field_masks = [], [], [], [] for query_id in field_query_ids: mask = candidates.valid_mask[sample_index, query_id] if int(mask.sum()) > self.max_candidates: raise ValueError("Record candidate count exceeds Core ML bucket") field_states.append(candidates.candidate_states[sample_index, query_id][mask]) field_spans.append(candidates.indices[sample_index, query_id][mask]) field_logits.append(candidates.pair_logits[sample_index, query_id][mask]) field_masks.append(torch.ones(int(mask.sum()), dtype=torch.bool)) instance_seed = [] instance_spans = [] if spec.mode == "natural": anchor = field_query_ids.index(spec.anchor_query_id) instances = field_states[anchor] for index, span in enumerate(field_spans[anchor]): instance_seed.append((anchor, index)) instance_spans.append((int(span[0]), int(span[1]))) elif spec.mode == "latent": instances = ( torch.cat(field_states, 0) if field_states else query_states.new_zeros((0, query_states.shape[-1])) ) for field_index, spans in enumerate(field_spans): for index, span in enumerate(spans): instance_seed.append((field_index, index)) instance_spans.append((int(span[0]), int(span[1]))) else: context_states = ( torch.cat(field_states, 0) if field_states else query_states.new_zeros((0, query_states.shape[-1])) ) context_capacity = self.max_fields * self.max_candidates context_mask = torch.zeros(context_capacity, dtype=torch.float32) context_mask[: context_states.shape[0]] = 1.0 output = self.anchorless_model.predict( { "context_states": padded(context_states, context_capacity).numpy().astype(np.float32), "context_mask": context_mask.numpy(), } )["instance_states"] instances = to_tensor(output) instance_seed = [None] * instances.shape[0] instance_spans = [None] * instances.shape[0] count = instances.shape[0] if count > self.max_instances: raise ValueError("Record instance count exceeds Core ML bucket") hidden = query_states.shape[-1] candidate_states = torch.zeros(self.max_fields, self.max_candidates, hidden) for index, states in enumerate(field_states): candidate_states[index, : states.shape[0]] = states fields = ( padded(instances, self.max_instances), padded(query_states[field_query_ids], self.max_fields), candidate_states, ) predicted = self.assignment_model.predict( { name: value.numpy().astype(np.float32) for name, value in zip(("instance_states", "field_queries", "field_candidate_states"), fields) } ) assignment = to_tensor(predicted["assignment_logits"]) assignment_by_field = [ assignment[:count, field_index, : 1 + states.shape[0]] for field_index, states in enumerate(field_states) ] if spec.mode == "natural": object_logits = field_logits[field_query_ids.index(spec.anchor_query_id)] elif spec.mode == "latent": object_logits = to_tensor(predicted["latent_seed_logits"])[:count] else: object_logits = to_tensor(predicted["object_logits"])[:count] return RecordGroupOutput( spec=spec, object_logits=object_logits, assign_logits=assignment_by_field, field_query_ids=field_query_ids, field_specs=field_specs, field_spans=field_spans, field_cand_mask=field_masks, field_cand_logits=field_logits, instance_seed=instance_seed, instance_spans=instance_spans, ) class CoreMLBoundaryExtractor(BoundaryExtractor): """Native GLiNER2 schema/decoder with all trained extraction heads in Core ML.""" def __init__( self, model_dir: str, precision: str = "fp32", compute_units=ct.ComputeUnit.CPU_ONLY, feature_package: str | None = None, ): if precision not in ("fp16", "fp32"): raise ValueError("precision must be fp16 or fp32") torch.nn.Module.__init__(self) folder = Path(model_dir) config = json.loads((folder / "config.json").read_text()) self.boundary_settings = BoundaryHeadSettings(**config["boundary_head"]) self.processor = load_processor(str(folder / "tokenizer")) self.enable_records = self.boundary_settings.enable_records self.enable_relations = self.boundary_settings.enable_relations self.strict_extraction = True self.length, self.max_words, self.max_queries = 128, 64, 8 self._context = ContextVar("gliner2_coreml_extraction_context") suffix = f"{precision}_L128_W64_Q8" feature_path = ( Path(feature_package) if feature_package else Path(f"{MODEL_PREFIX}_extraction_features_{suffix}.mlpackage") ) if not feature_path.is_absolute(): feature_path = folder / feature_path self.features_model = ct.models.MLModel(str(feature_path), compute_units=compute_units) self.scorer_model = ct.models.MLModel( str(folder / f"{MODEL_PREFIX}_extraction_scorer_{suffix}.mlpackage"), compute_units=compute_units ) explicit = ct.models.MLModel( str(folder / f"{MODEL_PREFIX}_explicit_{precision}_W64_Q8_S64.mlpackage"), compute_units=compute_units ) relation = ct.models.MLModel( str(folder / f"{MODEL_PREFIX}_relation_{precision}_W64_R4_P256.mlpackage"), compute_units=compute_units ) assignment = ct.models.MLModel( str(folder / f"{MODEL_PREFIX}_record_assignment_{precision}_F8_C192_I1536.mlpackage"), compute_units=compute_units, ) anchorless = ct.models.MLModel( str(folder / f"{MODEL_PREFIX}_record_anchorless_{precision}_F8_C192_I1536.mlpackage"), compute_units=compute_units, ) self.boundary_head = CoreMLBoundaryHead(self._context, explicit, self.max_queries, 64) self.relation_scorer = CoreMLRelationScorer(self._context, relation, self.max_words, 4, 256) self.record_decoder = CoreMLRecordHead(assignment, anchorless) self.relation_pair_generator = TypedRelationPairGenerator( RelationProposalSettings( heads_per_relation=self.boundary_settings.relation_heads_per_type, tails_per_relation=self.boundary_settings.relation_tails_per_type, pair_cap=self.boundary_settings.relation_pair_cap, argument_threshold=self.boundary_settings.relation_argument_proposal_threshold, ) ) def _encode_core(self, batch): features = self._context.get()["features"] query_states = features["query_states"] text_states = features["text_states"] query_mask = to_tensor(self._context.get()["arrays"]["query_mask"]).bool() query_count = int(query_mask.sum()) query_states = query_states[:, :query_count] query_mask = query_mask[:, :query_count] text_mask = to_tensor(self._context.get()["arrays"]["text_mask"]).bool() ext_specs, cls_specs, rel_specs, word_offsets = [], [], [], [] for sample_index in range(len(batch)): specs = [ { "group_index": item.task_index, "field_index": item.role_index, "task_type": item.task_type, "task_name": item.task_name, "field_name": item.role_name, } for item in batch.query_layouts[sample_index].queries ] ext_specs.append(specs) classifications = [] choice_offset = 0 for group_index in range(batch.schema_counts[sample_index]): if batch.task_types[sample_index][group_index] != "classifications": continue count = max(len(batch.schema_special_indices[sample_index][group_index]) - 1, 0) schema_tokens = batch.schema_tokens_list[sample_index][group_index] if count: classifications.append( { "group_index": group_index, "task_name": schema_tokens[2], "schema_tokens": schema_tokens, "group_embs": features["classification_logits"][ sample_index, choice_offset : choice_offset + count ], } ) choice_offset += count cls_specs.append(classifications) word_offsets.append( max(int(batch.text_word_counts[sample_index]) - len(batch.start_mappings[sample_index]), 0) ) groups = {} for query_id, spec in enumerate(specs): groups.setdefault(spec["group_index"], []).append(query_id) relations = [] for group_index, role_ids in groups.items(): if batch.task_types[sample_index][group_index] != "relations" or len(role_ids) < 2: continue head_id, tail_id = role_ids[:2] role_states = query_states[sample_index, [head_id, tail_id]] state = ( torch.cat((role_states[0], role_states[1]), -1) if self.boundary_settings.directional_relation_states else role_states.mean(0) ) relations.append( { "group_index": group_index, "relation_type": specs[head_id]["task_name"], "spec": RelationTypeSpec( specs[head_id]["task_name"], head_query_ids=(head_id,), tail_query_ids=(tail_id,) ), "query_state": state, } ) rel_specs.append(relations) return { "text_states": text_states, "text_mask": text_mask, "text_lengths": text_mask.sum(-1).long(), "query_states": query_states, "query_mask": query_mask, "ext_specs": ext_specs, "cls_specs": cls_specs, "rel_specs": rel_specs, "word_offsets": word_offsets, } def _extract_classification_result(self, results, schema_name, schema, embs, schema_tokens, temperature=1.0): cls_config = self._resolve_classification_config(schema_tokens[2], schema.get("classifications", [])) if cls_config is None: return if temperature <= 0: raise ValueError("Classification temperature must be positive") logits = embs / temperature activation = cls_config.get("class_act", "auto") multi_label = cls_config.get("multi_label", False) if activation == "sigmoid" or (activation != "softmax" and multi_label): probabilities = torch.sigmoid(logits) else: probabilities = torch.softmax(logits, dim=-1) labels = cls_config["labels"] if multi_label: threshold = cls_config.get("cls_threshold", 0.5) chosen = [ (labels[index], float(probabilities[index])) for index in range(len(labels)) if float(probabilities[index]) >= threshold ] if not chosen: best = int(probabilities.argmax()) chosen = [(labels[best], float(probabilities[best]))] results[cls_config["task"]] = chosen return best = int(probabilities.argmax()) results[cls_config["task"]] = (labels[best], float(probabilities[best])) def extract( self, text: str, schema, threshold: float = 0.5, format_results: bool = True, include_confidence: bool = False, include_spans: bool = False, overlap_policy=None, ): schema_dicts, metadata_list = self._build_schema_dicts_and_metadata([schema]) if overlap_policy is not None: metadata_list[0]["_overlap_policy"] = self._resolved_overlap_policy(overlap_policy) arrays, batch = prepare_extraction( self.processor, text, schema_dicts[0], self.length, self.max_words, self.max_queries ) predicted = self.features_model.predict(arrays) features = {name: to_tensor(predicted[name]) for name in FEATURE_NAMES} query_mask = to_tensor(arrays["query_mask"]).bool() candidates = None if bool(query_mask.any()): head = self.boundary_settings pool = select_candidates( features["pool_start_projection"], features["pool_end_projection"], features["boundary_mask"].bool(), query_mask, features["start_logits"], features["end_logits"], boundary_top_k=head.pool_boundary_top_k, pool_size=head.pool_size, min_pool_per_query=head.min_pool_per_query, ) values = ( features["text_states"], to_tensor(arrays["text_mask"]), features["query_states"], to_tensor(arrays["query_mask"]), features["boundary_states"], features["start_logits"], features["end_logits"], features["inside_prefix"], features["inside_prefix_mean"], pool.indices.int(), pool.mask.float(), pool.compat_logits, ) scored = self.scorer_model.predict( { name: value.numpy().astype(np.int32 if name == "candidate_indices" else np.float32) for name, value in zip(SCORE_INPUT_NAMES, values) } ) candidates = CandidateTensorBatch( indices=pool.indices.unsqueeze(1).expand(1, self.max_queries, -1, 2), proposal_logits=pool.proposal_logits.unsqueeze(1).expand(1, self.max_queries, -1), pair_logits=to_tensor(scored["pair_logits"]), valid_mask=pool.mask.unsqueeze(1).expand(1, self.max_queries, -1), query_mask=query_mask, candidate_states=to_tensor(scored["candidate_states"]).unsqueeze(1).expand(1, self.max_queries, -1, -1), ) token = self._context.set({"arrays": arrays, "features": features, "candidates": candidates}) try: raw = self._extract_from_batch(batch, threshold, metadata_list, include_confidence, include_spans)[0] if format_results: return self.format_results( raw, include_confidence, metadata_list[0].get("relation_order", []), metadata_list[0].get("classification_tasks", []), ) return raw finally: self._context.reset(token) class CoreMLAdaptiveBoundaryExtractor: """Use an ANE-capable FP16 path and load FP32 for sensitive latent records.""" def __init__( self, model_dir: str, fp16_feature_package: str = "gliner2_multi_extraction_features_w8_embedding_linear_fp16_L128_W64_Q8.mlpackage", fp32_feature_package: str = "gliner2_multi_extraction_features_w8_embedding_fp32_L128_W64_Q8.mlpackage", fp16_compute_units=ct.ComputeUnit.CPU_AND_NE, fp32_compute_units=ct.ComputeUnit.ALL, ): self.model_dir = model_dir self.feature_packages = {"fp16": fp16_feature_package, "fp32": fp32_feature_package} self.compute_units = {"fp16": fp16_compute_units, "fp32": fp32_compute_units} self._lock = threading.Lock() self._active_precision = None self._runtime = None @staticmethod def selected_precision(schema) -> str: built = schema.build() if hasattr(schema, "build") else schema records = built.get("record_metadata", {}) return "fp32" if any(details.get("mode") == "latent" for details in records.values()) else "fp16" def extract(self, text: str, schema, **kwargs): precision = self.selected_precision(schema) with self._lock: if self._active_precision != precision: self._runtime = None self._active_precision = None gc.collect() self._runtime = CoreMLBoundaryExtractor( self.model_dir, precision=precision, compute_units=self.compute_units[precision], feature_package=self.feature_packages[precision], ) self._active_precision = precision return self._runtime.extract(text, schema, **kwargs)