D1A-E2B v0.6 Core ML: text model, 5 int8 chunks, L=512, with the reference runner (#224)

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,91 @@
+---
+license: apache-2.0
+base_model: google/gemma-4-E2B
+tags: [coreml, apple-neural-engine, decision-model, d1a]
+---
+# D1A-E2B v0.6 · Core ML (text, ≤ 512 tokens)
+
+D1A-E2B v0.6 for the Apple Neural Engine. D1A is a small decision model: a document and typed questions go in, and a
+calibrated probability for every option comes out, in one pass, with no generated text ([jonpol01/d1a](https://github.com/jonpol01/d1a)).
+This repository is the text model converted to Core ML; the conversion is described in jonpol01/d1a#224.
+
+## Status
+
+- **Correct at ≤ 512 tokens.** On 28 real requests (42 questions), no clear-cut answer changes against the fp32 PyTorch
+  rebuild of the same model, on an M1 Max and an M4.
+- **Speed (median request):** M4 Neural Engine 0.96 s (MLX on the GPU: 0.20 s); M1 Max 9.7 s (MLX: 0.17 s). The M1's Neural Engine does not take these chunks, so on M1-class Macs they run on the CPU; the M4 runs them on its Neural Engine.
+- **Use it** where the GPU should stay free or power matters (iPhone, iPad). On a Mac, the MLX build is faster.
+- **Not included yet:** photo, voice and video (their encoders are separate models, still to convert), and inputs longer than
+  512 tokens (2,048-token programs give wrong answers on the M4 Neural Engine).
+
+## Run it
+
+```bash
+pip install coremltools numpy torch transformers safetensors "d1a @ git+https://github.com/jonpol01/d1a"
+hf download JohnP1/d1a-e2b-coreml --revision v0.6 --local-dir d1a-e2b-coreml
+python d1a-e2b-coreml/reference/run.py d1a-e2b-coreml request.json
+```
+
+`request.json` is a System One request, for example:
+
+```json
+{"state": "Shoes arrived two weeks late and in the wrong size. Also I see two charges on my card.",
+ "questions": {"team": {"type": "choice", "instructions": "Which team should handle this?", "criteria": {"returns": null, "shipping": null, "billing": null}},
+               "urgent": {"type": "noul", "instructions": "Does this need a reply today?"}}}
+```
+
+The first run compiles the chunks into `coreml/.compiled/` (a few minutes); later runs load them.
+
+## Files
+
+| Path | What |
+|---|---|
+| `coreml/L512_<a>-<b>.mlpackage` | 5 chunks of the 35-layer text model, int8 weights (per block of 64). Chunk 7-15 outputs the shared keys and values that chunks 15-35 read. |
+| `embeddings.safetensors` | the token and per-layer embedding tables (MLX affine-quantized, from JohnP1/d1a-e2b-mlx-q8@v0.6), read row by row on the host |
+| `head.safetensors`, `d1a_config.json` | the pointer head, the temperature and the pad and delimiter ids |
+| `config.json`, `tokenizer.json`, `tokenizer_config.json` | the Gemma 4 model config and tokenizer |
+| `io.json` | each chunk's input names |
+| `reference/` | Python reference: `run.py` (one request), `workers.py` (one process per chunk), `d1a_coreml.py` (host-side embeddings and masks) |
+
+## How a request runs
+
+1. Encode it with D1A's layout (`d1a.core.encoding.encode`) and right-pad it to 512 tokens with the pad token.
+2. Look up the token embeddings (scaled by √hidden) and the per-layer embeddings (scaled by √256) on the host.
+3. Build the block-causal masks (full and sliding), clamped to −1e4.
+4. Run the chunks in order, each in its own process. Several large Neural Engine programs in one process can crash the
+   runtime.
+5. Read the hidden states at each question's `<decide>` and `</opt>` positions, score them with the head, divide by the
+   temperature, and take the softmax.
+
+macOS 15 / iOS 18 or later.
+
+## D1A builds compared
+
+Measured on 2026-10-11 on the same 28 real requests (42 questions, all ≤ 512 tokens), on an Apple M1 Max and a Mac mini M4.
+
+| | E2B · MLX | E2B · Core ML | E4B · MLX | E4B · Core ML |
+|---|---|---|---|---|
+| Version | v0.6 | v0.6 | v0.6 | v0.6 |
+| Repository | JohnP1/d1a-e2b-mlx-q8 | JohnP1/d1a-e2b-coreml | JohnP1/d1a-e4b-mlx-q8 | JohnP1/d1a-e4b-coreml |
+| Runs on | Mac GPU | Neural Engine (iPhone, iPad, M4-class Macs) | Mac GPU | Neural Engine |
+| Text | yes | yes | yes | yes |
+| Photo, voice, video (16 frames, no sound) | yes | not yet (next) | yes | not yet |
+| Input length | long documents | ≤ 512 tokens | long documents | ≤ 512 tokens |
+| Questions | yes/no, choice, score, with a calibrated probability for each option | same | same | same |
+| Languages | questions in English; documents in English or Japanese | same | same | same |
+| PR labels, 487 newest PRs: type / severity | 87.5% / 77.2%; 0 of 30 P0/P1 found | same model | 90.6% / 79.1%; 13 of 30 P0/P1 | same model |
+| Size | 3.8 GB + 1.0 GB media | 3.5 GB (1.8 GB chunks + 1.6 GB embeddings) | 6.6 GB + 1.0 GB media | 6.1 GB (3.9 GB chunks + 2.1 GB embeddings) |
+| Median request, M4 (p95) | **0.20 s** (0.26) | 0.96 s (0.97) | 0.63 s (0.89) | 1.80 s (1.91) |
+| Median request, M1 Max (p95) | **0.17 s** (0.23) | 9.7 s (10.0) ² | 0.81 s (1.33) | 2.93 s (2.95) |
+| Load, M4 / M1 Max | 2 s / 9 s | 40 s / 247 s | 3 s / – | 52 s / 106 s |
+| Answer changes vs fp32 PyTorch ¹ | 0 of 37 clear-cut (M4: 1 near-tie of 42) | 0 of 37 clear-cut (M4: 1 near-tie of 42) | 0 of 38 clear-cut | 0 of 38 clear-cut |
+| Photo / voice / video, median (M1 Max) | 0.83 s over the 14 demo requests | – | – | – |
+
+¹ Each build is compared with the fp32 PyTorch rebuild of the same v0.6 model. "Clear-cut" means the reference's top
+answer leads by more than 0.05; a near-tie can flip either way.
+² On the M1 Max's Neural Engine, the E2B chunks fall back to the CPU: each chunk takes the same time with or without the
+Neural Engine. The M4 runs the same files on its Neural Engine.
+
+## License
+
+Apache-2.0. Base model: Gemma 4 by Google (Apache-2.0). D1A is built on Kev by Jared Palmer (Apache-2.0).

config.json

@@ -0,0 +1,164 @@
+{
+    "architectures": [
+        "Gemma4ForConditionalGeneration"
+    ],
+    "audio_config": {
+        "_name_or_path": "",
+        "architectures": null,
+        "attention_chunk_size": 12,
+        "attention_context_left": 13,
+        "attention_context_right": 0,
+        "attention_invalid_logits_value": -1000000000.0,
+        "attention_logit_cap": 50.0,
+        "chunk_size_feed_forward": 0,
+        "conv_kernel_size": 5,
+        "dtype": "bfloat16",
+        "gradient_clipping": 10000000000.0,
+        "hidden_act": "silu",
+        "hidden_size": 1024,
+        "id2label": {
+            "0": "LABEL_0",
+            "1": "LABEL_1"
+        },
+        "initializer_range": 0.02,
+        "is_encoder_decoder": false,
+        "label2id": {
+            "LABEL_0": 0,
+            "LABEL_1": 1
+        },
+        "model_type": "gemma4_audio",
+        "num_attention_heads": 8,
+        "num_hidden_layers": 12,
+        "output_attentions": false,
+        "output_hidden_states": false,
+        "output_proj_dims": 1536,
+        "problem_type": null,
+        "residual_weight": 0.5,
+        "return_dict": true,
+        "rms_norm_eps": 1e-06,
+        "subsampling_conv_channels": [
+            128,
+            32
+        ],
+        "use_clipped_linears": true
+    },
+    "audio_token_id": 258881,
+    "boa_token_id": 256000,
+    "boi_token_id": 255999,
+    "dtype": "bfloat16",
+    "eoa_token_id": 258883,
+    "eoa_token_index": 258883,
+    "eoi_token_id": 258882,
+    "eos_token_id": 1,
+    "image_token_id": 258880,
+    "initializer_range": 0.02,
+    "model_type": "gemma4",
+    "quantization": {
+        "group_size": 64,
+        "bits": 8,
+        "mode": "affine",
+        "language_model.model.embed_tokens_per_layer": {
+            "group_size": 64,
+            "bits": 4,
+            "mode": "affine"
+        }
+    },
+    "quantization_config": {
+        "group_size": 64,
+        "bits": 8,
+        "mode": "affine",
+        "language_model.model.embed_tokens_per_layer": {
+            "group_size": 64,
+            "bits": 4,
+            "mode": "affine"
+        }
+    },
+    "text_config": {
+        "attention_bias": false,
+        "attention_dropout": 0.0,
+        "attention_k_eq_v": false,
+        "bos_token_id": 2,
+        "dtype": "bfloat16",
+        "enable_moe_block": false,
+        "eos_token_id": 1,
+        "expert_intermediate_size": null,
+        "final_logit_softcapping": 30.0,
+        "global_head_dim": 512,
+        "head_dim": 256,
+        "hidden_activation": "gelu_pytorch_tanh",
+        "hidden_size": 1536,
+        "hidden_size_per_layer_input": 256,
+        "initializer_range": 0.02,
+        "intermediate_size": 6144,
+        "layer_types": [
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "full_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "full_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "full_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "full_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "full_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "full_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "full_attention"
+        ],
+        "max_position_embeddings": 131072,
+        "model_type": "gemma4_text",
+        "num_attention_heads": 8,
+        "num_experts": null,
+        "num_global_key_value_heads": null,
+        "num_hidden_layers": 35,
+        "num_key_value_heads": 1,
+        "num_kv_shared_layers": 20,
+        "pad_token_id": 0,
+        "rms_norm_eps": 1e-06,
+        "rope_parameters": {
+            "full_attention": {
+                "partial_rotary_factor": 0.25,
+                "rope_theta": 1000000.0,
+                "rope_type": "proportional"
+            },
+            "sliding_attention": {
+                "rope_theta": 10000.0,
+                "rope_type": "default"
+            }
+        },
+        "sliding_window": 512,
+        "tie_word_embeddings": true,
+        "top_k_experts": null,
+        "use_bidirectional_attention": null,
+        "use_cache": true,
+        "use_double_wide_mlp": true,
+        "vocab_size": 262144,
+        "vocab_size_per_layer_input": 262144
+    },
+    "tie_word_embeddings": true,
+    "transformers_version": "5.5.0.dev0",
+    "video_token_id": 258884,
+    "vision_soft_tokens_per_image": 280
+}

coreml/L512_00-07.mlpackage/Data/com.apple.CoreML/model.mlmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:63cdc56003526dc6bc1922b1f5245e14fda32e472ff1c22f0b1bace80ba61235
+size 254034

coreml/L512_00-07.mlpackage/Data/com.apple.CoreML/weights/weight.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4932ca078e73e04eb9bacb44aafee7356ff9b418068c7a0d2bb3325dbcadfe34
+size 287097152

coreml/L512_00-07.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "15E6ABF0-3B7A-45EB-AEFE-575087ED5B8B": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        },
+        "186F9637-EC03-4170-A083-0E8F08B64E70": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        }
+    },
+    "rootModelIdentifier": "15E6ABF0-3B7A-45EB-AEFE-575087ED5B8B"
+}

