File size: 3,255 Bytes
628e2fd | 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 | """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())):
# positions[0] is the [P] prompt marker; the rest are one [L] marker per label.
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,
}
|