File size: 3,154 Bytes
0deb31c | 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 | """Standalone Jeff classification inference from the Core ML package.
No original PyTorch checkpoint is loaded. Tokenization and the classification
prompt formatting use the pinned GLiFormer processor and tokenizer files.
"""
from __future__ import annotations
import json
from pathlib import Path
import coremltools as ct
import numpy as np
from gliformer.config import GLiFormerConfig
from gliformer.processing.collator import resolve_gliformer_collator_class
from gliformer.processing.processor import resolve_gliformer_processor_class
from gliner.data_processing.tokenizer import WordsSplitter
from transformers import AutoTokenizer
from jeff_decision import marker_positions
class JeffCoreML:
"""One-group, one-text Jeff classifier for 1–8 labels and at most 128 tokens."""
def __init__(self, asset_dir: str | Path, package: str | Path, compute_units=ct.ComputeUnit.ALL):
asset_dir = Path(asset_dir)
config = GLiFormerConfig(**json.loads((asset_dir / "gliner_config.json").read_text()))
tokenizer = AutoTokenizer.from_pretrained(asset_dir, local_files_only=True)
splitter = WordsSplitter(config.words_splitter_type)
processor_cls = resolve_gliformer_processor_class(config)
processor = processor_cls(config, tokenizer, splitter)
collator_cls = resolve_gliformer_collator_class(config)
self.collator = collator_cls(config, data_processor=processor, return_tokens=True, prepare_labels=False)
self.splitter = splitter
self.config = config
self.model = ct.models.MLModel(str(package), compute_units=compute_units)
self.output_name = self.model.get_spec().description.output[0].name
def score(self, text: str, labels: list[str], name: str = "", description: str = "") -> list[float]:
if not 1 <= len(labels) <= 8:
raise ValueError("JeffCoreML accepts 1 to 8 labels")
tokens = [word for word, _, _ in self.splitter(text)]
batch = self.collator([{
"tokenized_text": tokens,
"classification": [{
"name": name,
"description": description,
"all_labels": labels,
"true_labels": [],
}],
}])
ids = batch["input_ids"]
mask = batch["attention_mask"]
if ids.shape[1] > 128:
raise ValueError(f"JeffCoreML L128 bucket cannot fit {ids.shape[1]} tokens")
parent, children, count = marker_positions(ids, self.config, 8)
if count != len(labels):
raise ValueError("tokenizer label markers disagree with the supplied label count")
inputs = {
"input_ids": np.pad(ids.numpy().astype(np.int32), ((0, 0), (0, 128 - ids.shape[1]))),
"attention_mask": np.pad(mask.numpy().astype(np.int32), ((0, 0), (0, 128 - mask.shape[1]))),
"parent_position": parent.numpy(),
"category_positions": children.numpy(),
}
logits = self.model.predict(inputs)[self.output_name][0, :count].astype(np.float32)
probabilities = 1.0 / (1.0 + np.exp(-logits))
return probabilities.tolist()
|