coreml/L512_07-15.mlpackage/Data/com.apple.CoreML/model.mlmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9f64f99c757d484f0d57e5746f3ede35e2d44d377de494e72f8e9d532b9afb29
+size 286085

coreml/L512_07-15.mlpackage/Data/com.apple.CoreML/weights/weight.bin

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

coreml/L512_07-15.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "C34EF27C-FE7F-43D9-B3CB-EF07BA6ED05E": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        },
+        "E1CCD0C2-4CB7-4016-8EA6-D0C9512C59E7": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        }
+    },
+    "rootModelIdentifier": "C34EF27C-FE7F-43D9-B3CB-EF07BA6ED05E"
+}

coreml/L512_15-22.mlpackage/Data/com.apple.CoreML/model.mlmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:98899d9bcfbe1941ae4533d9b2db80647f205ebe4d14a9a039f0850be19db094
+size 183109

coreml/L512_15-22.mlpackage/Data/com.apple.CoreML/weights/weight.bin

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

coreml/L512_15-22.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "8BF3EE60-129B-43DC-ADBC-76FCC2AE0F22": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        },
+        "D588AC66-1980-4F48-A80E-D22E2DC5FC04": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        }
+    },
+    "rootModelIdentifier": "8BF3EE60-129B-43DC-ADBC-76FCC2AE0F22"
+}

