| """Native GLiNER2 span-architecture schema preprocessing for a fixed Core ML bucket.""" |
| import numpy as np |
| from gliner2 import Schema |
| from gliner2.models.base import load_extractor_tokenizer |
| from gliner2.processor import SchemaTransformer |
| from gliner2.training.trainer import ExtractorCollator |
|
|
|
|
| def load_processor(tokenizer_dir: str): |
| """Load only the tokenizer and schema formatter needed by the Core ML model.""" |
| return SchemaTransformer(tokenizer=load_extractor_tokenizer(tokenizer_dir), token_pooling="first") |
|
|
|
|
| def classification_schema(tasks: dict) -> Schema: |
| """Same task-dict handling as native ``classify_text``.""" |
| schema = Schema() |
| for name, config in tasks.items(): |
| if isinstance(config, dict) and "labels" in config: |
| cfg = config.copy() |
| labels = cfg.pop("labels") |
| schema.classification(name, labels, **cfg) |
| else: |
| schema.classification(name, config) |
| return schema |
|
|
|
|
| def task_labels(tasks: dict) -> dict[str, list[str]]: |
| result = {} |
| for name, config in tasks.items(): |
| labels = config["labels"] if isinstance(config, dict) and "labels" in config else config |
| result[name] = list(labels.keys()) if isinstance(labels, dict) else list(labels) |
| return result |
|
|
|
|
| def prepare_decision(processor, text: str, tasks: dict, length: int, max_heads: int, max_options: int): |
| """Tokenize ``tasks`` exactly as the native span collator does and pad into the bucket.""" |
| if not 1 <= len(tasks) <= max_heads: |
| raise ValueError(f"Expected 1..{max_heads} decision heads, got {len(tasks)}") |
| labels = task_labels(tasks) |
| for name, values in labels.items(): |
| if not 1 <= len(values) <= max_options: |
| raise ValueError(f"Head {name!r} has {len(values)} labels; bucket holds 1..{max_options}") |
| collator = ExtractorCollator(processor, is_training=False, max_len=None, architecture="span") |
| batch = collator([(text, classification_schema(tasks).build())]) |
| ids = batch.input_ids.numpy() |
| if ids.shape[1] > length: |
| raise ValueError(f"Schema and text require {ids.shape[1]} subwords; bucket holds {length}") |
| groups = batch.schema_special_indices[0] |
| if len(groups) != len(labels): |
| raise ValueError("Schema head count does not match the requested tasks") |
| indices = np.zeros((1, max_heads, max_options), dtype=np.int32) |
| mask = np.zeros((1, max_heads, max_options), dtype=np.float32) |
| for head, (positions, values) in enumerate(zip(groups, labels.values())): |
| |
| markers = list(positions[1:]) |
| if len(markers) != len(values): |
| raise ValueError("Label markers were truncated or merged") |
| indices[0, head, : len(markers)] = markers |
| mask[0, head, : len(markers)] = 1.0 |
| attention = batch.attention_mask.numpy() |
| pad = processor.tokenizer.pad_token_id |
| return { |
| "input_ids": np.pad(ids, ((0, 0), (0, length - ids.shape[1])), constant_values=pad).astype(np.int32), |
| "attention_mask": np.pad(attention, ((0, 0), (0, length - attention.shape[1]))).astype(np.int32), |
| "marker_indices": indices, |
| "marker_mask": mask, |
| } |
|
|
