| """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() |
|
|