coreml/L512_22-29.mlpackage/Data/com.apple.CoreML/model.mlmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5d99163080dcb2cdf8fd583083db5e1b11a165a6ae47856186e5a45e455e6586
+size 183113

coreml/L512_22-29.mlpackage/Data/com.apple.CoreML/weights/weight.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:99d4cc60a0b5d645595238fd279e7d2eb0c939de902e7f80bc9eb5f005ed3503
+size 473557760

coreml/L512_22-29.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "5CBC619A-F5EC-4532-AE9A-E730428289EC": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        },
+        "E85695E6-50B8-4D34-AFD3-8C9EAF6EF828": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        }
+    },
+    "rootModelIdentifier": "E85695E6-50B8-4D34-AFD3-8C9EAF6EF828"
+}

coreml/L512_29-35.mlpackage/Data/com.apple.CoreML/model.mlmodel

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

coreml/L512_29-35.mlpackage/Data/com.apple.CoreML/weights/weight.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:81a752d4ccecd24523c1e5bc5c5859c07d44233a75d229fe846912e44ef90cea
+size 413443264

coreml/L512_29-35.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "4D7E2B93-F4E3-4337-A629-BB3697573F32": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        },
+        "64D6F739-0E0A-4E79-97A2-DCFB02714BFA": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        }
+    },
+    "rootModelIdentifier": "64D6F739-0E0A-4E79-97A2-DCFB02714BFA"
+}

d1a_config.json

@@ -0,0 +1,142 @@
+{
+  "format": "d1a-mlx",
+  "format_version": 1,
+  "base": "google/gemma-4-E2B",
+  "base_revision": "d29ff6b45f081a49ee2733a859c9c9c2d95d1a6f",
+  "source": "JohnP1/d1a-e2b",
+  "source_revision": "v0.6",
+  "adapter_sha256": "247f222314c8ceee158140d38a0be1c304da81719034ef6e71ee02506e334c1b",
+  "lora": 16,
+  "head_dim": 256,
+  "hidden_size": 1536,
+  "temperature": 1.447269237440378,
+  "leading_ids": [
+    2
+  ],
+  "bos_id": 2,
+  "delimiter_ids": [
+    6,
+    7,
+    8,
+    9,
+    10
+  ],
+  "pad_id": 0,
+  "dtype": "bfloat16",
+  "quantization": {
+    "group_size": 64,
+    "bits": 8,
+    "mode": "affine",
+    "language_model.model.embed_tokens_per_layer": {
+      "group_size": 64,
+      "bits": 4,
+      "mode": "affine"
+    },
+    "embeddings": true
+  },
+  "head": "head.safetensors",
+  "temperature_fit": {
+    "rows": [
+      "cal/rows.json",
+      "calpr/rows.json"
+    ],
+    "n": 3178,
+    "method": "min micro mean NLL over a 121-point log grid 0.25..4",
+    "cross_validation": {
+      "raw": {
+        "n": 3178,
+        "ece": 0.0505561988092814,
+        "brier": 0.2952905648331223,
+        "nll": 0.5648973549274169,
+        "confident_error_rate": 0.02957835116425425,
+        "coverage_at_5pct_error": 0.5613593455003146
+      },
+      "out_of_fold": {
+        "n": 3178,
+        "ece": 0.023704548579686884,
+        "brier": 0.2896952220920997,
+        "nll": 0.5324486179012022,
+        "confident_error_rate": 0.017935808684707365,
+        "coverage_at_5pct_error": 0.5556954059156702
+      },
+      "ece_ci95": {
+        "raw": [
+          0.04102750015575881,
+          0.06377306530700594
+        ],
+        "out_of_fold": [
+          0.016511069955034523,
+          0.03860363764230206
+        ],
+        "delta": [
+          -0.04016350379193325,
+          -0.010727297381358629
+        ]
+      },
+      "separated": true,
+      "temperatures": [
+        1.447269237440378,
+        1.447269237440378,
+        1.4142135623730947,
+        1.4810975522865641,
+        1.447269237440378
+      ],
+      "folds": 5,
+      "seed": 0,
+      "samples": 1000,
+      "groups": 1670,
+      "unit": "source-stratified (source, group); sibling questions and variants stay together"
+    },
+    "interval": {
+      "level": 0.9,
+      "samples": 2000,
+      "seed": 0,
+      "lower": 1.381912879967776,
+      "upper": 1.5157165665103978,
+      "questions": 3178
+    },
+    "fit_rows": [
+      {
+        "rows": "cal/rows.json",
+        "suite": "evals/v7/decision-v7",
+        "split": "calibration",
+        "questions": 1148
+      },
+      {
+        "rows": "calpr/rows.json",
+        "suite": "evals/d1a/pr-labels",
+        "split": "development",
+        "questions": 2030
+      }
+    ],
+    "training_suites": [
+      "evals/devtools-v1",
+      "evals/documents-v1",
+      "evals/hard-v1",
+      "evals/v7/decision-v7"
+    ],
+    "training_partitions": [
+      "evals/d1a/ja-jglue:train",
+      "evals/d1a/pr-labels:train",
+      "evals/d1a/pr-labels:train-blast",
+      "evals/d1a/pr-labels:train-ja",
+      "evals/d1a/routing:factory-train",
+      "evals/d1a/routing:generic-train"
+    ],
+    "same_corpus": [
+      "calpr/rows.json: SAME CORPUS: evals/d1a/pr-labels:development is an eval partition of evals/d1a/pr-labels, whose train partitions ['evals/d1a/pr-labels:train', 'evals/d1a/pr-labels:train-blast', 'evals/d1a/pr-labels:train-ja'] the checkpoint trained on"
+    ],
+    "in_distribution": {
+      "allowed": true,
+      "problems": [
+        "cal/rows.json: evals/v7/decision-v7 is training data of the checkpoint",
+        "cal/rows.json: it pools 12 training source(s) ['agnews', 'amazon', 'banking77', 'boolq', 'compositional', 'dbpedia14', 'imdb', 'legacy_policy'] ...",
+        "cal/rows.json: it reads the calibration partition of evals/v7/decision-v7, a training corpus",
+        "calpr/rows.json: SAME CORPUS: evals/d1a/pr-labels:development is an eval partition of evals/d1a/pr-labels, whose train partitions ['evals/d1a/pr-labels:train', 'evals/d1a/pr-labels:train-blast', 'evals/d1a/pr-labels:train-ja'] the checkpoint trained on",
+        "cannot list the sources of --data /work/data/train-e2b-v06.jsonl (outside evals/, no sidecar train-e2b-v06.jsonl.json naming its sources), training data of d1"
+      ]
+    }
+  },
+  "use_case_temperatures": {},
+  "use_case_temperature_fits": {}
+}

