Publish validated Jeff GLiFormer Large L128 FP16 Core ML classifier

JeffDecision-L128-FP16.mlpackage/Data/com.apple.CoreML/model.mlmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:87c0960087a6b39d6aac2e0c2511340f0197f19d471b77b2e894ca265204c1d3
+size 423429

JeffDecision-L128-FP16.mlpackage/Data/com.apple.CoreML/weights/weight.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:25988597bb4cc8b488970f7c6a15671dcf5f422c88c01d244a35a99a6aa45ba3
+size 922387968

JeffDecision-L128-FP16.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "A8DBF8BD-E4B2-4DC7-9A54-336D47C02F68": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        },
+        "AB3C25D4-080D-4C1A-9C41-B16757CD95F1": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        }
+    },
+    "rootModelIdentifier": "AB3C25D4-080D-4C1A-9C41-B16757CD95F1"
+}

README.md

@@ -0,0 +1,44 @@
+---
+license: apache-2.0
+library_name: coreml
+tags:
+  - coreml
+  - text-classification
+  - jeff
+  - gliformer
+  - on-device
+---
+
+# Jeff / GLiFormer Large classification for Core ML
+
+This is a fixed L128, batch-1 Core ML conversion of the **trained classification path** in [knowledgator/gliformer-large-v1](https://huggingface.co/knowledgator/gliformer-large-v1), as used by [Jeff](https://github.com/logan-markewich/jeff). It includes the 24-layer DeBERTa encoder, learned prompt-marker embeddings, and trained linear classification head in one FP16 ML Program. The original checkpoint contains 575,637,510 parameters; this Core ML package is about 880 MB on disk. The source checkpoint, tokenizer and decision-engine revisions are pinned in `assets.lock.json`.
+
+The model scores one classification group with 1–8 labels and up to 128 tokenized input tokens. Choice, yes/no and ordinal labels use the same trained classification head; the probabilities are independent sigmoid scores, not a normalized softmax. Only the first `len(labels)` logits are returned. Other GLiFormer tasks such as NER, layout extraction, vision, audio and structuring are **not** exported here. This is not a Decision Index benchmark result.
+
+The original model's classifier uses CLS pooling, a parent prompt anchor, linear parent+category fusion and dot scoring. Its word-level RNN is evaluated upstream but cannot affect these classification logits under this checkpoint's settings. `jeff_decision.py` checks those settings and carries the trained encoder and head; `trace_compat.py` makes fixed-shape DeBERTa attention convertible without changing trained weights. The trace-only finite mask substitutes -10000 for fp32-min before FP16 conversion; patched PyTorch logits match native logits on the checked cases.
+
+## Local inference
+
+On macOS 15 or later with Python 3.12:
+
+```bash
+uv sync --frozen
+uv run python - <<'PY'
+from runtime import JeffCoreML
+
+model = JeffCoreML(".", "JeffDecision-L128-FP16.mlpackage")
+print(model.score(
+    "The invoice was charged twice and the customer asks for a refund.",
+    ["billing: invoice or payment issue", "support: technical product issue"],
+    name="Choose the correct support queue",
+))
+PY
+```
+
+`JeffCoreML` uses the bundled tokenizer/config and GLiFormer processor, and loads **no original PyTorch weights**. Inputs exceeding 128 tokens or eight labels fail explicitly. Build from the pinned source checkpoint with `uv run python export.py --convert --precision fp16`; the source checkpoint must be in the local Hugging Face cache. `uv run python verify.py --precision fp16` compares Core ML logits to native PyTorch.
+
+## Validation
+
+Four real source-checkpoint fixtures cover billing, technical support, a three-label intent choice and yes/no classification. The mathematical decision wrapper matched native logits within **3.82e-6**. Tracing-only patches matched native logits exactly on those fixtures. The exported FP16 Core ML model preserved all four chosen labels with maximum absolute logit error **0.1139**. FP32 Core ML is a diagnostic control with four-of-four agreement and maximum logit error **3.44e-5**; the FP32 package is not included because it is much larger. Full Decision Index quality, additional input lengths, other task heads and broad latency/ANE performance have not been evaluated. See `native-parity.json` and `coreml-parity-fp16.json`.
+
+This conversion is derived from the Apache-2.0 GLiFormer weights and [Transformers](https://github.com/huggingface/transformers) DeBERTa implementation. Jeff's decision adapter code is MIT licensed; the helper source here retains attribution. The model card makes no claim that this conversion is faster or more accurate than another model on a benchmark.

assets.lock.json

@@ -0,0 +1,16 @@
+{
+  "decision_engine": {
+    "repo": "https://github.com/logan-markewich/jeff",
+    "revision": "34b32f99a727c47b679adde33f4702a001e02979",
+    "license": "MIT"
+  },
+  "checkpoint": {
+    "repo": "knowledgator/gliformer-large-v1",
+    "revision": "d0a4e53d09cebe6bc963dd9be319d4279084bb2d",
+    "weights_file": "pytorch_model.bin",
+    "weights_bytes": 2302735855,
+    "weights_sha256": "f80b29199d66f878669f283703e4dba9fd726755dcc20aba1ed0d24fce4a23f1",
+    "license": "Apache-2.0"
+  },
+  "scope": "Jeff typed-decision classification path; not all GLiFormer extraction heads"
+}

coreml-parity-fp16.json

@@ -0,0 +1,77 @@
+[
+  {
+    "name": "billing",
+    "labels": [
+      "billing: invoice or payment issue",
+      "support: technical product issue"
+    ],
+    "native_logits": [
+      10.187458038330078,
+      -6.71131706237793
+    ],
+    "coreml_logits": [
+      10.125,
+      -6.65625
+    ],
+    "max_logit_error": 0.062458038330078125,
+    "top_label_agreement": true,
+    "coreml_wall_ms": 165.4915830004029
+  },
+  {
+    "name": "technical",
+    "labels": [
+      "billing: invoice or payment issue",
+      "support: technical product issue"
+    ],
+    "native_logits": [
+      -12.322779655456543,
+      13.379047393798828
+    ],
+    "coreml_logits": [
+      -12.296875,
+      13.4140625
+    ],
+    "max_logit_error": 0.035015106201171875,
+    "top_label_agreement": true,
+    "coreml_wall_ms": 82.55887497216463
+  },
+  {
+    "name": "three_way",
+    "labels": [
+      "schedule: appointment request",
+      "billing: payment issue",
+      "support: technical issue"
+    ],
+    "native_logits": [
+      5.41853141784668,
+      -12.076669692993164,
+      -9.080745697021484
+    ],
+    "coreml_logits": [
+      5.40234375,
+      -12.09375,
+      -9.1171875
+    ],
+    "max_logit_error": 0.036441802978515625,
+    "top_label_agreement": true,
+    "coreml_wall_ms": 53.515166975557804
+  },
+  {
+    "name": "boolean",
+    "labels": [
+      "yes",
+      "no"
+    ],
+    "native_logits": [
+      1.0098832845687866,
+      -1.0589547157287598
+    ],
+    "coreml_logits": [
+      0.89599609375,
+      -1.001953125
+    ],
+    "max_logit_error": 0.11388719081878662,
+    "top_label_agreement": true,
+    "coreml_wall_ms": 53.68766700848937
+  }
+]

export.py

@@ -0,0 +1,166 @@
+"""Validate Jeff's trained classifier graph and export its complete decision path."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import warnings
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from huggingface_hub import snapshot_download
+from jeff.backends.torch_backend import TorchBackend
+from jeff.core.backend import Group
+
+from jeff_decision import JeffDecision, marker_positions
+from trace_compat import finite_fp16_mask, install_trace_compatibility
+
+SOURCE = "knowledgator/gliformer-large-v1"
+REVISION = "d0a4e53d09cebe6bc963dd9be319d4279084bb2d"
+BUCKET = 128
+MAX_CATEGORIES = 8
+
+FIXTURES = (
+    (
+        "billing",
+        "The invoice was charged twice and the customer asks for a refund.",
+        Group(key="route", labels=("billing: invoice or payment issue", "support: technical product issue"),
+              name="Choose the correct support queue"),
+    ),
+    (
+        "technical",
+        "The app crashes when I save my project. Please help me recover the file.",
+        Group(key="route", labels=("billing: invoice or payment issue", "support: technical product issue"),
+              name="Choose the correct support queue"),
+    ),
+    (
+        "three_way",
+        "Tomorrow at 9 a.m. works well for the appointment.",
+        Group(key="intent", labels=("schedule: appointment request", "billing: payment issue",
+                                     "support: technical issue"), name="Classify the user intent"),
+    ),
+    (
+        "boolean",
+        "I cannot sign in after resetting my password.",
+        Group(key="answer", labels=("yes", "no"), name="Is this a technical support request?"),
+    ),
+)
+
+
+def make_batch(backend: TorchBackend, text: str, group: Group) -> dict:
+    tokens, _, _ = backend.model.prepare_inputs([text])
+    return backend._collator([{
+        "tokenized_text": tokens[0],
+        "classification": [{
+            "name": group.name,
+            "description": group.description,
+            "all_labels": list(group.labels),
+            "true_labels": [],
+        }],
+    }])
+
+
+def model_inputs(batch: dict, config) -> tuple[torch.Tensor, ...]:
+    ids = batch["input_ids"]
+    mask = batch["attention_mask"]
+    if ids.shape[1] > BUCKET:
+        raise ValueError(f"input has {ids.shape[1]} tokens; L{BUCKET} cannot serve it")
+    parent, children, count = marker_positions(ids, config, MAX_CATEGORIES)
+    if count != len(batch["classes_mapping"].cat_mapping[0].cat_class_to_id[0].class_to_id):
+        raise ValueError("collator category mapping does not match marker count")
+    pad = BUCKET - ids.shape[1]
+    return (
+        F.pad(ids.to(torch.int32), (0, pad)),
+        F.pad(mask.to(torch.int32), (0, pad)),
+        parent,
+        children,
+    )
+
+
+@torch.inference_mode()
+def native_report(backend: TorchBackend, decision: JeffDecision) -> tuple[list[dict], tuple[torch.Tensor, ...]]:
+    records = []
+    first = None
+    for name, text, group in FIXTURES:
+        batch = make_batch(backend, text, group)
+        tensors = model_inputs(batch, backend.model.config)
+        native = backend.model.model(**batch, include_media=False).cat_logits.detach().float().numpy()[0]
+        converted = decision(*tensors).detach().float().numpy()[0, :len(group.labels)]
+        error = float(np.max(np.abs(native - converted)))
+        records.append({
+            "name": name,
+            "token_count": int(batch["attention_mask"].sum()),
+            "labels": list(group.labels),
+            "native_logits": native.tolist(),
+            "decision_logits": converted.tolist(),
+            "max_logit_error": error,
+        })
+        if error > 1e-3:
+            raise AssertionError(f"{name}: trained graph differs from native logits by {error}")
+        if first is None:
+            first = tensors
+    assert first is not None
+    return records, first
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--convert", action="store_true")
+    parser.add_argument("--precision", choices=("fp16", "fp32"), default="fp16")
+    args = parser.parse_args()
+    torch.set_num_threads(2)
+    checkpoint = snapshot_download(SOURCE, revision=REVISION, local_files_only=True)
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", message=r"attn_kernel=.*flashdeberta")
+        backend = TorchBackend(checkpoint, device="cpu", dtype="float32", attn_kernel="eager", batch_size=1)
+    decision = JeffDecision(backend.model.model).eval()
+    report, first = native_report(backend, decision)
+    out = Path("build")
+    out.mkdir(exist_ok=True)
+    (out / "native-parity.json").write_text(json.dumps(report, indent=2) + "\n")
+    print(json.dumps({"native_parity": report}, indent=2), flush=True)
+    if not args.convert:
+        return
+
+    import coremltools as ct
+    install_trace_compatibility()
+    with finite_fp16_mask():
+        patched_report, _ = native_report(backend, decision)
+        patched_error = max(
+            abs(before - after)
+            for baseline, patched in zip(report, patched_report)
+            for before, after in zip(baseline["decision_logits"], patched["decision_logits"])
+        )
+        if patched_error > 1e-3:
+            raise AssertionError(f"tracing-only mask/attention scale changes native logits by {patched_error}")
+        print(f"Patched mask/attention max native logit error: {patched_error:.8f}", flush=True)
+        traced = torch.jit.trace(decision, first, check_trace=False).eval()
+    traced.save(str(out / "jeff-decision-L128.pt"))
+    with torch.inference_mode():
+        expected = decision(*first).detach().numpy()
+        actual = traced(*first).detach().numpy()
+    trace_error = float(np.max(np.abs(expected - actual)))
+    if trace_error > 1e-3:
+        raise AssertionError(f"TorchScript trace mismatch: {trace_error}")
+    print(f"TorchScript trace max logit error: {trace_error:.8f}", flush=True)
+    mlmodel = ct.convert(
+        traced,
+        convert_to="mlprogram",
+        minimum_deployment_target=ct.target.macOS15,
+        compute_precision=ct.precision.FLOAT16 if args.precision == "fp16" else ct.precision.FLOAT32,
+        inputs=[
+            ct.TensorType(name="input_ids", shape=(1, BUCKET), dtype=np.int32),
+            ct.TensorType(name="attention_mask", shape=(1, BUCKET), dtype=np.int32),
+            ct.TensorType(name="parent_position", shape=(1, 1), dtype=np.int32),
+            ct.TensorType(name="category_positions", shape=(1, MAX_CATEGORIES), dtype=np.int32),
+        ],
+    )
+    package = out / f"JeffDecision-L128-{args.precision.upper()}.mlpackage"
+    mlmodel.save(str(package))
+    print(f"Saved {package}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

gliner_config.json

@@ -0,0 +1,441 @@
+{
+  "adjacency_loss_coef": 1.0,
+  "anchor_num_heads": 4,
+  "anchor_num_layers": 2,
+  "backbone_type": "deberta_2d",
+  "bos_token_id": 1,
+  "cat_loss_coef": 1.0,
+  "cat_parent_token": "[SCHEMA]",
+  "cat_token": "[CLASS]",
+  "cat_token_index": 128003,
+  "child_token": "[FIELD]",
+  "child_token_index": 128005,
+  "class_token_index": 128001,
+  "classification_config": {
+    "anchor_context_gate_init": 0.1,
+    "anchor_context_gate_trainable": true,
+    "anchor_cross_attention_bias": null,
+    "anchor_layer": null,
+    "anchor_memory_position": null,
+    "anchor_memory_position_usage": null,
+    "anchor_mode": "parent",
+    "anchor_modeling": "linear",
+    "anchor_normalization": "none",
+    "anchor_query_position": null,
+    "anchor_refine_heads": 8,
+    "anchor_refine_layer_scale_init": null,
+    "anchor_refine_layers": 0,
+    "anchor_refine_norm": "post_norm",
+    "anchor_refinement": null,
+    "anchor_self_attention_bias": null,
+    "cat_token_index": 128003,
+    "embed_cat_token": true,
+    "embed_parent_token": true,
+    "feature_anchor_mlp": false,
+    "feature_anchor_mlp_hidden_multiplier": 1,
+    "focal_loss_alpha": 0.8,
+    "focal_loss_gamma": -1.0,
+    "focal_loss_prob_margin": 0.0,
+    "loss_coef": 1.0,
+    "parent_token_index": 128002,
+    "pooling_type": "cls",
+    "scorer_type": "dot"
+  },
+  "classifier_layer": null,
+  "count_config": null,
+  "count_layer": null,
+  "count_loss_coef": 1.0,
+  "count_mode": "regression",
+  "dropout": 0.25,
+  "embed_cat_token": true,
+  "embed_child_token": true,
+  "embed_ent_token": true,
+  "embed_parent_token": true,
+  "embed_rel_token": true,
+  "embedding_config": {
+    "encoder_dropout": 0.0,
+    "loss_coef": 1.0,
+    "loss_fn": "mse",
+    "margin": 0.0,
+    "pooling_type": "cls",
+    "projection_dim": 1024,
+    "projection_dropout": 0.0,
+    "similarity_fn": "cosine"
+  },
+  "embedding_loss_coef": 1.0,
+  "encoder_config": {
+    "_name_or_path": "knowledgator/DeBERTa-large-joint-3000",
+    "architectures": [
+      "FlashDebertaV2ForContextPredictionTraining"
+    ],
+    "attention_probs_dropout_prob": 0.0,
+    "attn_kernel": "flash",
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "coordinate_size": 128,
+    "dtype": "bfloat16",
+    "eos_token_id": null,
+    "hidden_act": "gelu",
+    "hidden_dropout_prob": 0.0,
+    "hidden_size": 1024,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "intermediate_size": 4096,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_norm_eps": 1e-07,
+    "layout_bias_propagation": "all_layers",
+    "layout_embedding_type": "absolute",
+    "layout_max_relative_positions": 1024,
+    "layout_position_buckets": 32,
+    "layout_relative_attention": false,
+    "legacy": true,
+    "max_2d_position_embeddings": 1024,
+    "max_page_embeddings": 1024,
+    "max_position_embeddings": 512,
+    "max_relative_positions": -1,
+    "model_type": "layout-deberta",
+    "norm_rel_ebd": "layer_norm",
+    "num_attention_heads": 16,
+    "num_hidden_layers": 24,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "pad_token_id": 0,
+    "pooler_dropout": 0.0,
+    "pooler_hidden_act": "gelu",
+    "pooler_hidden_size": 1024,
+    "pos_att_type": [
+      "p2c",
+      "c2p"
+    ],
+    "position_biased_input": false,
+    "position_buckets": 256,
+    "problem_type": null,
+    "relative_attention": true,
+    "return_dict": true,
+    "shape_size": 64,
+    "share_att_key": true,
+    "spatial_embedding_type": "absolute",
+    "tie_word_embeddings": true,
+    "type_vocab_size": 0,
+    "vocab_size": 128008
+  },
+  "ent_token": "[ENTITY]",
+  "eos_token_id": 2,
+  "fine_tune": true,
+  "fuse_layers": false,
+  "groups_layer": "fixed",
+  "groups_loss_coef": 1.0,
+  "hidden_size": 1024,
+  "image_size": 224,
+  "joint_relex_config": {
+    "adjacency_loss_coef": 1.0,
+    "anchor_context_gate_init": 0.1,
+    "anchor_context_gate_trainable": true,
+    "anchor_cross_attention_bias": null,
+    "anchor_layer": null,
+    "anchor_memory_position": null,
+    "anchor_memory_position_usage": null,
+    "anchor_mode": "fixed",
+    "anchor_modeling": "linear",
+    "anchor_normalization": "none",
+    "anchor_num_heads": 4,
+    "anchor_num_layers": 2,
+    "anchor_query_position": null,
+    "anchor_refine_heads": 8,
+    "anchor_refine_layer_scale_init": null,
+    "anchor_refine_layers": 0,
+    "anchor_refine_norm": "post_norm",
+    "anchor_refinement": null,
+    "anchor_self_attention_bias": null,
+    "embed_parent_token": true,
+    "embed_rel_token": true,
+    "feature_anchor_mlp": false,
+    "feature_anchor_mlp_hidden_multiplier": 1,
+    "focal_loss_alpha": 0.8,
+    "focal_loss_gamma": -1.0,
+    "focal_loss_prob_margin": 0.0,
+    "loss_coef": 1.0,
+    "max_count": 20,
+    "max_relation_entities": null,
+    "max_relation_span_width": null,
+    "neg_spans_ratio": 1.0,
+    "num_fixed_slots": 10,
+    "parent_token_index": -1,
+    "rel_token_index": 128004,
+    "relation_focal_loss_alpha": 0.8,
+    "relation_focal_loss_gamma": 2,
+    "relation_focal_loss_prob_margin": 0.0,
+    "relation_loss_coef": 1.0,
+    "relation_loss_reduction": "mean",
+    "relation_neighbor_chunk_size": 512,
+    "relation_span_nms": false,
+    "relation_top_k_neighbors": null,
+    "relations_layer": null,
+    "represent_spans": false,
+    "span_loss_coef": 1.0,
+    "span_loss_reduction": "mean",
+    "triples_layer": null
+  },
+  "labels_encoder": null,
+  "labels_encoder_config": null,
+  "layout_image_tokens": true,
+  "max_count": 20,
+  "max_len": 16384,
+  "max_neg_type_ratio": 1,
+  "max_page_embeddings": 1024,
+  "max_types": 128,
+  "max_width": 12,
+  "media_parent_embedding_source": "fixed",
+  "model_name": "knowledgator/DeBERTa-large-joint-3000",
+  "model_type": "gliformer-layout",
+  "model_variant": "layout",
+  "multimodal_fusion": "uni-encoder",
+  "name": "gliformer",
+  "neg_spans_ratio": 1.0,
+  "ner_config": {
+    "anchor_context_gate_init": 0.1,
+    "anchor_context_gate_trainable": true,
+    "anchor_cross_attention_bias": null,
+    "anchor_layer": null,
+    "anchor_memory_position": null,
+    "anchor_memory_position_usage": null,
+    "anchor_mode": "parent",
+    "anchor_modeling": "linear",
+    "anchor_normalization": "none",
+    "anchor_query_position": null,
+    "anchor_refine_heads": 8,
+    "anchor_refine_layer_scale_init": null,
+    "anchor_refine_layers": 0,
+    "anchor_refine_norm": "post_norm",
+    "anchor_refinement": null,
+    "anchor_self_attention_bias": null,
+    "embed_parent_token": true,
+    "feature_anchor_mlp": false,
+    "feature_anchor_mlp_hidden_multiplier": 1,
+    "focal_loss_alpha": 0.8,
+    "focal_loss_gamma": -1,
+    "focal_loss_prob_margin": 0.0,
+    "loss_coef": 1.0,
+    "neg_spans_ratio": 1.0,
+    "parent_token_index": 128002,
+    "represent_spans": false,
+    "span_loss_coef": 1.0,
+    "span_loss_reduction": "mean"
+  },
+  "ner_loss_coef": 1.0,
+  "ner_parent_token": "[SCHEMA]",
+  "num_post_fusion_layers": 1,
+  "num_rnn_layers": 1,
+  "obj_token": "[OBJECT]",
+  "open_rel_parent_token": "[SCHEMA]",
+  "open_relex_config": null,
+  "pad_token_id": 0,
+  "parent_token": "[SCHEMA]",
+  "parent_token_index": 128002,
+  "per_task_parents": false,
+  "post_fusion_schema": "",
+  "projector_hidden_act": "gelu",
+  "rel_mode": "adjacency",
+  "rel_token": "[RELATION]",
+  "rel_token_index": 128004,
+  "relation_loss_coef": 1.0,
+  "relations_config": {
+    "adjacency_loss_coef": 1.0,
+    "anchor_context_gate_init": 0.1,
+    "anchor_context_gate_trainable": true,
+    "anchor_cross_attention_bias": null,
+    "anchor_layer": null,
+    "anchor_memory_position": null,
+    "anchor_memory_position_usage": null,
+    "anchor_mode": "fixed",
+    "anchor_modeling": "linear",
+    "anchor_normalization": "none",
+    "anchor_num_heads": 4,
+    "anchor_num_layers": 2,
+    "anchor_query_position": null,
+    "anchor_refine_heads": 8,
+    "anchor_refine_layer_scale_init": null,
+    "anchor_refine_layers": 0,
+    "anchor_refine_norm": "post_norm",
+    "anchor_refinement": null,
+    "anchor_self_attention_bias": null,
+    "embed_parent_token": true,
+    "embed_rel_token": true,
+    "feature_anchor_mlp": false,
+    "feature_anchor_mlp_hidden_multiplier": 1,
+    "focal_loss_alpha": 0.8,
+    "focal_loss_gamma": -1.0,
+    "focal_loss_prob_margin": 0.0,
+    "loss_coef": 1.0,
+    "max_count": 20,
+    "max_relation_entities": null,
+    "max_relation_span_width": null,
+    "neg_spans_ratio": 1.0,
+    "num_fixed_slots": 10,
+    "parent_token_index": -1,
+    "rel_token_index": 128004,
+    "relation_focal_loss_alpha": 0.8,
+    "relation_focal_loss_gamma": 2,
+    "relation_focal_loss_prob_margin": 0.0,
+    "relation_loss_coef": 1.0,
+    "relation_loss_reduction": "mean",
+    "relation_neighbor_chunk_size": 512,
+    "relation_span_nms": false,
+    "relation_top_k_neighbors": null,
+    "relations_layer": null,
+    "represent_spans": false,
+    "span_loss_coef": 1.0,
+    "span_loss_reduction": "mean",
+    "triples_layer": null
+  },
+  "relations_layer": null,
+  "represent_spans": false,
+  "sep_token": "[SEP]",
+  "seq_token": "[SEQ]",
+  "shared_anchor_modeling": null,
+  "shared_anchor_refine_heads": 8,
+  "shared_anchor_refine_layers": 0,
+  "shared_anchor_refinement": null,
+  "span_loss_coef": 1.0,
+  "span_mode": "markerV0",
+  "struct_parent_token": "[SCHEMA]",
+  "structuring_child_token": "<<CHILD>>",
+  "structuring_config": {
+    "anchor_context_gate_init": 0.1,
+    "anchor_context_gate_trainable": true,
+    "anchor_cross_attention_bias": {
+      "params": {
+        "sigma": 0.5,
+        "units": "query_steps",
+        "weight": 1.0
+      },
+      "type": "gaussian_distance"
+    },
+    "anchor_layer": {
+      "params": {
+        "num_slots": 100
+      },
+      "type": "fixed"
+    },
+    "anchor_memory_position": {
+      "params": {},
+      "type": "fourier"
+    },
+    "anchor_memory_position_usage": "keys_and_values",
+    "anchor_mode": "fixed_transformer",
+    "anchor_modeling": "linear",
+    "anchor_normalization": "center_rms",
+    "anchor_num_heads": 4,
+    "anchor_num_layers": 2,
+    "anchor_objectness": true,
+    "anchor_objectness_loss_coef": 1.0,
+    "anchor_objectness_threshold": 0.5,
+    "anchor_query_position": {
+      "params": {},
+      "type": "fourier"
+    },
+    "anchor_refine_heads": 8,
+    "anchor_refine_layer_scale_init": null,
+    "anchor_refine_layers": 0,
+    "anchor_refine_norm": "post_norm",
+    "anchor_refinement": {
+      "params": {
+        "layer_scale_init": 0.1,
+        "norm_style": "pre_norm",
+        "num_heads": 8,
+        "num_layers": 2
+      },
+      "type": "cross_attention"
+    },
+    "anchor_relations_focal_loss_alpha": 0.75,
+    "anchor_relations_focal_loss_gamma": null,
+    "anchor_relations_focal_loss_prob_margin": null,
+    "anchor_relations_layer": "mlp",
+    "anchor_relations_loss_coef": 1.0,
+    "anchor_relations_threshold": 0.5,
+    "anchor_self_attention_bias": null,
+    "assignment_loss_coef": 1.0,
+    "bio_loss_reduction": "mean",
+    "child_token_index": 128005,
+    "embed_child_token": true,
+    "embed_parent_token": true,
+    "entity_loss_coef": 1.0,
+    "feature_anchor_mlp": false,
+    "feature_anchor_mlp_hidden_multiplier": 1,
+    "focal_loss_alpha": 0.75,
+    "focal_loss_gamma": -1,
+    "focal_loss_prob_margin": 0.0,
+    "head_type": "structuring",
+    "log_loss_stats": false,
+    "log_loss_stats_every": 50,
+    "loss_coef": 1.0,
+    "masking": "none",
+    "matcher_dice_cost": 1.0,
+    "matcher_membership_cost": 1.0,
+    "matcher_membership_temperature": 1.0,
+    "matcher_objectness_cost": 0.5,
+    "matcher_objectness_temperature": 1.0,
+    "matching_focal_loss_alpha": 0.99,
+    "matching_focal_loss_gamma": null,
+    "matching_focal_loss_prob_margin": null,
+    "max_count": 100,
+    "memory_position_embedding_kwargs": {},
+    "memory_position_embedding_type": "none",
+    "memory_position_in_values": false,
+    "multi_level": true,
+    "neg_spans_ratio": 1.0,
+    "negatives": 1.0,
+    "ner_focal_loss_alpha": 0.75,
+    "ner_focal_loss_gamma": null,
+    "ner_focal_loss_prob_margin": null,
+    "num_fixed_slots": 10,
+    "objectness_focal_loss_alpha": 0.5,
+    "objectness_focal_loss_gamma": null,
+    "objectness_focal_loss_prob_margin": null,
+    "parent_token_index": 128002,
+    "position_bucket_attention_bias_type": "none",
+    "position_bucket_attention_bias_weight": 1.0,
+    "position_bucket_attention_sigma": 0.5,
+    "position_bucket_normalization": "none",
+    "query_position_embedding_kwargs": {},
+    "query_position_embedding_type": "none",
+    "represent_spans": true,
+    "reuse_ner_head": true,
+    "span_loss_coef": 1.0,
+    "span_loss_reduction": "mean",
+    "structure_mode": {
+      "decoder": null,
+      "decoder_options": {},
+      "params": {},
+      "processor": null,
+      "processor_options": {},
+      "type": "multi_level"
+    },
+    "use_anchor_matching": true
+  },
+  "structuring_end_token": "<<END>>",
+  "structuring_loss_coef": 1.0,
+  "subtoken_pooling": "first",
+  "token_loss_coef": 1.0,
+  "transformers_version": "5.16.1",
+  "triples_layer": null,
+  "use_layout": true,
+  "vision_center_crop_size": null,
+  "vision_do_normalize": false,
+  "vision_do_rescale": true,
+  "vision_image_mean": null,
+  "vision_image_std": null,
+  "vision_interpolation": "bilinear",
+  "vision_processor_name": null,
+  "vision_processor_type": "custom",
+  "vision_resize_size": null,
+  "vocab_size": 128008,
+  "words_splitter_type": "whitespace"
+}

jeff_decision.py

@@ -0,0 +1,79 @@
+"""The trained GLiFormer Large classification path used by Jeff.
+
+This checkpoint's classification config uses CLS pooling, parent anchors,
+no anchor refinement/normalization, linear anchor modeling, and dot scoring.
+Under those exact settings the word-level RNN is computed upstream but cannot
+influence classification logits. We still assert the settings at construction.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+class JeffDecision(nn.Module):
+    def __init__(self, model: nn.Module):
+        super().__init__()
+        config = model.config.classification_config
+        expected = {
+            "pooling_type": "cls",
+            "anchor_mode": "parent",
+            "anchor_modeling": "linear",
+            "anchor_normalization": "none",
+            "scorer_type": "dot",
+            "anchor_refine_layers": 0,
+        }
+        for name, value in expected.items():
+            actual = getattr(config, name)
+            if actual != value:
+                raise ValueError(f"unsupported classification config {name}={actual!r}; expected {value!r}")
+        if not config.embed_parent_token or not config.embed_cat_token:
+            raise ValueError("this path requires embeddings at the parent and category marker tokens")
+        if model.config.hidden_size != 1024:
+            raise ValueError("this fixed classifier requires the pinned 1024-wide checkpoint")
+        head = model.heads["classification"]
+        if hasattr(head, "anchor_refine"):
+            raise ValueError("classification anchor refinement cannot be omitted")
+        if type(head.anchor_layer).__name__ != "ParentAnchorLayer":
+            raise ValueError("unsupported anchor layer")
+        if type(head.anchor_modeling).__name__ != "LinearAnchorModeling":
+            raise ValueError("unsupported anchor model")
+        self.encoder = model.token_rep_layer
+        self.projection = head.anchor_modeling.proj
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        parent_position: torch.Tensor,
+        category_positions: torch.Tensor,
+    ) -> torch.Tensor:
+        """Return unnormalized logits for one classification group, padded to C=8."""
+        encoded = self.encoder(input_ids.long(), attention_mask.long())
+        cls = encoded[:, 0, :]
+        parent_index = parent_position.long().unsqueeze(-1).expand(1, 1, 1024)
+        parent = torch.gather(encoded, 1, parent_index)
+        child_index = category_positions.long().unsqueeze(-1).expand(1, 8, 1024)
+        children = torch.gather(encoded, 1, child_index)
+        combined = torch.cat((parent.expand(1, 8, 1024), children), dim=-1)
+        fused = self.projection(combined)
+        return (cls.unsqueeze(1) * fused).sum(dim=-1)
+
+
+def marker_positions(
+    input_ids: torch.Tensor, config, max_categories: int = 8
+) -> tuple[torch.Tensor, torch.Tensor, int]:
+    """Find the actual learned prompt markers in a one-row Jeff collator batch."""
+    if input_ids.shape[0] != 1:
+        raise ValueError("one classification group per call is required")
+    parent = torch.nonzero(input_ids[0] == config.classification_config.parent_token_index).flatten()
+    children = torch.nonzero(input_ids[0] == config.classification_config.cat_token_index).flatten()
+    if parent.numel() != 1 or not (1 <= children.numel() <= max_categories):
+        raise ValueError(
+            f"expected 1 parent and 1..{max_categories} categories; got {parent.numel()}, {children.numel()}"
+        )
+    count = int(children.numel())
+    category_positions = torch.zeros((1, max_categories), dtype=torch.int32)
+    category_positions[0, :count] = children.to(torch.int32)
+    return parent.to(torch.int32).view(1, 1), category_positions, count

native-parity.json

@@ -0,0 +1,73 @@
+[
+  {
+    "name": "billing",
+    "token_count": 39,
+    "labels": [
+      "billing: invoice or payment issue",
+      "support: technical product issue"
+    ],
+    "native_logits": [
+      10.187458038330078,
+      -6.71131706237793
+    ],
+    "decision_logits": [
+      10.187459945678711,
+      -6.71131706237793
+    ],
+    "max_logit_error": 1.9073486328125e-06
+  },
+  {
+    "name": "technical",
+    "token_count": 42,
+    "labels": [
+      "billing: invoice or payment issue",
+      "support: technical product issue"
+    ],
+    "native_logits": [
+      -12.322779655456543,
+      13.379047393798828
+    ],
+    "decision_logits": [
+      -12.322783470153809,
+      13.379046440124512
+    ],
+    "max_logit_error": 3.814697265625e-06
+  },
+  {
+    "name": "three_way",
+    "token_count": 41,
+    "labels": [
+      "schedule: appointment request",
+      "billing: payment issue",
+      "support: technical issue"
+    ],
+    "native_logits": [
+      5.41853141784668,
+      -12.076669692993164,
+      -9.080745697021484
+    ],
+    "decision_logits": [
+      5.418530464172363,
+      -12.076668739318848,
+      -9.080747604370117
+    ],
+    "max_logit_error": 1.9073486328125e-06
+  },
+  {
+    "name": "boolean",
+    "token_count": 29,
+    "labels": [
+      "yes",
+      "no"
+    ],
+    "native_logits": [
+      1.0098832845687866,
+      -1.0589547157287598
+    ],
+    "decision_logits": [
+      1.009883165359497,
+      -1.0589548349380493
+    ],
+    "max_logit_error": 1.1920928955078125e-07
+  }
+]

probe-native.py

@@ -0,0 +1,65 @@
+"""Inspect Jeff's pinned, trained GLiFormer decision boundary using a real request."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+from huggingface_hub import snapshot_download
+from jeff.backends.torch_backend import TorchBackend
+from jeff.core.backend import Group
+
+SOURCE_REPO = "knowledgator/gliformer-large-v1"
+SOURCE_REVISION = "d0a4e53d09cebe6bc963dd9be319d4279084bb2d"
+
+
+def main() -> None:
+    torch.set_num_threads(2)
+    checkpoint = snapshot_download(SOURCE_REPO, revision=SOURCE_REVISION)
+    backend = TorchBackend(checkpoint, device="cpu", dtype="float32", attn_kernel="eager", batch_size=1)
+    text = "The invoice was charged twice and the customer asks for a refund."
+    group = Group(
+        key="route",
+        labels=("billing: invoice or payment issue", "support: technical product issue"),
+        name="Choose the correct support queue",
+    )
+    native = backend.score([text], [[group]])[0]
+    tokens, _, _ = backend.model.prepare_inputs([text])
+    batch = backend._collator(
+        [
+            {
+                "tokenized_text": tokens[0],
+                "classification": [
+                    {
+                        "name": group.name,
+                        "description": group.description,
+                        "all_labels": list(group.labels),
+                        "true_labels": [],
+                    }
+                ],
+            }
+        ]
+    )
+    report = {
+        "source_repo": SOURCE_REPO,
+        "source_revision": SOURCE_REVISION,
+        "backend": backend.info(),
+        "scores": native.scores,
+        "input_tokens": native.input_tokens,
+        "model_parameters": sum(parameter.numel() for parameter in backend.model.model.parameters()),
+        "batch_tensors": {
+            name: {"shape": list(value.shape), "dtype": str(value.dtype)}
+            for name, value in batch.items()
+            if isinstance(value, torch.Tensor)
+        },
+        "batch_other": {name: str(type(value)) for name, value in batch.items() if not isinstance(value, torch.Tensor)},
+    }
+    output = Path("build/native-probe.json")
+    output.parent.mkdir(parents=True, exist_ok=True)
+    output.write_text(json.dumps(report, indent=2) + "\n")
+    print(json.dumps(report, indent=2), flush=True)
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -0,0 +1,25 @@
+[project]
+name = "jeff-coreml"
+version = "0.1.0"
+requires-python = ">=3.12,<3.13"
+dependencies = [
+  "coremltools==9.0",
+  "gliformer==0.1.2",
+  "jeff @ git+https://github.com/logan-markewich/jeff.git@34b32f99a727c47b679adde33f4702a001e02979",
+  "numpy>=2.5.3",
+  "torch==2.7.0",
+]
+
+[dependency-groups]
+dev = ["pytest>=9.1", "ruff>=0.13"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+pythonpath = ["."]
+
+[tool.ruff]
+line-length = 120
+target-version = "py312"
+
+[tool.ruff.lint]
+select = ["E", "F", "I"]

runtime.py

@@ -0,0 +1,68 @@
+"""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()

tokenizer_config.json

@@ -0,0 +1,28 @@
+{
+  "add_prefix_space": true,
+  "backend": "tokenizers",
+  "bos_token": "[CLS]",
+  "clean_up_tokenization_spaces": false,
+  "cls_token": "[CLS]",
+  "do_lower_case": false,
+  "eos_token": "[SEP]",
+  "is_local": true,
+  "local_files_only": false,
+  "mask_token": "[MASK]",
+  "max_length": 1024,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_to_multiple_of": null,
+  "pad_token": "[PAD]",
+  "pad_token_type_id": 0,
+  "padding_side": "right",
+  "sep_token": "[SEP]",
+  "sp_model_kwargs": {},
+  "split_by_punct": false,
+  "stride": 0,
+  "tokenizer_class": "DebertaV2Tokenizer",
+  "truncation_side": "right",
+  "truncation_strategy": "longest_first",
+  "unk_id": 3,
+  "unk_token": "[UNK]",
+  "vocab_type": "spm"
+}

trace_compat.py

@@ -0,0 +1,138 @@
+"""Tracing-only DeBERTa relative attention for fixed batch 1.
+
+This is adapted from Hugging Face Transformers' Apache-2.0
+``DisentangledSelfAttention.disentangled_attention_bias``. The trained weights
+are untouched. The only semantic change is a literal repeat count of one for
+the fixed B=1 Core ML export, avoiding an aten::Int conversion failure.
+"""
+
+from __future__ import annotations
+
+import math
+from contextlib import contextmanager
+
+import torch
+from transformers.models.deberta_v2.modeling_deberta_v2 import build_relative_position
+
+
+def constant_attention_scale(query_layer: torch.Tensor, scale_factor: int) -> torch.Tensor:
+    """Static head-width sqrt, equivalent to Transformers' fp32 calculation."""
+    return torch.tensor(
+        math.sqrt(query_layer.shape[-1] * scale_factor),
+        dtype=torch.float32,
+        device=query_layer.device,
+    )
+
+
+def batch_one_disentangled_attention_bias(
+    self, query_layer, key_layer, relative_pos, rel_embeddings, scale_factor
+):
+    if query_layer.shape[0] != self.num_attention_heads:
+        raise ValueError("this tracing path requires batch size one")
+    if relative_pos is None:
+        relative_pos = build_relative_position(
+            query_layer,
+            key_layer,
+            bucket_size=self.position_buckets,
+            max_position=self.max_relative_positions,
+        )
+    if relative_pos.dim() == 2:
+        relative_pos = relative_pos.unsqueeze(0).unsqueeze(0)
+    elif relative_pos.dim() == 3:
+        relative_pos = relative_pos.unsqueeze(1)
+    elif relative_pos.dim() != 4:
+        raise ValueError(f"relative position ids must have 2, 3 or 4 dims; got {relative_pos.dim()}")
+
+    att_span = self.pos_ebd_size
+    relative_pos = relative_pos.to(device=query_layer.device, dtype=torch.long)
+    rel_embeddings = rel_embeddings[: att_span * 2, :].unsqueeze(0)
+    if self.share_att_key:
+        pos_query_layer = self.transpose_for_scores(
+            self.query_proj(rel_embeddings), self.num_attention_heads
+        ).repeat(1, 1, 1)
+        pos_key_layer = self.transpose_for_scores(
+            self.key_proj(rel_embeddings), self.num_attention_heads
+        ).repeat(1, 1, 1)
+    else:
+        if "c2p" in self.pos_att_type:
+            pos_key_layer = self.transpose_for_scores(
+                self.pos_key_proj(rel_embeddings), self.num_attention_heads
+            ).repeat(1, 1, 1)
+        if "p2c" in self.pos_att_type:
+            pos_query_layer = self.transpose_for_scores(
+                self.pos_query_proj(rel_embeddings), self.num_attention_heads
+            ).repeat(1, 1, 1)
+
+    score = 0
+    if "c2p" in self.pos_att_type:
+        scale = constant_attention_scale(pos_key_layer, scale_factor)
+        c2p_att = torch.bmm(query_layer, pos_key_layer.transpose(-1, -2))
+        c2p_pos = torch.clamp(relative_pos + att_span, 0, att_span * 2 - 1)
+        c2p_att = torch.gather(
+            c2p_att,
+            dim=-1,
+            index=c2p_pos.squeeze(0).expand(
+                [query_layer.size(0), query_layer.size(1), relative_pos.size(-1)]
+            ),
+        )
+        score += c2p_att / scale.to(dtype=c2p_att.dtype)
+
+    if "p2c" in self.pos_att_type:
+        scale = constant_attention_scale(pos_query_layer, scale_factor)
+        if query_layer.shape[-2] != key_layer.shape[-2]:
+            raise ValueError("fixed classification encoder requires equal query and key lengths")
+        r_pos = relative_pos
+        p2c_pos = torch.clamp(-r_pos + att_span, 0, att_span * 2 - 1)
+        p2c_att = torch.bmm(key_layer, pos_query_layer.transpose(-1, -2))
+        p2c_att = torch.gather(
+            p2c_att,
+            dim=-1,
+            index=p2c_pos.squeeze(0).expand(
+                [query_layer.size(0), key_layer.size(-2), key_layer.size(-2)]
+            ),
+        ).transpose(-1, -2)
+        score += p2c_att / scale.to(dtype=p2c_att.dtype)
+    return score
+
+
+def install_trace_compatibility() -> None:
+    """Install fixed-shape tracing helpers; call only after native baseline."""
+    import gliformer.backbones.deberta_2d as gliformer_deberta
+    import transformers.models.deberta_v2.modeling_deberta_v2 as hf_deberta
+
+    gliformer_deberta.scaled_size_sqrt = constant_attention_scale
+    hf_deberta.scaled_size_sqrt = constant_attention_scale
+    gliformer_deberta.LayoutDisentangledSelfAttention.disentangled_attention_bias = (
+        batch_one_disentangled_attention_bias
+    )
+
+
+@contextmanager
+def finite_fp16_mask():
+    """Trace finite attention-mask fill instead of fp32 minimum overflowing to FP16 -inf.
+
+    Valid attention rows retain the same softmax in fp32. The unmasked native
+    baseline and patched model are compared on every parity fixture.
+    """
+    original = torch.finfo
+
+    class FiniteFinfo:
+        def __init__(self, real):
+            self.real = real
+
+        def __getattr__(self, name):
+            return getattr(self.real, name)
+
+        @property
+        def min(self):
+            return -10000.0
+
+    def patched(dtype):
+        info = original(dtype)
+        return FiniteFinfo(info) if dtype.is_floating_point else info
+
+    torch.finfo = patched
+    try:
+        yield
+    finally:
+        torch.finfo = original

verify.py

@@ -0,0 +1,68 @@
+"""Compare the exported Jeff FP16 Core ML package with its trained native model."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import time
+from pathlib import Path
+
+import coremltools as ct
+import numpy as np
+import torch
+from huggingface_hub import snapshot_download
+from jeff.backends.torch_backend import TorchBackend
+
+from export import FIXTURES, REVISION, SOURCE, make_batch, model_inputs
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--precision", choices=("fp16", "fp32"), default="fp16")
+    args = parser.parse_args()
+    torch.set_num_threads(2)
+    checkpoint = snapshot_download(SOURCE, revision=REVISION, local_files_only=True)
+    backend = TorchBackend(checkpoint, device="cpu", dtype="float32", attn_kernel="eager", batch_size=1)
+    package = Path(f"build/JeffDecision-L128-{args.precision.upper()}.mlpackage")
+    model = ct.models.MLModel(str(package), compute_units=ct.ComputeUnit.CPU_ONLY)
+    output_name = model.get_spec().description.output[0].name
+    results = []
+    for name, text, group in FIXTURES:
+        batch = make_batch(backend, text, group)
+        native = backend.model.model(**batch, include_media=False).cat_logits.detach().float().numpy()[0]
+        tensor_inputs = model_inputs(batch, backend.model.config)
+        inputs = {
+            key: tensor.numpy()
+            for key, tensor in zip(
+                ("input_ids", "attention_mask", "parent_position", "category_positions"),
+                tensor_inputs,
+            )
+        }
+        start = time.perf_counter()
+        converted = model.predict(inputs)[output_name][0, :len(group.labels)].astype(np.float32)
+        elapsed_ms = (time.perf_counter() - start) * 1000
+        error = float(np.max(np.abs(native - converted))) if np.isfinite(converted).all() else float("inf")
+        agreement = int(np.argmax(native) == np.argmax(converted))
+        results.append({
+            "name": name,
+            "labels": list(group.labels),
+            "native_logits": native.tolist(),
+            "coreml_logits": converted.tolist(),
+            "max_logit_error": error,
+            "top_label_agreement": bool(agreement),
+            "coreml_wall_ms": elapsed_ms,
+        })
+        print(
+            f"{name}: logits={converted.tolist()}, error={error:.6f}, "
+            f"top_label={bool(agreement)}, wall_ms={elapsed_ms:.1f}",
+            flush=True,
+        )
+    Path(f"build/coreml-parity-{args.precision}.json").write_text(json.dumps(results, indent=2) + "\n")
+    if not all(row["top_label_agreement"] for row in results):
+        raise AssertionError("Core ML changed the chosen label on a parity fixture")
+    if max(row["max_logit_error"] for row in results) > 0.25:
+        raise AssertionError("Core ML logit error exceeds the 0.25 tolerance")
+
+
+if __name__ == "__main__":
+    main()