embeddings.safetensors

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

head.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b56c4f9da06b75a359da342e6002c6c73896033925b88e9e13c3757878e7447
+size 3148080

io.json

@@ -0,0 +1 @@
+{"chunks": ["00-07", "07-15", "15-22", "22-29", "29-35"], "inputs": {"00-07": ["emb", "ple_tok", "pos", "mask_full", "mask_slide"], "07-15": ["h_in", "ple", "pos", "mask_full", "mask_slide"], "15-22": ["h_in", "ple", "pos", "mask_full", "mask_slide", "k_full_attention", "v_full_attention", "k_sliding_attention", "v_sliding_attention"], "22-29": ["h_in", "ple", "pos", "mask_full", "mask_slide", "k_full_attention", "v_full_attention", "k_sliding_attention", "v_sliding_attention"], "29-35": ["h_in", "ple", "pos", "mask_full", "mask_slide", "k_full_attention", "v_full_attention", "k_sliding_attention", "v_sliding_attention", "picks"]}}

reference/d1a_coreml.py

@@ -0,0 +1,184 @@
+"""D1A (Gemma 4 text) -> Core ML pilot (jonpol01/d1a#224): rebuild a D1A MLX export as a torch module, split into
+  - host side: the two embedding tables (token embedding, per-layer embedding: 2.75B of E2B's 4.6B parameters) read
+    row by row from the export's safetensors (dequantized per request, never held whole), and
+  - Core: everything that computes (per-layer projection, 35 decoder layers, final norm, pointer head), which is what
+    gets converted to Core ML.
+Inputs follow D1A's packed encoding (d1a.core.encoding.encode): one record, block-causal mask, positions restarting per
+question after the state, readout at each question's <decide> and </opt> tokens."""
+import json
+import math
+from pathlib import Path
+
+import numpy as np
+import torch
+from safetensors import safe_open
+from torch import nn
+from transformers import AutoConfig
+from transformers.models.gemma4 import modeling_gemma4 as G
+
+PREFIX = "language_model.model."
+
+
+def dequant(packed, scales, biases, bits, group=64):
+    """MLX affine dequantization: packed uint32 [out, in*bits/32] (lowest bits first) -> float32 [out, in]."""
+    p = packed.view(torch.int32).to(torch.int64) & 0xFFFFFFFF
+    per = 32 // bits
+    shifts = torch.arange(per, dtype=torch.int64) * bits
+    q = ((p[..., None] >> shifts) & ((1 << bits) - 1)).reshape(p.shape[0], -1).to(torch.float32)
+    s, b = scales.to(torch.float32), biases.to(torch.float32)
+    return (q.view(q.shape[0], -1, group) * s[..., None] + b[..., None]).reshape(q.shape[0], -1)
+
+
+class _Shards:
+    """model.safetensors or model-0000N-of-0000M.safetensors (via model.safetensors.index.json), read like one safe_open."""
+
+    def __init__(self, path):
+        idx = Path(path) / "model.safetensors.index.json"
+        if idx.exists():
+            wm = json.loads(idx.read_text())["weight_map"]
+            files = {f: safe_open(str(Path(path) / f), "pt") for f in sorted(set(wm.values()))}
+            self.where = {k: files[f] for k, f in wm.items()}
+        else:   # a released Core ML bundle ships only the embedding tables, as embeddings.safetensors
+            name = "model.safetensors" if (Path(path) / "model.safetensors").exists() else "embeddings.safetensors"
+            f = safe_open(str(Path(path) / name), "pt")
+            self.where = {k: f for k in f.keys()}
+
+    def keys(self):
+        return self.where.keys()
+
+    def get_tensor(self, k):
+        return self.where[k].get_tensor(k)
+
+    def get_slice(self, k):
+        return self.where[k].get_slice(k)
+
+
+class Export:
+    """A D1A MLX export directory (config.json, d1a_config.json, model.safetensors, head.safetensors)."""
+
+    def __init__(self, path):
+        self.path = Path(path)
+        self.cfg = AutoConfig.from_pretrained(self.path)
+        self.tc = self.cfg.text_config
+        self.tc._attn_implementation = "eager"
+        self.d1a = json.loads((self.path / "d1a_config.json").read_text())
+        q = self.d1a["quantization"]
+        self.bits = {"default": (q["bits"], q["group_size"]), "embed_tokens_per_layer": (q[PREFIX.replace("model.", "") + "model.embed_tokens_per_layer"]["bits"] if False else q["language_model.model.embed_tokens_per_layer"]["bits"], q["language_model.model.embed_tokens_per_layer"]["group_size"])}
+        self.f = _Shards(self.path)
+
+    def bits_of(self, name):
+        return self.bits["embed_tokens_per_layer"] if name.endswith("embed_tokens_per_layer") else self.bits["default"]
+
+    def tensor(self, name):
+        """A float32 tensor for `name` (without PREFIX), dequantized if it is stored quantized."""
+        full = PREFIX + name
+        keys = set(self.f.keys())
+        if full + ".weight" in keys and full + ".scales" in keys:
+            bits, group = self.bits_of(name)
+            return dequant(self.f.get_tensor(full + ".weight"), self.f.get_tensor(full + ".scales"), self.f.get_tensor(full + ".biases"), bits, group)
+        return self.f.get_tensor(full).to(torch.float32)
+
+    def text_model(self, dtype=torch.float32):
+        """The Gemma 4 text trunk with every non-embedding weight loaded (embedding tables stay empty: host side)."""
+        with torch.device("meta"):
+            m = G.Gemma4TextModel(self.tc)
+        m = m.to_empty(device="cpu")
+        state = {}
+        keys = set(self.f.keys())
+        for name, p in m.named_parameters():
+            if name.startswith("embed_tokens"):
+                continue
+            base = name[: -len(".weight")] if name.endswith(".weight") else name
+            if PREFIX + base + ".scales" in keys:
+                state[name] = self.tensor(base)
+            else:
+                state[name] = self.f.get_tensor(PREFIX + name).to(torch.float32)
+            assert state[name].shape == p.shape, (name, state[name].shape, p.shape)
+        for name, b in m.named_buffers():
+            if PREFIX + name in keys:
+                state[name] = self.f.get_tensor(PREFIX + name).to(torch.float32)
+        missing, unexpected = m.load_state_dict(state, strict=False)
+        missing = [k for k in missing if not k.startswith("embed_tokens")]
+        assert not missing and not unexpected, (missing, unexpected)
+        # buffers transformers computes itself (rotary inv_freq, embed_scale): re-run their init on CPU
+        m.rotary_emb = G.Gemma4TextRotaryEmbedding(self.tc)
+        return m.to(dtype).eval()
+
+    def head(self):
+        h = safe_open(str(self.path / "head.safetensors"), "pt")
+        dp, d = h.get_tensor("q.weight").shape                     # pointer width, hidden size (E2B 1536, E4B 2560)
+        q, k = nn.Linear(d, dp), nn.Linear(d, dp)
+        q.weight.data, q.bias.data = h.get_tensor("q.weight").float(), h.get_tensor("q.bias").float()
+        k.weight.data, k.bias.data = h.get_tensor("k.weight").float(), h.get_tensor("k.bias").float()
+        return q, k
+
+
+class HostEmbed:
+    """Row lookups of the two embedding tables, dequantized per request: emb = E[ids] * sqrt(hidden) and
+    ple = E_pl[ids] * sqrt(ple_dim), reshaped [L, layers, ple_dim] (transformers' get_per_layer_inputs)."""
+
+    def __init__(self, ex):
+        self.ex, self.tc = ex, ex.tc
+        self.f = ex.f
+
+    def _rows(self, name, ids):
+        full = PREFIX + name
+        bits, group = self.ex.bits_of(name)
+        w, s, b = (self.f.get_slice(full + suf) for suf in (".weight", ".scales", ".biases"))
+        uniq, inv = np.unique(np.asarray(ids), return_inverse=True)
+        rows = [dequant(w[int(i): int(i) + 1], s[int(i): int(i) + 1], b[int(i): int(i) + 1], bits, group) for i in uniq]
+        return torch.cat(rows)[torch.as_tensor(inv)]
+
+    def __call__(self, ids):
+        emb = self._rows("embed_tokens", ids) * (self.tc.hidden_size ** 0.5)
+        ple = self._rows("embed_tokens_per_layer", ids) * (self.tc.hidden_size_per_layer_input ** 0.5)
+        return emb[None], ple.view(1, len(ids), self.tc.num_hidden_layers, self.tc.hidden_size_per_layer_input)
+
+
+def masks_for(enc, window, L=None, dtype=torch.float32):
+    """D1A's packed masks (d1a.backends.torch.branch_masks) for one record, padded to L: additive [1, 1, L, L]."""
+    seg, pos = enc["seg"], enc["pos"]
+    n = len(seg)
+    L = L or n
+    s = torch.full((L,), -1)
+    s[:n] = torch.tensor(seg)
+    p = torch.zeros(L, dtype=torch.long)
+    p[:n] = torch.tensor(pos)
+    causal = torch.tril(torch.ones(L, L, dtype=torch.bool))
+    same = (s[None, :] == s[:, None]) | (s[None, :] == 0)
+    allow = (causal & same & (s != -1)[None, :]) | torch.eye(L, dtype=torch.bool)
+    neg = torch.finfo(dtype).min
+    full = torch.zeros(L, L, dtype=dtype).masked_fill(~allow, neg)
+    far = (p[:, None] - p[None, :] >= window) & ~torch.eye(L, dtype=torch.bool)
+    slide = full.masked_fill(far, neg)
+    return full[None, None], slide[None, None], p[None]
+
+
+class Core(nn.Module):
+    """Everything that computes: per-layer projection, decoder layers, final norm, pointer head.
+    forward(emb [1,L,d], ple_tok [1,L,layers,256], pos [1,L] int, mask_full, mask_slide [1,1,L,L],
+            decide [Q] int, opts [Q,K] int, opt_valid [Q,K] float (1 used, 0 padding)) -> logits [Q,K]
+    (raw, temperature not applied; padded options get a large negative)."""
+
+    def __init__(self, text, head_q, head_k, silu_tanh=True):
+        super().__init__()
+        self.text, self.hq, self.hk = text, head_q, head_k
+        self.types = text.config.layer_types
+        self.scale = 1 / math.sqrt(head_q.out_features)
+
+    def forward(self, emb, ple_tok, pos, mask_full, mask_slide, decide, opts, opt_valid):
+        t = self.text
+        ple = t.project_per_layer_inputs(emb, ple_tok)
+        rope = {lt: t.rotary_emb(emb, pos, lt) for lt in t.unique_layer_types}
+        masks = {"full_attention": mask_full, "sliding_attention": mask_slide}
+        from collections import UserDict
+        kv = UserDict()
+        h = emb
+        for i, layer in enumerate(t.layers):
+            h = layer(h, ple[:, :, i, :], shared_kv_states=kv, position_embeddings=rope[self.types[i]],
+                      attention_mask=masks[self.types[i]], position_ids=pos)
+        h = t.norm(h)[0]                                    # [L, d]
+        q = self.hq(h[decide])                              # [Q, dp]
+        k = self.hk(h[opts.reshape(-1)]).view(*opts.shape, -1)   # [Q, K, dp]
+        z = (k * q[:, None, :]).sum(-1) * self.scale
+        return z + (opt_valid - 1.0) * 1e4

reference/pipeline.py

@@ -0,0 +1,80 @@
+"""Run a chunked D1A Core ML trunk end to end: host embeddings -> chunk 0 .. chunk N (h, per-layer inputs and the shared
+keys/values passed along) -> hidden states at the picks -> pointer head on the host.
+Usage as a library: Pipeline(export_dir, chunk_dir, L, units).probs(enc) -> [per-question probabilities]."""
+import glob
+import os
+import re
+import sys
+import time
+
+import coremltools as ct
+import numpy as np
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))   # d1a_coreml.py sits next to this file
+from d1a_coreml import Export, HostEmbed, masks_for  # noqa: E402
+
+P_PICKS = 128
+
+
+class Pipeline:
+    def __init__(self, export_dir, chunk_dir, L, units="CPU_AND_NE"):
+        self.ex = Export(export_dir)
+        self.host = HostEmbed(self.ex)
+        self.hq, self.hk = self.ex.head()
+        self.T = self.ex.d1a["temperature"]
+        self.L = L
+        cu = getattr(ct.ComputeUnit, units)
+        self.chunks = []
+        if chunk_dir.endswith(".mlpackage"):   # one multifunction package: functions c00_07, c07_15, ...
+            names = ["c00_07", "c07_15", "c15_22", "c22_29", "c29_35"]
+            for name in names:
+                a, b = map(int, name[1:].split("_"))
+                m = ct.models.MLModel(chunk_dir, compute_units=cu, function_name=name)
+                fn = next(f for f in m.get_spec().description.functions if f.name == name)
+                self.chunks.append((a, b, m, [i.name for i in fn.input], [o.name for o in fn.output]))
+            return
+        files = sorted(glob.glob(f"{chunk_dir}/L{L}_*.mlpackage"))
+        for f in files:
+            a, b = map(int, re.search(r"_(\d+)-(\d+)\.mlpackage$", f).groups())
+            m = ct.models.MLModel(f, compute_units=cu)
+            spec = m.get_spec()
+            self.chunks.append((a, b, m, [i.name for i in spec.description.input], [o.name for o in spec.description.output]))
+
+    def hidden(self, enc):
+        L, n = self.L, len(enc["ids"])
+        # right-pad with the pad token's own embeddings (as d1a's torch backend pads ids), never zero rows: an all-zero row
+        # makes every RMSNorm compute 0 / sqrt(eps), and eps (1e-6) is subnormal in fp16, which the Neural Engine flushes
+        ids = list(enc["ids"]) + [self.ex.d1a.get("pad_id", 0)] * (L - n)
+        emb, ple = self.host(ids)
+        e = emb.numpy().astype(np.float16)
+        p = ple.numpy().astype(np.float16)
+        mf, ms, pos = masks_for(enc, self.ex.tc.sliding_window, L=L)
+        picks = list(enc["decide_idx"]) + [o for oo in enc["opt_idx"] for o in oo]
+        pk = np.zeros(P_PICKS, np.int32); pk[: len(picks)] = picks
+        common = {"pos": pos.numpy().astype(np.int32), "mask_full": mf.clamp(min=-1e4).numpy().astype(np.float16),
+                  "mask_slide": ms.clamp(min=-1e4).numpy().astype(np.float16)}
+        state = {"emb": e, "ple_tok": p, "picks": pk}
+        for a, b, m, ins, outs in self.chunks:
+            feed = {k: np.ascontiguousarray(common[k] if k in common else state[k]) for k in ins}
+            feed = {k: (v.astype(np.float16) if v.dtype == np.float32 else v) for k, v in feed.items()}   # chunk outputs come back fp32
+            out = {k: np.array(v, copy=True) for k, v in m.predict(feed).items()}   # Core ML may reuse its output buffers on the next call
+            if "hidden" in out:
+                return torch.tensor(out["hidden"][: len(picks)], dtype=torch.float32)
+            state["h_in"] = out["h"]
+            if "ple_out" in out:
+                state["ple"] = out["ple_out"]
+            for k, v in out.items():
+                if k.endswith("_out") and k != "ple_out":
+                    state[k[: -len("_out")]] = v
+        raise RuntimeError("no chunk returned hidden")
+
+    def probs(self, enc):
+        h = self.hidden(enc)
+        Q = len(enc["decide_idx"])
+        out, j = [], Q
+        with torch.no_grad():
+            for qi, oo in enumerate(enc["opt_idx"]):
+                z = (self.hk(h[j: j + len(oo)]) @ self.hq(h[qi])) / (self.hq.out_features ** 0.5)
+                out.append(torch.softmax(z / self.T, -1).numpy()); j += len(oo)
+        return out

reference/run.py

@@ -0,0 +1,38 @@
+"""Answer one System One request with a D1A Core ML bundle (this repository, downloaded), on the Neural Engine.
+
+    pip install coremltools numpy torch transformers safetensors "d1a @ git+https://github.com/jonpol01/d1a"
+    python reference/run.py <bundle dir> request.json
+
+request.json is a /v1/systemone request: {"state": "...", "questions": {"id": {"type": "noul" | "choice" | "score", ...}}}.
+The request must encode to 512 tokens or fewer. The first run compiles the chunks into coreml/.compiled/ (a few minutes);
+later runs load them. Prints each question's options with their probabilities."""
+import json
+import os
+import sys
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))   # workers.py and d1a_coreml.py sit next to this file
+
+
+def main():
+    from transformers import AutoTokenizer
+
+    from d1a.core.api import SystemOneRequest, to_record
+    from d1a.core.encoding import encode
+    from workers import WorkerPipeline
+
+    bundle, request = sys.argv[1], sys.argv[2]
+    req = SystemOneRequest.model_validate({"model": "d1a-latest", **json.load(open(request))})
+    rec, meta = to_record(req)
+    enc = encode(AutoTokenizer.from_pretrained(bundle), rec)
+    if len(enc["ids"]) > 512:
+        raise SystemExit(f"the request encodes to {len(enc['ids'])} tokens; this build takes 512 or fewer")
+    pipe = WorkerPipeline(bundle, os.path.join(bundle, "coreml"), 512)
+    try:
+        for m, p in zip(meta, pipe.probs(enc)):
+            print(m["id"], {k: round(float(v), 4) for k, v in zip(m["keys"], p)})
+    finally:
+        pipe.close()
+
+
+if __name__ == "__main__":
+    main()

reference/workers.py

@@ -0,0 +1,123 @@
+"""One process per Core ML chunk (workaround: several long-context chunks in one process crash the Neural Engine runtime
+on the second request on this Mac, while each chunk alone, or all on the GPU, runs fine). The parent runs the host side
+(embeddings, masks, pointer head) and passes tensors along the chain through pipes.
+  WorkerPipeline(export_dir, chunk_dir, L, units).probs(enc)"""
+import glob
+import os
+import multiprocessing as mp
+import re
+import shutil
+import sys
+
+import numpy as np
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))   # d1a_coreml.py sits next to this file
+
+
+BUCKETS = tuple(int(x) for x in os.environ.get("D1A_COREML_BUCKETS", "512").split(","))   # each length is its own ANE program; fewer keeps the ANE within its memory
+
+
+def compiled_chunks(chunk_dir):
+    """-> [(first layer, end layer, .mlmodelc, its .mlpackage)]. A released bundle's coreml/L512_<a>-<b>.mlpackage are
+    compiled once, each as a one-function package (L512), into coreml/.compiled/; a build tree's compiled/*.mlmodelc
+    (multifunction, one function per length) are used as they are."""
+    import coremltools as ct
+    out = []
+    for f in sorted(glob.glob(f"{chunk_dir}/*.mlmodelc")):
+        a, b = map(int, re.search(r"(\d+)-(\d+)\.mlmodelc$", f).groups())
+        out.append((a, b, f, f.replace("compiled/", "buckets/").replace(".mlmodelc", ".mlpackage")))
+    if out:
+        return out
+    os.makedirs(f"{chunk_dir}/.compiled", exist_ok=True)
+    for src in sorted(glob.glob(f"{chunk_dir}/L*_*.mlpackage")):
+        L, a, b = map(int, re.search(r"L(\d+)_(\d+)-(\d+)\.mlpackage$", src).groups())
+        dst = f"{chunk_dir}/.compiled/{a:02d}-{b:02d}.mlmodelc"
+        if not os.path.exists(dst):
+            d = ct.utils.MultiFunctionDescriptor(); d.add_function(src, src_function_name="main", target_function_name=f"L{L}")
+            d.default_function_name = f"L{L}"; tmp = f"{chunk_dir}/.compiled/{a:02d}-{b:02d}.mlpackage"
+            ct.utils.save_multifunction(d, tmp); shutil.move(ct.models.utils.compile_model(tmp), dst); shutil.rmtree(tmp)
+        out.append((a, b, dst, src))
+    return out
+
+
+def _worker(path, spec_path, units, conn):
+    import coremltools as ct
+    cu = getattr(ct.ComputeUnit, units)
+    # a precompiled .mlmodelc (compiled once by coremltools.models.utils.compile_model): loading each length function
+    # from it never recompiles, where MLModel(.mlpackage) compiles the whole package again for every function
+    models = {L: ct.models.CompiledMLModel(path, compute_units=cu, function_name=f"L{L}") for L in BUCKETS}
+    spec = ct.utils.load_spec(spec_path)
+    fns = {f.name: f for f in spec.description.functions}
+    fn = fns.get(f"L{BUCKETS[0]}") or spec.description   # a one-function package keeps its inputs in the description
+    conn.send(([i.name for i in fn.input], [o.name for o in fn.output]))
+    while True:
+        msg = conn.recv()
+        if msg is None:
+            break
+        L, feed = msg
+        import time as _t
+        t0 = _t.perf_counter()
+        out = models[L].predict(feed)
+        t1 = _t.perf_counter()
+        res = {k: np.array(v, dtype=np.float16 if v.dtype == np.float32 else v.dtype, copy=True) for k, v in out.items()}
+        res["_predict_ms"] = np.array((t1 - t0) * 1000)
+        conn.send(res)
+
+
+class WorkerPipeline:
+    def __init__(self, export_dir, chunk_dir, L=None, units="CPU_AND_NE"):
+        from d1a_coreml import Export, HostEmbed
+        self.ex = Export(export_dir)
+        self.host = HostEmbed(self.ex)
+        self.hq, self.hk = self.ex.head()
+        self.T, self.L = self.ex.d1a["temperature"], L
+        ctx = mp.get_context("spawn")
+        self.chunks = []
+        for a, b, f, spec_path in compiled_chunks(chunk_dir):
+            parent, child = ctx.Pipe()
+            p = ctx.Process(target=_worker, args=(f, spec_path, units, child), daemon=True)
+            p.start()
+            self.chunks.append((a, b, p, parent))
+            self.io = getattr(self, "io", []) + [parent.recv()]   # one at a time: concurrent first loads race in Core ML's compiler
+
+    def close(self):
+        for _, _, p, conn in self.chunks:
+            conn.send(None); p.join(5)
+
+    def hidden(self, enc):
+        from d1a_coreml import masks_for
+        self.last_predict_ms = []
+        n = len(enc["ids"])
+        L = next(b for b in BUCKETS if b >= n)   # the smallest compiled length that fits
+        ids = list(enc["ids"]) + [self.ex.d1a.get("pad_id", 0)] * (L - n)
+        emb, ple = self.host(ids)
+        mf, ms, pos = masks_for(enc, self.ex.tc.sliding_window, L=L)
+        picks = list(enc["decide_idx"]) + [o for oo in enc["opt_idx"] for o in oo]
+        pk = np.zeros(128, np.int32); pk[: len(picks)] = picks
+        common = {"pos": pos.numpy().astype(np.int32), "mask_full": mf.clamp(min=-1e4).numpy().astype(np.float16),
+                  "mask_slide": ms.clamp(min=-1e4).numpy().astype(np.float16)}
+        state = {"emb": emb.numpy().astype(np.float16), "ple_tok": ple.numpy().astype(np.float16), "picks": pk}
+        for (a, b, p, conn), (ins, outs) in zip(self.chunks, self.io):
+            conn.send((L, {k: (common[k] if k in common else state[k]) for k in ins}))
+            out = conn.recv()
+            self.last_predict_ms = getattr(self, "last_predict_ms", []) + [float(out.pop("_predict_ms"))]
+            if "hidden" in out:
+                return torch.tensor(out["hidden"][: len(picks)].astype(np.float32))
+            state["h_in"] = out["h"]
+            if "ple_out" in out:
+                state["ple"] = out["ple_out"]
+            for k, v in out.items():
+                if k.endswith("_out") and k != "ple_out":
+                    state[k[: -len("_out")]] = v
+        raise RuntimeError("no chunk returned hidden")
+
+    def probs(self, enc):
+        h = self.hidden(enc)
+        Q = len(enc["decide_idx"])
+        out, j = [], Q
+        with torch.no_grad():
+            for qi, oo in enumerate(enc["opt_idx"]):
+                z = (self.hk(h[j: j + len(oo)]) @ self.hq(h[qi])) / (self.hq.out_features ** 0.5)
+                out.append(torch.softmax(z / self.T, -1).numpy()); j += len(oo)
+        return out

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:12bac982b793c44b03d52a250a9f0d0b666813da566b910c24a6da0695fd11e6
+size 32170070

tokenizer_config.json

@@ -0,0 +1,55 @@
+{
+  "audio_token": "<|audio|>",
+  "backend": "tokenizers",
+  "boa_token": "<|audio>",
+  "boi_token": "<|image>",
+  "bos_token": "<bos>",
+  "eoa_token": "<audio|>",
+  "eoc_token": "<channel|>",
+  "eoi_token": "<image|>",
+  "eos_token": "<eos>",
+  "eot_token": "<turn|>",
+  "escape_token": "<|\"|>",
+  "etc_token": "<tool_call|>",
+  "etd_token": "<tool|>",
+  "etr_token": "<tool_response|>",
+  "extra_special_tokens": [
+    "<|video|>"
+  ],
+  "image_token": "<|image|>",
+  "is_local": false,
+  "local_files_only": false,
+  "mask_token": "<mask>",
+  "model_max_length": 1000000000000000019884624838656,
+  "model_specific_special_tokens": {
+    "audio_token": "<|audio|>",
+    "boa_token": "<|audio>",
+    "boi_token": "<|image>",
+    "eoa_token": "<audio|>",
+    "eoc_token": "<channel|>",
+    "eoi_token": "<image|>",
+    "eot_token": "<turn|>",
+    "escape_token": "<|\"|>",
+    "etc_token": "<tool_call|>",
+    "etd_token": "<tool|>",
+    "etr_token": "<tool_response|>",
+    "image_token": "<|image|>",
+    "soc_token": "<|channel>",
+    "sot_token": "<|turn>",
+    "stc_token": "<|tool_call>",
+    "std_token": "<|tool>",
+    "str_token": "<|tool_response>",
+    "think_token": "<|think|>"
+  },
+  "pad_token": "<pad>",
+  "padding_side": "left",
+  "processor_class": "Gemma4Processor",
+  "soc_token": "<|channel>",
+  "sot_token": "<|turn>",
+  "stc_token": "<|tool_call>",
+  "std_token": "<|tool>",
+  "str_token": "<|tool_response>",
+  "think_token": "<|think|>",
+  "tokenizer_class": "GemmaTokenizer",
+  "unk_token": "<unk>"
+}