zico: checkpoint, code, model card

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* 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/tokenizer.json filter=lfs diff=lfs merge=lfs -text

LICENSE-CODE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

LICENSE.md

@@ -0,0 +1,7 @@
+# Licences
+
+* **Model weights** (`backbone/`, `head.safetensors`): Creative Commons Attribution-NonCommercial 4.0 International
+  (CC-BY-NC-4.0), https://creativecommons.org/licenses/by-nc/4.0/legalcode. The base model is Apache-2.0, but part of the
+  training data is licensed for non-commercial or research use only; see the model card.
+* **Code** (`zico/`, `scripts/`, `tests/`, `examples/`): Apache License 2.0, see `LICENSE-CODE`.
+* **`vision/` and `tokenizer/`**: unmodified from `Qwen/Qwen3.5-4B-Base`, Apache License 2.0, see `NOTICE`.

NOTICE

@@ -0,0 +1,7 @@
+zico
+Copyright 2026 AmeenAhmed2
+
+This checkpoint is a fine-tune of Qwen/Qwen3.5-4B-Base (https://huggingface.co/Qwen/Qwen3.5-4B-Base),
+Copyright Alibaba Cloud, licensed under the Apache License, Version 2.0. The files in `vision/` (vision tower and image
+preprocessor configuration) and `tokenizer/` are redistributed from that model without modification; the language-model
+weights in `backbone/` were modified by fine-tuning.

README.md

@@ -0,0 +1,325 @@
+---
+license: cc-by-nc-4.0
+base_model: Qwen/Qwen3.5-4B-Base
+base_model_relation: finetune
+language:
+- en
+pipeline_tag: zero-shot-classification
+tags:
+- classification
+- zero-shot
+- calibration
+- decision-model
+- system-one
+- routing
+- guardrails
+- image-classification
+- qwen3.5
+- mlx
+- pytorch
+datasets:
+- google/boolq
+- fancyzhx/ag_news
+- Yelp/yelp_review_full
+- google-research-datasets/go_emotions
+- mteb/banking77
+- nyu-mll/glue
+- ehovy/race
+- google/civil_comments
+- allenai/openbookqa
+- tau/commonsense_qa
+- allenai/qasc
+- allenai/sciq
+- Rowan/hellaswag
+- allenai/winogrande
+- aps/super_glue
+- facebook/anli
+- stanfordnlp/snli
+- google-research-datasets/paws
+- tals/vitaminc
+- rajpurkar/squad_v2
+- clinc/clinc_oos
+- fancyzhx/dbpedia_14
+- community-datasets/yahoo_answers_topics
+- SetFit/20_newsgroups
+- cardiffnlp/tweet_eval
+- SetFit/sst5
+- SetFit/amazon_reviews_multi_en
+- deepset/prompt-injections
+- jackhhao/jailbreak-classification
+- lmsys/toxic-chat
+- microsoft/ms_marco
+- Team-ACE/ToolACE
+- openlifescienceai/medmcqa
+- derek-thomas/ScienceQA
+- bitext/Bitext-customer-support-llm-chatbot-training-dataset
+- benayas/snips
+- benayas/atis
+- zeroshot/twitter-financial-news-sentiment
+- Muennighoff/natural-instructions
+- chainyo/rvl-cdip
+- clane9/imagenet-100
+- uoft-cs/cifar100
+- tanganke/stanford_cars
+- bentrevett/caltech-ucsd-birds-200-2011
+- Multimodal-Fatima/FGVC_Aircraft_train
+- tanganke/dtd
+- tanganke/eurosat
+- tanganke/gtsrb
+- tanganke/resisc45
+- tanganke/sun397
+- ashraq/fashion-product-images-small
+---
+
+# zico
+
+**zico is a 4B decision model.** You give it a *state* (a message, a document, a JSON object, an image) and a set of
+*typed questions*. It answers every question with a **calibrated probability distribution over the options you
+listed** — in one forward pass, without generating any text.
+
+```jsonc
+// request                                                   // response
+{ "state": { "ticket": { "body": "Charged twice, refund me    { "answers": {
+             or I cancel." } },                                   "wants_refund": { "type": "noul", "noul": 0.89 },
+  "questions": {                                                  "team": { "type": "choice", "choice": "billing",
+    "wants_refund": { "type": "noul",                                       "probabilities": { "billing": 0.94, "technical": 0.02, … },
+        "instructions": "The customer asks for a refund." },                "confidence": 0.81 },
+    "team": { "type": "choice",                                   "frustration": { "type": "score", "score": 2.26,
+        "instructions": "Which team should handle this?",                   "probabilities": { "0": 0.04, "1": 0.08, "2": 0.54, … } } },
+        "criteria": { "billing": "Charges, refunds",            "usage": { "input_tokens": 240, "output_tokens": 0 } }
+                      "technical": "Bugs, outages", … } },
+    "frustration": { "type": "score",
+        "criteria": ["Calm", "Annoyed", "Frustrated", "Angry", "Furious"] } } }
+```
+
+It is meant for the many small decisions inside software — routing, triage, moderation, guardrails, tool selection,
+"is this answer grounded?", document and image sorting — where you want a fast, cheap, *typed* answer with a
+probability you can threshold, not a paragraph to parse.
+
+| Question type | You provide | You get |
+|---|---|---|
+| **Choice** | 2–255 named options, each with an optional description | the most likely option, the full distribution, a confidence |
+| **Score** | an ordered rubric of 2–10 level descriptions, lowest first | the expected level (e.g. `2.26`), the distribution over levels |
+| **Noul** | a statement that is true or false about the state | the probability that it is true |
+
+* **Type-safe by construction.** The output is a softmax over exactly the options in your request. An answer outside
+  your schema cannot happen, so there is nothing to parse, retry or validate.
+* **Calibrated.** Trained with strictly proper scoring rules; on held-out tasks its stated confidence tracks its accuracy
+  (expected calibration error ≈ 0.04–0.06, see below). "0.9" means roughly nine in ten.
+* **Many questions, one pass.** The state is encoded once and every question branches off it, so ten questions cost
+  little more than one. Answers are identical whether questions are asked together or alone.
+* **Text and images.** Images can sit anywhere inside the state, next to text fields.
+* **Fast.** ≈110 ms for 1–16 questions about a 500-token state on an H100 (64 questions: ≈260 ms); ≈300 ms for a
+  4-question request on an M4 Max laptop with the included MLX backend.
+
+## Quick start
+
+```bash
+pip install -U "huggingface_hub[cli]"
+hf download AmeenAhmed2/zico --local-dir zico && cd zico
+python -m venv .venv && source .venv/bin/activate
+pip install -e ".[serve]"            # Apple silicon: pip install -e ".[serve,mlx]"
+python scripts/infer.py . examples/request.json
+```
+
+The repository is the checkpoint *and* the code that runs it (`zico/`, a small Python package). The model cannot be
+loaded with `transformers.AutoModel` alone: the language backbone is a standard Qwen3.5 checkpoint, but the decision
+head and the input format are zico's own.
+
+```python
+import torch
+from zico import DecisionModel, Choice, Score, Noul
+
+model, renderer, _ = DecisionModel.load(".", device="cuda", dtype=torch.bfloat16)   # or "mps" / "cpu"
+model.autocast_dtype = torch.bfloat16
+
+answers = model.answer(
+    renderer,
+    state={"message": "hey can u move my 3pm with dr patel to next week, mornings are better"},
+    questions={
+        "intent": Choice("What does the patient want?", {"book": None, "reschedule": None, "cancel": None, "billing": None, "other": None}),
+        "time_preference": Noul("The patient states a preferred time of day."),
+        "urgency": Score("How urgent is this?", ["Whenever", "This week", "Today", "Emergency"]),
+    },
+)
+print(answers["answers"]["intent"])   # {'type': 'choice', 'choice': 'reschedule', 'probabilities': {...}, 'confidence': ...}
+```
+
+An image is a JSON object inside the state — `{"type": "image", "data": "<base64>"}`, optionally with
+`"detail": "low" | "high"` (64 to 1,024 image tokens; default ≤ 256):
+
+```python
+import base64
+photo = {"type": "image", "data": base64.b64encode(open("scan.png", "rb").read()).decode()}
+model.answer(renderer, {"scan": photo}, {"kind": Choice("What type of document is this?", {"invoice": None, "letter": None, "form": None, "resume": None})})
+```
+
+### HTTP server and test page
+
+```bash
+ZICO_CKPT=. scripts/serve.sh start        # also: stop | restart | status | log | key-path
+```
+
+This starts a small FastAPI server on `http://127.0.0.1:8787` (on Apple silicon it uses the MLX backend automatically):
+
+* `POST /v1/systemone` with `Authorization: Bearer <key>` takes the JSON request shown at the top (plus
+  `"model": "zico"`) and returns the response shown there. A key is generated on first start; `scripts/serve.sh key-path`
+  says where it is.
+* `GET /` is a **test page**: write a state, drop in an image, build questions, and watch the probability bars. "Live"
+  mode re-runs on every edit and shows how far each probability moved.
+
+The request and response shapes follow the *System One* wire format that TypeSafe AI documents publicly for its API, so
+client code written for that format can be pointed at this server (`ZICO_MODEL_ALIASES` adds any fixed model name such a
+client sends). See "Relation to other work" below.
+
+## How it works
+
+```
+row    =  ### State\n{state}\n\n   ### Choice\n{instructions}\n- billing: Charges…<OPT>\n- technical: Bugs…<OPT>\nAnswer:<ANS>
+          └──── shared prefix ────┘ └──────────────────────────── one question ──────────────────────────────────────┘
+
+logit_i = q(h[<ANS>]) · k(h[<OPT>_i]) / √d  +  u(h[<OPT>_i])          p = softmax over this question's own options
+```
+
+* **Backbone.** `Qwen/Qwen3.5-4B-Base` (a hybrid of Gated-DeltaNet and attention layers), fully fine-tuned, language-model
+  head removed. It is only ever used for *prefill*: there is no decoding loop.
+* **Slots, not tokens.** Each option ends in an `<OPT>` slot and each question in an `<ANS>` slot. Their input
+  embeddings come from a learned 4-row table, so the tokenizer and embedding matrix are untouched.
+* **Pointer head.** The hidden state at `<ANS>` is scored against the hidden state at every `<OPT>`; the softmax runs
+  over exactly those positions. Choice → arg-max key, Score → Σ i·pᵢ, Noul → p(true), confidence = 1 − H(p)/log K. The
+  head has 1.3M parameters.
+* **One state, many questions.** Training uses one `(state + question)` row per question. At inference the state is
+  prefilled once and its cache — attention KV *and* the recurrent / convolution state of the DeltaNet layers — is
+  forked across the questions, which run as one batch. Because the backbone is causal both compute the same function;
+  the test-suite asserts it.
+* **Images** enter as embeddings from Qwen3.5's own vision tower (frozen, shipped in `vision/`), placed where the image
+  sits in the state's JSON, with Qwen's 3-part rotary positions. They are part of the shared prefix, so an image is
+  encoded once for all questions.
+* **Loss.** Cross-entropy + 0.5 · Brier score, plus the ranked probability score for Score questions: all strictly proper
+  scoring rules, so the optimum is the true conditional distribution — accuracy and calibration are the same objective.
+  A per-type temperature fitted on validation data is stored in the checkpoint (1.09 / 1.10 / 1.23).
+
+## Training
+
+1. **Text pass** — 1 epoch over ≈40 public classification, NLI, reading-comprehension, multiple-choice, intent, moderation,
+   retrieval-relevance and tool-selection datasets plus ≈300 Super-NaturalInstructions tasks, all recast as typed
+   questions. To make the model *read the schema* rather than memorise label sets, records randomise option order,
+   option subsets, the presence of descriptions, instruction wording, whether the state is text or JSON, and add
+   "none of the above" cases. Whole datasets and whole instruction tasks are held out, and training text is de-duplicated
+   against every held-out set. ≈2.5 h on one H100 (bf16, 8-bit AdamW, ≈8k tokens/s).
+2. **Image pass** — 1 epoch over 90k images from 12 public image-classification sets (documents, objects, fine-grained
+   species / cars / aircraft, textures, satellite and aerial scenes, traffic signs, product attributes), vision tower
+   frozen. To stop the text skills drifting, 80k text rows were replayed with the *previous model's own probabilities*
+   as soft targets — not the gold labels. Replaying hard labels made the model markedly overconfident on tasks it had
+   never seen (ECE 0.045 → 0.131); anchoring to its own distributions avoided that (0.055) while keeping every image
+   gain. ≈85 min on one H100.
+
+Training and data-building code is in `scripts/` (`build_data.py`, `build_image_data.py`, `make_soft_targets.py`,
+`train.py`, `eval.py`).
+
+## Evaluation
+
+All numbers are this checkpoint, bf16, raw probabilities unless stated. "Held out" means the dataset (or the whole
+instruction task) never appeared in training, so the label set and the wording are new to the model. Sample sizes are
+small (150–500 per set; 5.5k for the unseen-task pool): differences under ≈3 points are noise.
+
+**Text, held-out (zero-shot)**
+
+| Set | Question | Accuracy | ECE |
+|---|---|---|---|
+| ARC-Easy | Choice, free-text options | 97.6 % | 0.031 |
+| IMDB sentiment | Noul | 94.0 % | 0.033 |
+| MASSIVE, 60 intents | Choice ×60 | 74.6 % | 0.090 |
+| 19 unseen Super-NaturalInstructions tasks | mixed | 75.2 % | 0.055 |
+| PubMedQA | Choice ×3 | 70.8 % | 0.088 |
+| dair-ai/emotion | Choice ×6 | 58.0 % | 0.200 |
+| AskUbuntu duplicate reranking (pick 1 of many) | Choice | 51.8 % | 0.161 |
+
+**Text, in-distribution test split:** 86.2 % over 37 sources (3,000 rows). Examples: Banking77 96.7 %, CLINC-150 with
+out-of-scope 95.3 %, AG News 93.0 %, MNLI 92.0 %, BoolQ 92.0 %, VitaminC fact-checking 92.9 %, SQuAD-v2 answerability
+88.0 %, jailbreak / prompt-injection detection 100 % (small sets), SST-5 67.3 %, MS MARCO passage selection 57.1 %.
+By question type over the whole test split (text and images): Choice 90.0 %, Noul 92.3 %, Score 68.5 % exact level.
+
+**Images, held-out sets (zero-shot: never trained on, full label set in every question)**
+
+| Set | Options per question | Accuracy | ECE |
+|---|---|---|---|
+| CIFAR-10 | 10 | 97.7 % | 0.024 |
+| MNIST digits | 10 | 95.3 % | 0.032 |
+| POPE object presence | Noul | 90.0 % | 0.073 |
+| Food-101 | 101 | 87.7 % | 0.043 |
+| Oxford-IIIT pets | 37 | 85.0 % | 0.064 |
+
+**Images, test splits of the training sets.** As in training, a question lists either the full label set or a random
+subset that contains the right answer, and about a quarter are "the image shows X" Noul questions — so these are easier
+than standard full-label benchmarks and are not comparable to published leaderboard numbers. 150–300 questions per set.
+
+| Set | Options per question | Accuracy | ECE |
+|---|---|---|---|
+| RVL-CDIP scanned document types | 6 or all 16 | 92.3 % | 0.036 |
+| Stanford Cars | 8–64 of 196 | 98.5 % | 0.023 |
+| EuroSAT land use | 5 or all 10 | 97.0 % | 0.023 |
+| SUN397 scenes | 8–64 of 397 | 96.6 % | 0.028 |
+| ImageNet-100 | 8–100 | 95.2 % | 0.017 |
+| RESISC45 aerial scenes | 8–45 | 94.5 % | 0.032 |
+| CIFAR-100 | 10–100 | 93.6 % | 0.026 |
+| DTD textures | 8–47 | 93.3 % | 0.046 |
+| CUB-200 birds | 8–64 of 200 | 92.8 % | 0.034 |
+| GTSRB traffic signs | 8–43 | 92.7 % | 0.035 |
+
+**Calibration.** Over all 9.9k held-out rows (text and images): ECE 0.055 raw, 0.041 with the stored temperatures. On the
+in-distribution test split: 0.011 raw. Full per-source results are in `eval.json`.
+
+**Consistency.** Forked (batched) answers equal single-question answers to ≈1e-6 in fp32 (6e-5 in bf16 on an H100). The MLX backend matches PyTorch:
+388 probabilities from real requests differ by 0.0003 on average, with the same top answer every time.
+
+## Limitations
+
+* **One forward pass, no reasoning.** Weak at counting, arithmetic, dates, multi-hop logic and double negatives. If a
+  question needs a chain of thought, this is the wrong tool.
+* **It classifies; it does not extract or generate.** No spans, no free text, no numbers beyond a 2–10 level rubric
+  (and on rubrics the exact level is right about two times in three; the expected value is the more useful output).
+* **Trained on rows up to ≈1k tokens.** The server accepts states up to 8k tokens, but accuracy and calibration on long
+  inputs are untested.
+* **Mostly English.** The backbone is multilingual; the fine-tuning data is almost entirely English.
+* **Fine-grained subjective labels are hard** (emotion 58 %, 5-star sentiment 67 %), and there the model is overconfident.
+  Calibration is a property of the distributions it was tested on; check it on your own data before relying on a threshold.
+* **Images:** single still images, classification-style questions. It is not an OCR engine and was not trained on
+  charts, screenshots, multi-page documents or video. Small text inside a photo may not be read.
+* **Safety-related questions** (toxicity, jailbreaks, prompt injection) were trained on small public sets. Treat the
+  output as one signal, not as a complete moderation or security system, and keep a human in the loop for decisions that
+  matter to people.
+* The training data carries the usual biases of public web, review and social-media corpora.
+
+## Licence and training data
+
+* **Weights: CC-BY-NC-4.0 (non-commercial).** The base model is Apache-2.0, but some of the training data is licensed
+  for non-commercial or research use only — notably MS MARCO and ToxicChat (text), and ImageNet, FGVC-Aircraft,
+  CUB-200, Stanford Cars, SUN397 and RVL-CDIP (images). The cautious reading is that a model trained on them should not
+  be used commercially. `scripts/build_data.py --exclude-nc` builds the text set without the two non-commercial text
+  sources if you want to train a variant.
+* **Code** (`zico/`, `scripts/`, `tests/`): Apache-2.0.
+* `vision/` is the unmodified vision tower of `Qwen/Qwen3.5-4B-Base` (Apache-2.0, © Alibaba Cloud), redistributed so the
+  checkpoint is self-contained. See `NOTICE`.
+* The datasets are listed in this card's metadata. They remain under their own licences; none of their content is
+  redistributed here.
+
+## Relation to other work
+
+* Built on **Qwen3.5-4B-Base** by the Qwen team; the Apple-silicon backend builds on **mlx-lm**.
+* The task format — a state plus typed Choice / Score / Noul questions answered with probabilities — and the HTTP wire
+  format follow the *System One* API that **TypeSafe AI** documents publicly. zico is an independent, from-scratch
+  project: it is not affiliated with or endorsed by TypeSafe AI, shares no code or weights with their models, and **no
+  output of their models was used to train or tune it**. The architecture described above is this project's own design
+  and says nothing about how theirs works.
+
+## Citation
+
+```bibtex
+@misc{zico2026,
+  title  = {zico: a calibrated, typed decision model in a single forward pass},
+  author = {AmeenAhmed2},
+  year   = {2026},
+  url    = {https://huggingface.co/AmeenAhmed2/zico}
+}
+```

backbone/config.json

@@ -0,0 +1,83 @@
+{
+  "architectures": [
+    "Qwen3_5TextModel"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "attn_output_gate": true,
+  "bos_token_id": null,
+  "dtype": "bfloat16",
+  "eos_token_id": 248044,
+  "full_attention_interval": 4,
+  "head_dim": 256,
+  "hidden_act": "silu",
+  "hidden_size": 2560,
+  "initializer_range": 0.02,
+  "intermediate_size": 9216,
+  "layer_types": [
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention"
+  ],
+  "linear_conv_kernel_dim": 4,
+  "linear_key_head_dim": 128,
+  "linear_num_key_heads": 16,
+  "linear_num_value_heads": 32,
+  "linear_value_head_dim": 128,
+  "mamba_ssm_dtype": "float32",
+  "max_position_embeddings": 262144,
+  "mlp_only_layers": [],
+  "model_type": "qwen3_5_text",
+  "mtp_num_hidden_layers": 1,
+  "mtp_use_dedicated_embeddings": false,
+  "num_attention_heads": 16,
+  "num_hidden_layers": 32,
+  "num_key_value_heads": 4,
+  "pad_token_id": null,
+  "partial_rotary_factor": 0.25,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "mrope_interleaved": true,
+    "mrope_section": [
+      11,
+      11,
+      10
+    ],
+    "partial_rotary_factor": 0.25,
+    "rope_theta": 10000000,
+    "rope_type": "default"
+  },
+  "tie_word_embeddings": true,
+  "transformers_version": "5.17.0",
+  "use_cache": true,
+  "vocab_size": 248320
+}

backbone/model.safetensors

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

eval.json

@@ -0,0 +1,1504 @@
+{
+ "meta": {
+  "base_model": "Qwen/Qwen3.5-4B-Base",
+  "d_head": 256,
+  "lora": false,
+  "renderer": {
+   "max_state_tokens": 768,
+   "max_instruction_tokens": 256,
+   "max_option_tokens": 64
+  },
+  "step": 734,
+  "val_nll": 0.27321311831474304
+ },
+ "temperatures": {
+  "choice": 1.0929638744012644,
+  "score": 1.1011288758276154,
+  "noul": 1.228405177254357
+ },
+ "test": {
+  "raw": {
+   "overall": {
+    "n": 5070,
+    "acc": 0.8964496850967407,
+    "nll": 0.276126503944397,
+    "brier": 0.1475353091955185,
+    "ece": 0.010809999694056847,
+    "mean_conf": 0.9048414826393127,
+    "score_mae": 0.41506025195121765
+   },
+   "by_type": {
+    "choice": {
+     "n": 3041,
+     "acc": 0.9003617167472839,
+     "nll": 0.2871852517127991,
+     "brier": 0.1437402367591858,
+     "ece": 0.013187985173783867,
+     "mean_conf": 0.9091227054595947
+    },
+    "score": {
+     "n": 273,
+     "acc": 0.6849817037582397,
+     "nll": 0.7681964039802551,
+     "brier": 0.4400661289691925,
+     "ece": 0.047672064850177796,
+     "mean_conf": 0.6837503910064697,
+     "score_mae": 0.41506025195121765
+    },
+    "noul": {
+     "n": 1756,
+     "acc": 0.9225512742996216,
+     "nll": 0.18047469854354858,
+     "brier": 0.10862862318754196,
+     "ece": 0.013375402451816704,
+     "mean_conf": 0.9318000674247742
+    }
+   },
+   "by_source": {
+    "ag_news": {
+     "n": 86,
+     "acc": 0.930232584476471,
+     "nll": 0.23873434960842133,
+     "brier": 0.10834771394729614,
+     "ece": 0.056332710342662184,
+     "mean_conf": 0.9181896448135376
+    },
+    "amazon_stars": {
+     "n": 86,
+     "acc": 0.5930232405662537,
+     "nll": 0.8701136112213135,
+     "brier": 0.5058151483535767,
+     "ece": 0.10280728207642115,
+     "mean_conf": 0.6330943703651428,
+     "score_mae": 0.4689413011074066
+    },
+    "anli": {
+     "n": 79,
+     "acc": 0.746835470199585,
+     "nll": 0.6076061129570007,
+     "brier": 0.3763708174228668,
+     "ece": 0.15590721699211704,
+     "mean_conf": 0.8249932527542114
+    },
+    "banking77": {
+     "n": 91,
+     "acc": 0.9670329689979553,
+     "nll": 0.109474316239357,
+     "brier": 0.05170481652021408,
+     "ece": 0.04274910372363963,
+     "mean_conf": 0.9463766813278198
+    },
+    "boolq": {
+     "n": 75,
+     "acc": 0.9200000166893005,
+     "nll": 0.22674815356731415,
+     "brier": 0.14536462724208832,
+     "ece": 0.07328386494371486,
+     "mean_conf": 0.908048689365387
+    },
+    "cifar100": {
+     "n": 297,
+     "acc": 0.936026930809021,
+     "nll": 0.19756226241588593,
+     "brier": 0.096889927983284,
+     "ece": 0.02602774308840307,
+     "mean_conf": 0.9325385689735413
+    },
+    "civil_comments": {
+     "n": 154,
+     "acc": 0.9545454382896423,
+     "nll": 0.1635725051164627,
+     "brier": 0.08106027543544769,
+     "ece": 0.08850163645923792,
+     "mean_conf": 0.8801278471946716
+    },
+    "clinc_oos": {
+     "n": 85,
+     "acc": 0.9529411792755127,
+     "nll": 0.10286867618560791,
+     "brier": 0.057649165391922,
+     "ece": 0.028260532641420116,
+     "mean_conf": 0.9765596389770508
+    },
+    "cola": {
+     "n": 89,
+     "acc": 0.8202247023582458,
+     "nll": 0.3538511097431183,
+     "brier": 0.23637528717517853,
+     "ece": 0.07856890498333513,
+     "mean_conf": 0.8428974747657776
+    },
+    "commonsense_qa": {
+     "n": 79,
+     "acc": 0.797468364238739,
+     "nll": 0.47439879179000854,
+     "brier": 0.2571961283683777,
+     "ece": 0.11401724499686991,
+     "mean_conf": 0.8429151177406311
+    },
+    "copa": {
+     "n": 19,
+     "acc": 0.9473684430122375,
+     "nll": 0.08234653621912003,
+     "brier": 0.056282687932252884,
+     "ece": 0.053942021972607,
+     "mean_conf": 0.9678375720977783
+    },
+    "cub_birds": {
+     "n": 209,
+     "acc": 0.9282296895980835,
+     "nll": 0.205092191696167,
+     "brier": 0.09457369148731232,
+     "ece": 0.0337111855618914,
+     "mean_conf": 0.9249444007873535
+    },
+    "dbpedia": {
+     "n": 83,
+     "acc": 0.9879518151283264,
+     "nll": 0.0267021507024765,
+     "brier": 0.01387640181928873,
+     "ece": 0.01773379953028531,
+     "mean_conf": 0.9858810305595398
+    },
+    "dtd": {
+     "n": 193,
+     "acc": 0.9326424598693848,
+     "nll": 0.1894713044166565,
+     "brier": 0.10953276604413986,
+     "ece": 0.0462280399812367,
+     "mean_conf": 0.9360920786857605
+    },
+    "eurosat": {
+     "n": 200,
+     "acc": 0.9700000286102295,
+     "nll": 0.09461035579442978,
+     "brier": 0.04988160729408264,
+     "ece": 0.022912378822553925,
+     "mean_conf": 0.9720562696456909
+    },
+    "go_emotions": {
+     "n": 156,
+     "acc": 0.8461538553237915,
+     "nll": 0.4075973331928253,
+     "brier": 0.2239874303340912,
+     "ece": 0.052582084089755954,
+     "mean_conf": 0.8886739015579224
+    },
+    "gtsrb": {
+     "n": 150,
+     "acc": 0.9266666769981384,
+     "nll": 0.2583926022052765,
+     "brier": 0.10076691210269928,
+     "ece": 0.0349075827405604,
+     "mean_conf": 0.9396733641624451
+    },
+    "hellaswag": {
+     "n": 71,
+     "acc": 0.9718309640884399,
+     "nll": 0.07258392125368118,
+     "brier": 0.03423016890883446,
+     "ece": 0.037721984350516546,
+     "mean_conf": 0.9680467844009399
+    },
+    "imagenet100": {
+     "n": 209,
+     "acc": 0.9521530866622925,
+     "nll": 0.1411520093679428,
+     "brier": 0.06291362643241882,
+     "ece": 0.01678104336341813,
+     "mean_conf": 0.9670758247375488
+    },
+    "jailbreak": {
+     "n": 43,
+     "acc": 1.0,
+     "nll": 0.003266148502007127,
+     "brier": 0.0003119981265626848,
+     "ece": 0.003183770989621637,
+     "mean_conf": 0.9968161582946777
+    },
+    "mnli": {
+     "n": 88,
+     "acc": 0.9204545617103577,
+     "nll": 0.17574532330036163,
+     "brier": 0.10608046501874924,
+     "ece": 0.039473747201574705,
+     "mean_conf": 0.9106850028038025
+    },
+    "ms_marco": {
+     "n": 77,
+     "acc": 0.5714285969734192,
+     "nll": 1.1672861576080322,
+     "brier": 0.5488265156745911,
+     "ece": 0.1069421570546632,
+     "mean_conf": 0.5573300123214722
+    },
+    "newsgroups": {
+     "n": 72,
+     "acc": 0.8194444179534912,
+     "nll": 0.656489372253418,
+     "brier": 0.24557192623615265,
+     "ece": 0.10580966421453894,
+     "mean_conf": 0.8494588732719421
+    },
+    "openbookqa": {
+     "n": 73,
+     "acc": 0.9452054500579834,
+     "nll": 0.2053627371788025,
+     "brier": 0.10051807016134262,
+     "ece": 0.07717989225561772,
+     "mean_conf": 0.8888569474220276
+    },
+    "paws": {
+     "n": 75,
+     "acc": 0.8799999952316284,
+     "nll": 0.252999871969223,
+     "brier": 0.16616970300674438,
+     "ece": 0.07703955284745612,
+     "mean_conf": 0.957039475440979
+    },
+    "prompt_injection": {
+     "n": 19,
+     "acc": 1.0,
+     "nll": 0.07026468217372894,
+     "brier": 0.02881462872028351,
+     "ece": 0.06134651247871403,
+     "mean_conf": 0.9386534690856934
+    },
+    "qasc": {
+     "n": 78,
+     "acc": 0.8974359035491943,
+     "nll": 0.22243788838386536,
+     "brier": 0.11997074633836746,
+     "ece": 0.058650574809634426,
+     "mean_conf": 0.9030769467353821
+    },
+    "qnli": {
+     "n": 79,
+     "acc": 0.9620253443717957,
+     "nll": 0.12711185216903687,
+     "brier": 0.06058496609330177,
+     "ece": 0.02656769830988215,
+     "mean_conf": 0.9400109648704529
+    },
+    "qqp": {
+     "n": 67,
+     "acc": 0.8507462739944458,
+     "nll": 0.32717540860176086,
+     "brier": 0.20558716356754303,
+     "ece": 0.05535818888270405,
+     "mean_conf": 0.8827617168426514
+    },
+    "race": {
+     "n": 71,
+     "acc": 0.8169013857841492,
+     "nll": 0.5218772292137146,
+     "brier": 0.26671260595321655,
+     "ece": 0.08260945939768694,
+     "mean_conf": 0.868141770362854
+    },
+    "resisc45": {
+     "n": 201,
+     "acc": 0.9452736377716064,
+     "nll": 0.17759159207344055,
+     "brier": 0.07590000331401825,
+     "ece": 0.03162269839005066,
+     "mean_conf": 0.9709147810935974
+    },
+    "rte": {
+     "n": 42,
+     "acc": 0.9047619104385376,
+     "nll": 0.21782515943050385,
+     "brier": 0.12485634535551071,
+     "ece": 0.05018191372233538,
+     "mean_conf": 0.9342974424362183
+    },
+    "rvl_cdip": {
+     "n": 207,
+     "acc": 0.9227052927017212,
+     "nll": 0.24853183329105377,
+     "brier": 0.1186482235789299,
+     "ece": 0.035717121188601514,
+     "mean_conf": 0.9394121170043945
+    },
+    "sciq": {
+     "n": 70,
+     "acc": 0.9571428298950195,
+     "nll": 0.08418832719326019,
+     "brier": 0.05511568486690521,
+     "ece": 0.05332319163331367,
+     "mean_conf": 0.9662564992904663
+    },
+    "snli": {
+     "n": 78,
+     "acc": 0.9615384340286255,
+     "nll": 0.13747791945934296,
+     "brier": 0.07794130593538284,
+     "ece": 0.06147045345163815,
+     "mean_conf": 0.9255906343460083
+    },
+    "squad_v2": {
+     "n": 83,
+     "acc": 0.8795180916786194,
+     "nll": 0.3885398805141449,
+     "brier": 0.18468858301639557,
+     "ece": 0.08604057614296412,
+     "mean_conf": 0.9535071849822998
+    },
+    "sst5": {
+     "n": 101,
+     "acc": 0.6732673048973083,
+     "nll": 0.8455938696861267,
+     "brier": 0.4966614544391632,
+     "ece": 0.10307709063944676,
+     "mean_conf": 0.6207720041275024,
+     "score_mae": 0.45541465282440186
+    },
+    "stanford_cars": {
+     "n": 199,
+     "acc": 0.9849246144294739,
+     "nll": 0.06568868458271027,
+     "brier": 0.033190734684467316,
+     "ece": 0.02284979613965564,
+     "mean_conf": 0.9779478907585144
+    },
+    "sun397": {
+     "n": 205,
+     "acc": 0.9658536314964294,
+     "nll": 0.13169661164283752,
+     "brier": 0.059237346053123474,
+     "ece": 0.027937838926008762,
+     "mean_conf": 0.9778258800506592
+    },
+    "toxic_chat": {
+     "n": 103,
+     "acc": 0.9805825352668762,
+     "nll": 0.04486699774861336,
+     "brier": 0.02695567160844803,
+     "ece": 0.011975512422466195,
+     "mean_conf": 0.9775879383087158
+    },
+    "tweet_irony": {
+     "n": 73,
+     "acc": 0.7397260069847107,
+     "nll": 0.47673916816711426,
+     "brier": 0.32497861981391907,
+     "ece": 0.10002183849680613,
+     "mean_conf": 0.8234904408454895
+    },
+    "tweet_offensive": {
+     "n": 95,
+     "acc": 0.800000011920929,
+     "nll": 0.42309972643852234,
+     "brier": 0.2671454846858978,
+     "ece": 0.11630791231693371,
+     "mean_conf": 0.8510119915008545
+    },
+    "tweet_sentiment": {
+     "n": 96,
+     "acc": 0.7604166865348816,
+     "nll": 0.6009823679924011,
+     "brier": 0.3709958493709564,
+     "ece": 0.08339798031360968,
+     "mean_conf": 0.7406583428382874
+    },
+    "vitaminc": {
+     "n": 84,
+     "acc": 0.9285714030265808,
+     "nll": 0.20010635256767273,
+     "brier": 0.11288775503635406,
+     "ece": 0.055092808689364614,
+     "mean_conf": 0.8885759711265564
+    },
+    "winogrande": {
+     "n": 82,
+     "acc": 0.8292682766914368,
+     "nll": 0.3929581940174103,
+     "brier": 0.2536102533340454,
+     "ece": 0.09293118971498115,
+     "mean_conf": 0.9037488698959351
+    },
+    "yahoo_topics": {
+     "n": 71,
+     "acc": 0.7746478915214539,
+     "nll": 0.6088139414787292,
+     "brier": 0.28299495577812195,
+     "ece": 0.06163965007121763,
+     "mean_conf": 0.7804079055786133
+    },
+    "yelp": {
+     "n": 127,
+     "acc": 0.834645688533783,
+     "nll": 0.4275256395339966,
+     "brier": 0.2361738681793213,
+     "ece": 0.047331056411305306,
+     "mean_conf": 0.858354926109314,
+     "score_mae": 0.31378617882728577
+    }
+   }
+  },
+  "calibrated": {
+   "overall": {
+    "n": 5070,
+    "acc": 0.8964496850967407,
+    "nll": 0.2756565809249878,
+    "brier": 0.14737455546855927,
+    "ece": 0.0079650055739778,
+    "mean_conf": 0.8922349810600281,
+    "score_mae": 0.42143115401268005
+   },
+   "by_type": {
+    "choice": {
+     "n": 3041,
+     "acc": 0.9003617167472839,
+     "nll": 0.28574398159980774,
+     "brier": 0.14335514605045319,
+     "ece": 0.007999743385107894,
+     "mean_conf": 0.8997042179107666
+    },
+    "score": {
+     "n": 273,
+     "acc": 0.6849817037582397,
+     "nll": 0.7671821713447571,
+     "brier": 0.439628005027771,
+     "ece": 0.046648137280912105,
+     "mean_conf": 0.6634480953216553,
+     "score_mae": 0.42143115401268005
+    },
+    "noul": {
+     "n": 1756,
+     "acc": 0.9225512742996216,
+     "nll": 0.18177150189876556,
+     "brier": 0.10889948904514313,
+     "ece": 0.019957571878301106,
+     "mean_conf": 0.9148688316345215
+    }
+   },
+   "by_source": {
+    "ag_news": {
+     "n": 86,
+     "acc": 0.930232584476471,
+     "nll": 0.2372349351644516,
+     "brier": 0.10982950031757355,
+     "ece": 0.06062877569098146,
+     "mean_conf": 0.909740686416626
+    },
+    "amazon_stars": {
+     "n": 86,
+     "acc": 0.5930232405662537,
+     "nll": 0.8707610964775085,
+     "brier": 0.5034118294715881,
+     "ece": 0.07147133660142413,
+     "mean_conf": 0.6115179657936096,
+     "score_mae": 0.4738217294216156
+    },
+    "anli": {
+     "n": 79,
+     "acc": 0.746835470199585,
+     "nll": 0.5904105305671692,
+     "brier": 0.3706202208995819,
+     "ece": 0.12736005320452623,
+     "mean_conf": 0.8066393733024597
+    },
+    "banking77": {
+     "n": 91,
+     "acc": 0.9670329689979553,
+     "nll": 0.11517834663391113,
+     "brier": 0.05284213647246361,
+     "ece": 0.04865624303995164,
+     "mean_conf": 0.9390164613723755
+    },
+    "boolq": {
+     "n": 75,
+     "acc": 0.9200000166893005,
+     "nll": 0.230869323015213,
+     "brier": 0.14464128017425537,
+     "ece": 0.05709565622368751,
+     "mean_conf": 0.8859748244285583
+    },
+    "cifar100": {
+     "n": 297,
+     "acc": 0.936026930809021,
+     "nll": 0.19978056848049164,
+     "brier": 0.09738589078187943,
+     "ece": 0.03686069778541466,
+     "mean_conf": 0.9233465194702148
+    },
+    "civil_comments": {
+     "n": 154,
+     "acc": 0.9545454382896423,
+     "nll": 0.1945398449897766,
+     "brier": 0.09554325044155121,
+     "ece": 0.11094526093531587,
+     "mean_conf": 0.8504656553268433
+    },
+    "clinc_oos": {
+     "n": 85,
+     "acc": 0.9529411792755127,
+     "nll": 0.10213149338960648,
+     "brier": 0.05607043579220772,
+     "ece": 0.03434847114305911,
+     "mean_conf": 0.9722728133201599
+    },
+    "cola": {
+     "n": 89,
+     "acc": 0.8202247023582458,
+     "nll": 0.3592163622379303,
+     "brier": 0.2345142364501953,
+     "ece": 0.07689192332149841,
+     "mean_conf": 0.8134418725967407
+    },
+    "commonsense_qa": {
+     "n": 79,
+     "acc": 0.797468364238739,
+     "nll": 0.47140878438949585,
+     "brier": 0.25501152873039246,
+     "ece": 0.08131545514081495,
+     "mean_conf": 0.8276560306549072
+    },
+    "copa": {
+     "n": 19,
+     "acc": 0.9473684430122375,
+     "nll": 0.08365084230899811,
+     "brier": 0.05513031780719757,
+     "ece": 0.05683338684464534,
+     "mean_conf": 0.9632882475852966
+    },
+    "cub_birds": {
+     "n": 209,
+     "acc": 0.9282296895980835,
+     "nll": 0.20787712931632996,
+     "brier": 0.09552589058876038,
+     "ece": 0.03924435088578615,
+     "mean_conf": 0.9157379269599915
+    },
+    "dbpedia": {
+     "n": 83,
+     "acc": 0.9879518151283264,
+     "nll": 0.02914765290915966,
+     "brier": 0.013987519778311253,
+     "ece": 0.020033360317049276,
+     "mean_conf": 0.9829169511795044
+    },
+    "dtd": {
+     "n": 193,
+     "acc": 0.9326424598693848,
+     "nll": 0.19226685166358948,
+     "brier": 0.10932108014822006,
+     "ece": 0.04905586869892781,
+     "mean_conf": 0.926359236240387
+    },
+    "eurosat": {
+     "n": 200,
+     "acc": 0.9700000286102295,
+     "nll": 0.09493371099233627,
+     "brier": 0.04979003965854645,
+     "ece": 0.019272321772114326,
+     "mean_conf": 0.964755117893219
+    },
+    "go_emotions": {
+     "n": 156,
+     "acc": 0.8461538553237915,
+     "nll": 0.3998943567276001,
+     "brier": 0.2204635888338089,
+     "ece": 0.06369368834835576,
+     "mean_conf": 0.8688168525695801
+    },
+    "gtsrb": {
+     "n": 150,
+     "acc": 0.9266666769981384,
+     "nll": 0.25059908628463745,
+     "brier": 0.09959184378385544,
+     "ece": 0.04895997661024053,
+     "mean_conf": 0.9357502460479736
+    },
+    "hellaswag": {
+     "n": 71,
+     "acc": 0.9718309640884399,
+     "nll": 0.07610221207141876,
+     "brier": 0.034769345074892044,
+     "ece": 0.0420044008320708,
+     "mean_conf": 0.9625624418258667
+    },
+    "imagenet100": {
+     "n": 209,
+     "acc": 0.9521530866622925,
+     "nll": 0.140834778547287,
+     "brier": 0.06210394203662872,
+     "ece": 0.02685701278059019,
+     "mean_conf": 0.9609060883522034
+    },
+    "jailbreak": {
+     "n": 43,
+     "acc": 1.0,
+     "nll": 0.007107510231435299,
+     "brier": 0.0007506122929044068,
+     "ece": 0.006904559023125811,
+     "mean_conf": 0.9930953979492188
+    },
+    "mnli": {
+     "n": 88,
+     "acc": 0.9204545617103577,
+     "nll": 0.18531149625778198,
+     "brier": 0.10737426578998566,
+     "ece": 0.06684734037705664,
+     "mean_conf": 0.8966826796531677
+    },
+    "ms_marco": {
+     "n": 77,
+     "acc": 0.5714285969734192,
+     "nll": 1.1738694906234741,
+     "brier": 0.5511745810508728,
+     "ece": 0.1849330931448978,
+     "mean_conf": 0.5303053855895996
+    },
+    "newsgroups": {
+     "n": 72,
+     "acc": 0.8194444179534912,
+     "nll": 0.6413953304290771,
+     "brier": 0.2439313530921936,
+     "ece": 0.0672152123840676,
+     "mean_conf": 0.8347481489181519
+    },
+    "openbookqa": {
+     "n": 73,
+     "acc": 0.9452054500579834,
+     "nll": 0.21645799279212952,
+     "brier": 0.10464119911193848,
+     "ece": 0.07754900295525902,
+     "mean_conf": 0.8753600716590881
+    },
+    "paws": {
+     "n": 75,
+     "acc": 0.8799999952316284,
+     "nll": 0.23008832335472107,
+     "brier": 0.15555186569690704,
+     "ece": 0.07505809075017267,
+     "mean_conf": 0.9440792798995972
+    },
+    "prompt_injection": {
+     "n": 19,
+     "acc": 1.0,
+     "nll": 0.0893387570977211,
+     "brier": 0.0387185737490654,
+     "ece": 0.07711471421890259,
+     "mean_conf": 0.9228852987289429
+    },
+    "qasc": {
+     "n": 78,
+     "acc": 0.8974359035491943,
+     "nll": 0.2294393628835678,
+     "brier": 0.12008954584598541,
+     "ece": 0.07760370314998621,
+     "mean_conf": 0.8928042650222778
+    },
+    "qnli": {
+     "n": 79,
+     "acc": 0.9620253443717957,
+     "nll": 0.13672247529029846,
+     "brier": 0.0658198893070221,
+     "ece": 0.04174679779133661,
+     "mean_conf": 0.9221668243408203
+    },
+    "qqp": {
+     "n": 67,
+     "acc": 0.8507462739944458,
+     "nll": 0.3204824924468994,
+     "brier": 0.20365533232688904,
+     "ece": 0.05547176195417913,
+     "mean_conf": 0.8596541285514832
+    },
+    "race": {
+     "n": 71,
+     "acc": 0.8169013857841492,
+     "nll": 0.508931040763855,
+     "brier": 0.2651640772819519,
+     "ece": 0.09953184393917758,
+     "mean_conf": 0.855144202709198
+    },
+    "resisc45": {
+     "n": 201,
+     "acc": 0.9452736377716064,
+     "nll": 0.1719132363796234,
+     "brier": 0.07461114972829819,
+     "ece": 0.03276641192363214,
+     "mean_conf": 0.9654509425163269
+    },
+    "rte": {
+     "n": 42,
+     "acc": 0.9047619104385376,
+     "nll": 0.2149493396282196,
+     "brier": 0.12409208714962006,
+     "ece": 0.09782490985591286,
+     "mean_conf": 0.9137647747993469
+    },
+    "rvl_cdip": {
+     "n": 207,
+     "acc": 0.9227052927017212,
+     "nll": 0.24763381481170654,
+     "brier": 0.11820948868989944,
+     "ece": 0.024818566785942367,
+     "mean_conf": 0.9276553392410278
+    },
+    "sciq": {
+     "n": 70,
+     "acc": 0.9571428298950195,
+     "nll": 0.08641297370195389,
+     "brier": 0.0542028583586216,
+     "ece": 0.05665828745125506,
+     "mean_conf": 0.9611032605171204
+    },
+    "snli": {
+     "n": 78,
+     "acc": 0.9615384340286255,
+     "nll": 0.14562663435935974,
+     "brier": 0.07991129904985428,
+     "ece": 0.05743888654096996,
+     "mean_conf": 0.9150496125221252
+    },
+    "squad_v2": {
+     "n": 83,
+     "acc": 0.8795180916786194,
+     "nll": 0.3431330621242523,
+     "brier": 0.18107561767101288,
+     "ece": 0.08313956865378169,
+     "mean_conf": 0.9385043978691101
+    },
+    "sst5": {
+     "n": 101,
+     "acc": 0.6732673048973083,
+     "nll": 0.8503657579421997,
+     "brier": 0.49845677614212036,
+     "ece": 0.0914942533168388,
+     "mean_conf": 0.5979520082473755,
+     "score_mae": 0.46103909611701965
+    },
+    "stanford_cars": {
+     "n": 199,
+     "acc": 0.9849246144294739,
+     "nll": 0.0664161965250969,
+     "brier": 0.03385894373059273,
+     "ece": 0.02054634083136521,
+     "mean_conf": 0.9744680523872375
+    },
+    "sun397": {
+     "n": 205,
+     "acc": 0.9658536314964294,
+     "nll": 0.12713490426540375,
+     "brier": 0.05907195806503296,
+     "ece": 0.03913777098139376,
+     "mean_conf": 0.973029613494873
+    },
+    "toxic_chat": {
+     "n": 103,
+     "acc": 0.9805825352668762,
+     "nll": 0.04893926531076431,
+     "brier": 0.027793018147349358,
+     "ece": 0.0175806231729555,
+     "mean_conf": 0.971018373966217
+    },
+    "tweet_irony": {
+     "n": 73,
+     "acc": 0.7397260069847107,
+     "nll": 0.46547093987464905,
+     "brier": 0.3147864043712616,
+     "ece": 0.1386445264786409,
+     "mean_conf": 0.7930445075035095
+    },
+    "tweet_offensive": {
+     "n": 95,
+     "acc": 0.800000011920929,
+     "nll": 0.4134097099304199,
+     "brier": 0.263822078704834,
+     "ece": 0.12641781317580403,
+     "mean_conf": 0.8215466737747192
+    },
+    "tweet_sentiment": {
+     "n": 96,
+     "acc": 0.7604166865348816,
+     "nll": 0.6009461283683777,
+     "brier": 0.37088480591773987,
+     "ece": 0.09440179414827354,
+     "mean_conf": 0.7233335375785828
+    },
+    "vitaminc": {
+     "n": 84,
+     "acc": 0.9285714030265808,
+     "nll": 0.21254083514213562,
+     "brier": 0.11658700555562973,
+     "ece": 0.060456716703195634,
+     "mean_conf": 0.87336665391922
+    },
+    "winogrande": {
+     "n": 82,
+     "acc": 0.8292682766914368,
+     "nll": 0.3811376690864563,
+     "brier": 0.24891741573810577,
+     "ece": 0.07569023195409919,
+     "mean_conf": 0.8934732675552368
+    },
+    "yahoo_topics": {
+     "n": 71,
+     "acc": 0.7746478915214539,
+     "nll": 0.6164750456809998,
+     "brier": 0.2843913733959198,
+     "ece": 0.10321198388129582,
+     "mean_conf": 0.7571716904640198
+    },
+    "yelp": {
+     "n": 127,
+     "acc": 0.834645688533783,
+     "nll": 0.4197368919849396,
+     "brier": 0.23375587165355682,
+     "ece": 0.036530637687290035,
+     "mean_conf": 0.8449493646621704,
+     "score_mae": 0.32252421975135803
+    }
+   }
+  },
+  "reliability_raw": [
+   {
+    "lo": 0.133,
+    "hi": 0.2,
+    "conf": 0.17327187955379486,
+    "acc": 0.0,
+    "frac": 0.0007889546104706824
+   },
+   {
+    "lo": 0.2,
+    "hi": 0.267,
+    "conf": 0.23755492269992828,
+    "acc": 0.2222222238779068,
+    "frac": 0.001775147975422442
+   },
+   {
+    "lo": 0.267,
+    "hi": 0.333,
+    "conf": 0.3051908314228058,
+    "acc": 0.2800000011920929,
+    "frac": 0.004930966533720493
+   },
+   {
+    "lo": 0.333,
+    "hi": 0.4,
+    "conf": 0.36935192346572876,
+    "acc": 0.26923078298568726,
+    "frac": 0.0051282052882015705
+   },
+   {
+    "lo": 0.4,
+    "hi": 0.467,
+    "conf": 0.43677520751953125,
+    "acc": 0.42307692766189575,
+    "frac": 0.010256410576403141
+   },
+   {
+    "lo": 0.467,
+    "hi": 0.533,
+    "conf": 0.5065982937812805,
+    "acc": 0.5289855003356934,
+    "frac": 0.027218934148550034
+   },
+   {
+    "lo": 0.533,
+    "hi": 0.6,
+    "conf": 0.5669992566108704,
+    "acc": 0.5779221057891846,
+    "frac": 0.03037475422024727
+   },
+   {
+    "lo": 0.6,
+    "hi": 0.667,
+    "conf": 0.63744056224823,
+    "acc": 0.6171428561210632,
+    "frac": 0.03451676666736603
+   },
+   {
+    "lo": 0.667,
+    "hi": 0.733,
+    "conf": 0.7005413770675659,
+    "acc": 0.6280487775802612,
+    "frac": 0.03234713897109032
+   },
+   {
+    "lo": 0.733,
+    "hi": 0.8,
+    "conf": 0.767250657081604,
+    "acc": 0.7511312365531921,
+    "frac": 0.043589744716882706
+   },
+   {
+    "lo": 0.8,
+    "hi": 0.867,
+    "conf": 0.834570586681366,
+    "acc": 0.80694979429245,
+    "frac": 0.051084812730550766
+   },
+   {
+    "lo": 0.867,
+    "hi": 0.933,
+    "conf": 0.9034878015518188,
+    "acc": 0.9073033928871155,
+    "frac": 0.07021696120500565
+   },
+   {
+    "lo": 0.933,
+    "hi": 1.0,
+    "conf": 0.9904627203941345,
+    "acc": 0.9853742718696594,
+    "frac": 0.687771201133728
+   }
+  ],
+  "reliability_calibrated": [
+   {
+    "lo": 0.133,
+    "hi": 0.2,
+    "conf": 0.17058970034122467,
+    "acc": 0.0,
+    "frac": 0.0009861933067440987
+   },
+   {
+    "lo": 0.2,
+    "hi": 0.267,
+    "conf": 0.24141208827495575,
+    "acc": 0.23076923191547394,
+    "frac": 0.0025641026441007853
+   },
+   {
+    "lo": 0.267,
+    "hi": 0.333,
+    "conf": 0.30261722207069397,
+    "acc": 0.3076923191547394,
+    "frac": 0.0051282052882015705
+   },
+   {
+    "lo": 0.333,
+    "hi": 0.4,
+    "conf": 0.3718169331550598,
+    "acc": 0.29032257199287415,
+    "frac": 0.006114398594945669
+   },
+   {
+    "lo": 0.4,
+    "hi": 0.467,
+    "conf": 0.4343971312046051,
+    "acc": 0.5079365372657776,
+    "frac": 0.012426035478711128
+   },
+   {
+    "lo": 0.467,
+    "hi": 0.533,
+    "conf": 0.5043900012969971,
+    "acc": 0.5302013158798218,
+    "frac": 0.029388559982180595
+   },
+   {
+    "lo": 0.533,
+    "hi": 0.6,
+    "conf": 0.5663641691207886,
+    "acc": 0.5602409839630127,
+    "frac": 0.032741617411375046
+   },
+   {
+    "lo": 0.6,
+    "hi": 0.667,
+    "conf": 0.6334572434425354,
+    "acc": 0.6188119053840637,
+    "frac": 0.039842210710048676
+   },
+   {
+    "lo": 0.667,
+    "hi": 0.733,
+    "conf": 0.7039550542831421,
+    "acc": 0.6941747665405273,
+    "frac": 0.040631163865327835
+   },
+   {
+    "lo": 0.733,
+    "hi": 0.8,
+    "conf": 0.7678379416465759,
+    "acc": 0.7991452813148499,
+    "frac": 0.04615384712815285
+   },
+   {
+    "lo": 0.8,
+    "hi": 0.867,
+    "conf": 0.8343891501426697,
+    "acc": 0.8534798622131348,
+    "frac": 0.05384615436196327
+   },
+   {
+    "lo": 0.867,
+    "hi": 0.933,
+    "conf": 0.9034571051597595,
+    "acc": 0.9114832282066345,
+    "frac": 0.08244576305150986
+   },
+   {
+    "lo": 0.933,
+    "hi": 1.0,
+    "conf": 0.9883154630661011,
+    "acc": 0.9902557730674744,
+    "frac": 0.6477317810058594
+   }
+  ]
+ },
+ "heldout": {
+  "raw": {
+   "overall": {
+    "n": 9888,
+    "acc": 0.7774069309234619,
+    "nll": 0.6259722709655762,
+    "brier": 0.3122439980506897,
+    "ece": 0.055361537657564505,
+    "mean_conf": 0.8327684998512268
+   },
+   "by_type": {
+    "choice": {
+     "n": 9088,
+     "acc": 0.7644146084785461,
+     "nll": 0.6617501974105835,
+     "brier": 0.3292696177959442,
+     "ece": 0.05688635497695306,
+     "mean_conf": 0.8213009238243103
+    },
+    "noul": {
+     "n": 800,
+     "acc": 0.925000011920929,
+     "nll": 0.21953491866588593,
+     "brier": 0.11883323639631271,
+     "ece": 0.042017475422659586,
+     "mean_conf": 0.9630394577980042
+    }
+   },
+   "by_source": {
+    "arc_easy": {
+     "n": 500,
+     "acc": 0.9760000109672546,
+     "nll": 0.10126849263906479,
+     "brier": 0.04353426396846771,
+     "ece": 0.03141405401313018,
+     "mean_conf": 0.9497805237770081
+    },
+    "askubuntu": {
+     "n": 359,
+     "acc": 0.5181058645248413,
+     "nll": 1.715165376663208,
+     "brier": 0.6542366743087769,
+     "ece": 0.16121354500097346,
+     "mean_conf": 0.679319441318512
+    },
+    "cifar10": {
+     "n": 300,
+     "acc": 0.9766666889190674,
+     "nll": 0.06821396946907043,
+     "brier": 0.033433299511671066,
+     "ece": 0.0244025032659336,
+     "mean_conf": 0.9716125726699829
+    },
+    "emotion": {
+     "n": 500,
+     "acc": 0.5799999833106995,
+     "nll": 1.305448293685913,
+     "brier": 0.6054953932762146,
+     "ece": 0.1998288362110465,
+     "mean_conf": 0.7798287868499756
+    },
+    "food101": {
+     "n": 300,
+     "acc": 0.8766666650772095,
+     "nll": 0.35791701078414917,
+     "brier": 0.15910951793193817,
+     "ece": 0.04321668811345984,
+     "mean_conf": 0.8755965232849121
+    },
+    "imdb": {
+     "n": 500,
+     "acc": 0.9399999976158142,
+     "nll": 0.15117090940475464,
+     "brier": 0.0857759416103363,
+     "ece": 0.033096101656337,
+     "mean_conf": 0.9650362730026245
+    },
+    "massive_intent": {
+     "n": 500,
+     "acc": 0.7459999918937683,
+     "nll": 0.9858649969100952,
+     "brier": 0.3682084381580353,
+     "ece": 0.09032174787909839,
+     "mean_conf": 0.8352954983711243
+    },
+    "mnist": {
+     "n": 300,
+     "acc": 0.95333331823349,
+     "nll": 0.15937867760658264,
+     "brier": 0.06808914244174957,
+     "ece": 0.03201363153018352,
+     "mean_conf": 0.9333906769752502
+    },
+    "ni_unseen_tasks": {
+     "n": 5529,
+     "acc": 0.752215564250946,
+     "nll": 0.6323504447937012,
+     "brier": 0.34408876299858093,
+     "ece": 0.05498220219487805,
+     "mean_conf": 0.8071978688240051
+    },
+    "pets": {
+     "n": 300,
+     "acc": 0.8500000238418579,
+     "nll": 0.5367151498794556,
+     "brier": 0.22500383853912354,
+     "ece": 0.06403641861143206,
+     "mean_conf": 0.8197709918022156
+    },
+    "pope": {
+     "n": 300,
+     "acc": 0.8999999761581421,
+     "nll": 0.33347487449645996,
+     "brier": 0.17392873764038086,
+     "ece": 0.07330050353340761,
+     "mean_conf": 0.9597114324569702
+    },
+    "pubmed_qa": {
+     "n": 500,
+     "acc": 0.7080000042915344,
+     "nll": 0.738036036491394,
+     "brier": 0.4015096426010132,
+     "ece": 0.08785723579766086,
+     "mean_conf": 0.7890949249267578
+    }
+   }
+  },
+  "calibrated": {
+   "overall": {
+    "n": 9888,
+    "acc": 0.7774069309234619,
+    "nll": 0.6107043027877808,
+    "brier": 0.30924949049949646,
+    "ece": 0.040644236667579256,
+    "mean_conf": 0.8180511593818665
+   },
+   "by_type": {
+    "choice": {
+     "n": 9088,
+     "acc": 0.7644146084785461,
+     "nll": 0.6470264792442322,
+     "brier": 0.32634666562080383,
+     "ece": 0.041729311528364656,
+     "mean_conf": 0.8061439394950867
+    },
+    "noul": {
+     "n": 800,
+     "acc": 0.925000011920929,
+     "nll": 0.1980842649936676,
+     "brier": 0.11502549052238464,
+     "ece": 0.03511890266276535,
+     "mean_conf": 0.9533174633979797
+    }
+   },
+   "by_source": {
+    "arc_easy": {
+     "n": 500,
+     "acc": 0.9760000109672546,
+     "nll": 0.10715165734291077,
+     "brier": 0.04551118239760399,
+     "ece": 0.03949616313766471,
+     "mean_conf": 0.9417195320129395
+    },
+    "askubuntu": {
+     "n": 359,
+     "acc": 0.5181058645248413,
+     "nll": 1.6524147987365723,
+     "brier": 0.6454375982284546,
+     "ece": 0.1338557621591943,
+     "mean_conf": 0.651961624622345
+    },
+    "cifar10": {
+     "n": 300,
+     "acc": 0.9766666889190674,
+     "nll": 0.07086114585399628,
+     "brier": 0.03372351825237274,
+     "ece": 0.03409448859720851,
+     "mean_conf": 0.9671555757522583
+    },
+    "emotion": {
+     "n": 500,
+     "acc": 0.5799999833106995,
+     "nll": 1.2474008798599243,
+     "brier": 0.5944139957427979,
+     "ece": 0.17934864653363677,
+     "mean_conf": 0.7593486309051514
+    },
+    "food101": {
+     "n": 300,
+     "acc": 0.8766666650772095,
+     "nll": 0.3655674159526825,
+     "brier": 0.1599694937467575,
+     "ece": 0.05726011354152369,
+     "mean_conf": 0.8594070672988892
+    },
+    "imdb": {
+     "n": 500,
+     "acc": 0.9399999976158142,
+     "nll": 0.1412155032157898,
+     "brier": 0.08349451422691345,
+     "ece": 0.025615910650256563,
+     "mean_conf": 0.955990195274353
+    },
+    "massive_intent": {
+     "n": 500,
+     "acc": 0.7459999918937683,
+     "nll": 0.9505149722099304,
+     "brier": 0.36250603199005127,
+     "ece": 0.07815012133194074,
+     "mean_conf": 0.8151710033416748
+    },
+    "mnist": {
+     "n": 300,
+     "acc": 0.95333331823349,
+     "nll": 0.1676100343465805,
+     "brier": 0.07024406641721725,
+     "ece": 0.044357041122790236,
+     "mean_conf": 0.9220753908157349
+    },
+    "ni_unseen_tasks": {
+     "n": 5529,
+     "acc": 0.752215564250946,
+     "nll": 0.6206715106964111,
+     "brier": 0.34118887782096863,
+     "ece": 0.040373205554264546,
+     "mean_conf": 0.7925888299942017
+    },
+    "pets": {
+     "n": 300,
+     "acc": 0.8500000238418579,
+     "nll": 0.5368061065673828,
+     "brier": 0.22789603471755981,
+     "ece": 0.07454185713334927,
+     "mean_conf": 0.8035118579864502
+    },
+    "pope": {
+     "n": 300,
+     "acc": 0.8999999761581421,
+     "nll": 0.29286548495292664,
+     "brier": 0.1675771325826645,
+     "ece": 0.06612932776747837,
+     "mean_conf": 0.9488630890846252
+    },
+    "pubmed_qa": {
+     "n": 500,
+     "acc": 0.7080000042915344,
+     "nll": 0.7209590077400208,
+     "brier": 0.39785510301589966,
+     "ece": 0.07103049902030684,
+     "mean_conf": 0.7723875641822815
+    }
+   }
+  },
+  "reliability_raw": [
+   {
+    "lo": 0.133,
+    "hi": 0.2,
+    "conf": 0.18254680931568146,
+    "acc": 0.0,
+    "frac": 0.00040453075780533254
+   },
+   {
+    "lo": 0.2,
+    "hi": 0.267,
+    "conf": 0.24415171146392822,
+    "acc": 0.125,
+    "frac": 0.004045307636260986
+   },
+   {
+    "lo": 0.267,
+    "hi": 0.333,
+    "conf": 0.3000350594520569,
+    "acc": 0.2711864411830902,
+    "frac": 0.011933657340705395
+   },
+   {
+    "lo": 0.333,
+    "hi": 0.4,
+    "conf": 0.36993393301963806,
+    "acc": 0.2829268276691437,
+    "frac": 0.020732199773192406
+   },
+   {
+    "lo": 0.4,
+    "hi": 0.467,
+    "conf": 0.434598833322525,
+    "acc": 0.356890469789505,
+    "frac": 0.02862055040895939
+   },
+   {
+    "lo": 0.467,
+    "hi": 0.533,
+    "conf": 0.5052576661109924,
+    "acc": 0.4368070960044861,
+    "frac": 0.04561084136366844
+   },
+   {
+    "lo": 0.533,
+    "hi": 0.6,
+    "conf": 0.5669507384300232,
+    "acc": 0.4856114983558655,
+    "frac": 0.05622977390885353
+   },
+   {
+    "lo": 0.6,
+    "hi": 0.667,
+    "conf": 0.6334361433982849,
+    "acc": 0.5653104782104492,
+    "frac": 0.047228965908288956
+   },
+   {
+    "lo": 0.667,
+    "hi": 0.733,
+    "conf": 0.701598048210144,
+    "acc": 0.6368821263313293,
+    "frac": 0.05319579318165779
+   },
+   {
+    "lo": 0.733,
+    "hi": 0.8,
+    "conf": 0.7661896347999573,
+    "acc": 0.7110726833343506,
+    "frac": 0.058454692363739014
+   },
+   {
+    "lo": 0.8,
+    "hi": 0.867,
+    "conf": 0.8357764482498169,
+    "acc": 0.7466007471084595,
+    "frac": 0.08181634545326233
+   },
+   {
+    "lo": 0.867,
+    "hi": 0.933,
+    "conf": 0.9009120464324951,
+    "acc": 0.8152466416358948,
+    "frac": 0.1127629429101944
+   },
+   {
+    "lo": 0.933,
+    "hi": 1.0,
+    "conf": 0.9835764169692993,
+    "acc": 0.9503800868988037,
+    "frac": 0.4789643883705139
+   }
+  ],
+  "reliability_calibrated": [
+   {
+    "lo": 0.133,
+    "hi": 0.2,
+    "conf": 0.17538046836853027,
+    "acc": 0.0,
+    "frac": 0.0006067961221560836
+   },
+   {
+    "lo": 0.2,
+    "hi": 0.267,
+    "conf": 0.24302583932876587,
+    "acc": 0.16981132328510284,
+    "frac": 0.005360032431781292
+   },
+   {
+    "lo": 0.267,
+    "hi": 0.333,
+    "conf": 0.30060428380966187,
+    "acc": 0.25362318754196167,
+    "frac": 0.013956310227513313
+   },
+   {
+    "lo": 0.333,
+    "hi": 0.4,
+    "conf": 0.36858007311820984,
+    "acc": 0.3234042525291443,
+    "frac": 0.023766180500388145
+   },
+   {
+    "lo": 0.4,
+    "hi": 0.467,
+    "conf": 0.4340149462223053,
+    "acc": 0.3523809611797333,
+    "frac": 0.03185679763555527
+   },
+   {
+    "lo": 0.467,
+    "hi": 0.533,
+    "conf": 0.5055366158485413,
+    "acc": 0.45656564831733704,
+    "frac": 0.05006067827343941
+   },
+   {
+    "lo": 0.533,
+    "hi": 0.6,
+    "conf": 0.5651605129241943,
+    "acc": 0.5167548656463623,
+    "frac": 0.05734223127365112
+   },
+   {
+    "lo": 0.6,
+    "hi": 0.667,
+    "conf": 0.6337770223617554,
+    "acc": 0.5630081295967102,
+    "frac": 0.04975728318095207
+   },
+   {
+    "lo": 0.667,
+    "hi": 0.733,
+    "conf": 0.7002703547477722,
+    "acc": 0.6780238747596741,
+    "frac": 0.059364885091781616
+   },
+   {
+    "lo": 0.733,
+    "hi": 0.8,
+    "conf": 0.7682633399963379,
+    "acc": 0.7181528806686401,
+    "frac": 0.06351132690906525
+   },
+   {
+    "lo": 0.8,
+    "hi": 0.867,
+    "conf": 0.8352631330490112,
+    "acc": 0.7694915533065796,
+    "frac": 0.08950242400169373
+   },
+   {
+    "lo": 0.867,
+    "hi": 0.933,
+    "conf": 0.9024354219436646,
+    "acc": 0.8416013717651367,
+    "frac": 0.11620145291090012
+   },
+   {
+    "lo": 0.933,
+    "hi": 1.0,
+    "conf": 0.9815592169761658,
+    "acc": 0.9596588015556335,
+    "frac": 0.43871358036994934
+   }
+  ]
+ },
+ "fork_invariance_max_dp": 6.145238876342773e-05,
+ "latency_ms": {
+  "1": {
+   "fork": 108.40368270874023,
+   "separate": 959.8970413208008
+  },
+  "4": {
+   "fork": 107.42926597595215,
+   "separate": 1998.2314109802246
+  },
+  "16": {
+   "fork": 113.10410499572754,
+   "separate": 7827.5275230407715
+  },
+  "64": {
+   "fork": 256.7019462585449,
+   "separate": 16031.468868255615
+  }
+ }
+}

examples/request.json

@@ -0,0 +1,30 @@
+{
+  "state": {
+    "ticket": {
+      "id": 4412,
+      "plan": "pro",
+      "subject": "Charged twice??",
+      "body": "Hi, I just noticed I was charged twice for my subscription this month. This is the second time it has happened and honestly I'm pretty fed up. Please refund the extra charge today or I'm cancelling."
+    }
+  },
+  "questions": {
+    "wants_refund": {"type": "noul", "instructions": "The customer is asking for money back."},
+    "threatens_churn": {"type": "noul", "instructions": "The customer threatens to cancel or leave."},
+    "team": {
+      "type": "choice",
+      "instructions": "Which team should handle this ticket?",
+      "criteria": {
+        "billing": "Charges, invoices, refunds, payment methods",
+        "technical": "Bugs, outages, errors, performance problems",
+        "sales": "Upgrades, pricing questions, new purchases",
+        "account": "Login, password, profile and access issues"
+      }
+    },
+    "frustration": {
+      "type": "score",
+      "instructions": "How frustrated is the customer?",
+      "criteria": ["Calm and neutral", "Mildly annoyed", "Clearly frustrated", "Furious or abusive"]
+    },
+    "language": {"type": "choice", "instructions": "Which language is the ticket written in?", "criteria": ["english", "spanish", "german", "french"]}
+  }
+}

examples/systemone_request.json

@@ -0,0 +1,31 @@
+{
+  "model": "zico-latest",
+  "state": {
+    "ticket": {
+      "id": 4412,
+      "plan": "pro",
+      "subject": "Charged twice??",
+      "body": "Hi, I just noticed I was charged twice for my subscription this month. This is the second time it has happened and honestly I'm pretty fed up. Please refund the extra charge today or I'm cancelling."
+    }
+  },
+  "questions": {
+    "wants_refund": {"type": "noul", "instructions": "The customer is asking for money back."},
+    "threatens_churn": {"type": "noul", "instructions": "The customer threatens to cancel or leave."},
+    "team": {
+      "type": "choice",
+      "instructions": "Which team should handle this ticket?",
+      "criteria": {
+        "billing": "Charges, invoices, refunds, payment methods",
+        "technical": "Bugs, outages, errors, performance problems",
+        "sales": "Upgrades, pricing questions, new purchases",
+        "account": "Login, password, profile and access issues"
+      }
+    },
+    "frustration": {
+      "type": "score",
+      "instructions": "How frustrated is the customer?",
+      "criteria": ["Calm and neutral", "Mildly annoyed", "Clearly frustrated", "Furious or abusive"]
+    },
+    "language": {"type": "choice", "instructions": "Which language is the ticket written in?", "criteria": {"english": null, "spanish": null, "german": null, "french": null}}
+  }
+}

head.safetensors

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

pyproject.toml

@@ -0,0 +1,33 @@
+[project]
+name = "zico"
+version = "0.1.0"
+description = "zico: a typed decision model - state + typed questions -> calibrated probability distributions"
+requires-python = ">=3.10"
+dependencies = [
+    "torch",
+    "transformers>=5.2",
+    "datasets",
+    "peft",
+    "safetensors",
+    "numpy",
+    "pillow",  # images in the state
+    "torchvision",  # transformers' image processor for the vision tower
+]
+
+[project.optional-dependencies]
+dev = ["pytest", "httpx"]
+serve = ["fastapi", "uvicorn"]
+mlx = ["mlx", "mlx-lm"]  # Apple silicon serving backend
+
+[build-system]
+requires = ["setuptools>=68"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.packages.find]
+include = ["zico*"]
+
+[tool.setuptools.package-data]
+zico = ["ui.html"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

scripts/add_vision.py

@@ -0,0 +1,55 @@
+"""Package the base checkpoint's vision tower into a zico checkpoint:  <ckpt>/vision/
+
+    python scripts/add_vision.py runs/qwen35-4b-best-bf16
+
+Only the ViT + merger weights are copied (a few hundred MB), together with the image-processor
+settings and the special-token ids the renderer needs. The vision tower is used as shipped: our
+fine-tuning never touched it, and the zero-shot probe showed the alignment survived.
+"""
+import argparse
+import json
+from pathlib import Path
+
+import torch
+from huggingface_hub import snapshot_download
+from safetensors import safe_open
+from safetensors.torch import save_file
+from transformers import AutoConfig, AutoImageProcessor
+
+p = argparse.ArgumentParser()
+p.add_argument("ckpt")
+p.add_argument("--base", default=None, help="defaults to the checkpoint's base_model")
+p.add_argument("--dtype", default="bf16", choices=["bf16", "fp32"])
+args = p.parse_args()
+
+ckpt = Path(args.ckpt)
+base = args.base or json.loads((ckpt / "zico.json").read_text())["base_model"]
+cfg = AutoConfig.from_pretrained(base)
+src = Path(snapshot_download(base, allow_patterns=["*.safetensors", "*.json"]))
+dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float32
+
+weights = {}
+for shard in sorted(src.glob("*.safetensors")):
+    with safe_open(shard, framework="pt") as f:
+        for key in f.keys():
+            if key.startswith("model.visual."):
+                weights[key[len("model.visual."):]] = f.get_tensor(key).to(dtype).contiguous()
+if not weights:
+    raise SystemExit(f"no vision weights found in {base}")
+
+out = ckpt / "vision"
+out.mkdir(parents=True, exist_ok=True)
+save_file(weights, out / "model.safetensors")
+vision_cfg = {k: v for k, v in cfg.vision_config.to_dict().items() if k not in ("transformers_version", "architectures")}
+(out / "config.json").write_text(json.dumps(vision_cfg, indent=2))
+AutoImageProcessor.from_pretrained(base).save_pretrained(out)
+(out / "zico_vision.json").write_text(json.dumps({
+    "base_model": base,
+    "image_token_id": cfg.image_token_id,
+    "vision_start_token_id": cfg.vision_start_token_id,
+    "vision_end_token_id": cfg.vision_end_token_id,
+    "patch_size": cfg.vision_config.patch_size,
+    "spatial_merge_size": cfg.vision_config.spatial_merge_size,
+}, indent=2))  # fmt: skip
+n = sum(w.numel() for w in weights.values())
+print(f"packaged {len(weights)} tensors, {n/1e6:.0f}M parameters, {sum(w.numel()*w.element_size() for w in weights.values())/1e9:.2f} GB -> {out}")

scripts/bench_ours.py

@@ -0,0 +1,109 @@
+"""Score a zico checkpoint on the fixed benchmark rows (the same rows every other system saw).
+
+    python scripts/bench_ours.py runs/qwen35-4b/best --data data_ref --permutations 4
+
+--permutations K averages each Choice answer over K shuffled option orders (K=1 = as given).
+A long option list is read left to right by a causal backbone, so where the right answer sits
+matters a little; averaging over orders removes that bias. It needs no training and no labels,
+and with the forked state cache the extra orders cost only their own question tokens.
+"""
+import argparse
+import collections
+import json
+import random
+import time
+
+import torch
+
+from zico import DecisionModel
+from zico.backbone import pick_device
+from zico.benchmark import BENCH_SETS, format_line, load_rows, summarise
+from zico.data import render_records
+from zico.evaluate import predict
+
+p = argparse.ArgumentParser()
+p.add_argument("ckpt")
+p.add_argument("--data", default="data_ref")
+p.add_argument("--n", type=int, default=300)
+p.add_argument("--permutations", type=int, default=1)
+p.add_argument("--out", default=None)
+p.add_argument("--max-len", type=int, default=2048)
+p.add_argument("--errors", default="massive_intent", help="source to print a confusion breakdown for")
+args = p.parse_args()
+
+dev = pick_device()
+model, renderer, _ = DecisionModel.load(args.ckpt, device=dev)
+model.autocast_dtype = torch.bfloat16
+rng = random.Random(0)
+
+# one record per (row, option order); order[j] = original index of the option shown at position j
+records, orders, owner = [], [], []
+base_rows = []
+for split, source in BENCH_SETS:
+    for r in load_rows(args.data, split, source, args.n):
+        base_rows.append(r)
+        q = r["question"]
+        k = args.permutations if q["type"] == "choice" else 1
+        crit = q.get("criteria")
+        keys = list(crit) if crit is not None else []
+        for j in range(k):
+            order = list(range(len(keys)))
+            if j:
+                rng.shuffle(order)
+            if q["type"] == "choice":
+                new = {keys[i]: crit[keys[i]] for i in order} if isinstance(crit, dict) else [keys[i] for i in order]
+                rec = {**r, "sid": f"{r['sid']}#{j}", "question": {**q, "criteria": new}, "label": order.index(r["label"])}
+            else:
+                rec = r
+            records.append(rec)
+            orders.append(order)
+            owner.append(len(base_rows) - 1)
+
+t0 = time.time()
+rows, kept = render_records(records, renderer, args.max_len)
+assert len(kept) == len(records), "rows were dropped by max_len; raise --max-len"
+preds = predict(model, rows, records, renderer.slot_id)
+elapsed = time.time() - t0
+
+probs = [None] * len(base_rows)
+for pr, order, o in zip(preds, orders, owner):
+    p_orig = torch.zeros(len(pr["p"]))
+    if order:
+        p_orig[torch.tensor(order)] = pr["p"]  # position j showed original option order[j]
+    else:
+        p_orig = pr["p"]
+    probs[o] = p_orig if probs[o] is None else probs[o] + p_orig
+
+results = {}
+by_source = collections.defaultdict(list)
+for r, pv in zip(base_rows, probs):
+    pv = pv / pv.sum()
+    by_source[r["source"]].append((int(pv.argmax()), float(pv.max()), r))
+print(f"{len(records)} forward rows for {len(base_rows)} questions in {elapsed:.0f}s (permutations={args.permutations})")
+for _, source in BENCH_SETS:
+    items = by_source[source]
+    m = summarise([a for a, _, _ in items], [r["label"] for _, _, r in items], [c for _, c, _ in items], [])
+    results[source] = m
+    print(source)
+    print(format_line(f"ours x{args.permutations}", m))
+
+if args.errors in by_source:
+    items = by_source[args.errors]
+    conf = collections.Counter()
+    pos_hits = collections.defaultdict(lambda: [0, 0])
+    for pred, _, r in items:
+        keys = list(r["question"]["criteria"])
+        third = min(2, r["label"] * 3 // len(keys))
+        pos_hits[third][0] += pred == r["label"]
+        pos_hits[third][1] += 1
+        if pred != r["label"]:
+            conf[(keys[r["label"]], keys[pred])] += 1
+    print(f"\n{args.errors}: accuracy by where the gold option sits in the list: " +
+          ", ".join(f"{['first', 'middle', 'last'][t]} third {h}/{n} = {h / n:.2f}" for t, (h, n) in sorted(pos_hits.items())))
+    print("most common confusions (gold -> predicted):")
+    for (g, pd), c in conf.most_common(12):
+        print(f"   {c:2d}  {g} -> {pd}")
+
+if args.out:
+    with open(args.out, "w") as f:
+        json.dump({"permutations": args.permutations, "results": results}, f, indent=2)

scripts/build_data.py

@@ -0,0 +1,135 @@
+"""Download the source datasets and write data/{train,val,test,heldout}.jsonl.
+
+    python scripts/build_data.py --out data                       # original 8 sources (what the 0.8B run used)
+    python scripts/build_data.py --out data_a --sources all --n-train 15000 --include-ni
+
+Held-out sources are never trained on; training rows whose text also appears in a held-out
+set are dropped. With --include-ni, 10% of Super-NaturalInstructions *tasks* are held out
+whole, which measures zero-shot transfer to instructions the model has never seen.
+"""
+import argparse
+import collections
+import json
+import os
+import random
+import re
+import sys
+import traceback
+
+from zico.data import HELDOUT_SOURCES, TRAIN_SOURCES, label_context, take, to_records, write_jsonl
+from zico.sources_extra import (
+    EXTRA_HELDOUT_SOURCES,
+    EXTRA_TRAIN_SOURCES,
+    NONE_OK,
+    SOURCE_META,
+    add_none_option,
+    extra_label_context,
+    natural_instructions_records,
+)
+
+p = argparse.ArgumentParser()
+p.add_argument("--out", default="data")
+p.add_argument("--sources", default="core", choices=["core", "all"])
+p.add_argument("--only", default=None, help="comma-separated source names (for testing adapters)")
+p.add_argument("--exclude-nc", action="store_true", help="drop sources with non-commercial licences")
+p.add_argument("--include-ni", action="store_true", help="add Super-NaturalInstructions classification tasks")
+p.add_argument("--ni-scan", type=int, default=2_000_000, help="rows of the NI corpus to scan")
+p.add_argument("--n-train", type=int, default=4000, help="source examples per training dataset")
+p.add_argument("--n-val", type=int, default=300)
+p.add_argument("--n-test", type=int, default=500)
+p.add_argument("--seed", type=int, default=0)
+p.add_argument("--quick", action="store_true", help="stream every source with a small buffer (adapter smoke test, skewed samples)")
+args = p.parse_args()
+
+if args.quick:
+    _take = take
+
+    def take(hf_id, cfg, split, n, stream, seed):  # noqa: F811
+        return _take(hf_id, cfg, split, n, True, seed, buffer_size=1000)
+
+
+rng = random.Random(args.seed)
+train_sources, heldout_sources = dict(TRAIN_SOURCES), dict(HELDOUT_SOURCES)
+if args.sources == "all":
+    train_sources.update(EXTRA_TRAIN_SOURCES)
+    heldout_sources.update(EXTRA_HELDOUT_SOURCES)
+if args.exclude_nc:
+    train_sources = {k: v for k, v in train_sources.items() if SOURCE_META.get(k, {}).get("license") != "nc"}
+if args.only:
+    keep = set(args.only.split(","))
+    train_sources = {k: v for k, v in train_sources.items() if k in keep}
+    heldout_sources = {k: v for k, v in heldout_sources.items() if k in keep}
+
+ctx = label_context()
+if args.sources == "all":
+    ctx.update(extra_label_context())
+splits = collections.defaultdict(list)
+
+
+def text_key(state) -> str:
+    text = " ".join(str(v) for v in state.values()) if isinstance(state, dict) else str(state)
+    return re.sub(r"[^a-z0-9]+", " ", text.lower()).strip()[:200]
+
+
+# ------------------------------------------------------------------ held-out first, so training can be de-duplicated against it
+for name, (hf_id, cfg, _, test_split, stream, adapter) in heldout_sources.items():
+    try:
+        test = take(hf_id, cfg, test_split, args.n_test, stream, args.seed + 1)
+        splits["heldout"] += to_records(name, test, adapter, rng, ctx, "heldout")
+        print(f"ok   {name} (held out)", flush=True)
+    except Exception:  # noqa: BLE001 - one broken source should not sink the build
+        print(f"FAIL {name}\n{traceback.format_exc()}", flush=True)
+heldout_keys = {text_key(r["state"]) for r in splits["heldout"]}
+
+for name, (hf_id, cfg, train_split, test_split, stream, adapter) in train_sources.items():
+    try:
+        # civil_comments drops most non-toxic rows inside its adapter, so over-fetch for it
+        mult = 3 if name == "civil_comments" else 1
+        pool = take(hf_id, cfg, train_split, (args.n_train + args.n_val) * mult, stream, args.seed)
+        n_val = min(args.n_val * mult, len(pool) // 10)
+        parts = {"val": to_records(name, pool[:n_val], adapter, rng, ctx, "val"),
+                 "train": to_records(name, pool[n_val:], adapter, rng, ctx, "train")}  # fmt: skip
+        if test_split:
+            test = take(hf_id, cfg, test_split, args.n_test * mult, stream, args.seed + 1)
+            parts["test"] = to_records(name, test, adapter, rng, ctx, "test")
+        for split, recs in parts.items():
+            if name in NONE_OK:
+                recs = [add_none_option(r, rng) for r in recs]
+            splits[split] += recs
+        print(f"ok   {name}: {len(parts['train'])} train rows", flush=True)
+    except Exception:  # noqa: BLE001
+        print(f"FAIL {name}\n{traceback.format_exc()}", flush=True)
+
+if args.include_ni:
+    try:
+        blocked = ("imdb", "_arc_", "pubmedqa", "pubmed_qa", "massive", "askubuntu", "emotion")  # overlap with held-out sets
+        recs = [r for r in natural_instructions_records(rng, args.ni_scan) if not any(b in r["task"].lower() for b in blocked)]
+        tasks = sorted({r["task"] for r in recs})
+        rng.shuffle(tasks)
+        held = set(tasks[: len(tasks) // 10])  # whole tasks held out; none if fewer than 10 qualified
+        for r in recs:
+            if r["task"] in held:
+                splits["heldout"].append({**r, "source": "ni_unseen_tasks", "sid": r["sid"].replace("train-", "heldout-", 1)})
+            else:
+                splits["train"].append(r)
+        print(f"ok   natural_instructions: {len(tasks) - len(held)} train tasks, {len(held)} held-out tasks, {len(recs)} rows", flush=True)
+    except Exception:  # noqa: BLE001
+        print(f"FAIL natural_instructions\n{traceback.format_exc()}", flush=True)
+
+stats = {}
+for split, records in splits.items():
+    if split != "heldout":
+        before = len(records)
+        records = [r for r in records if text_key(r["state"]) not in heldout_keys]
+        if before != len(records):
+            print(f"{split}: dropped {before - len(records)} rows that also appear in a held-out set")
+    if split == "train":
+        rng.shuffle(records)
+    write_jsonl(f"{args.out}/{split}.jsonl", records)
+    by = collections.Counter((r["source"], r["question"]["type"]) for r in records)
+    stats[split] = {"rows": len(records), "by_source_type": {f"{s}/{t}": n for (s, t), n in sorted(by.items())}}
+    print(f"{split:8s} {len(records):7d} rows")
+with open(f"{args.out}/stats.json", "w") as f:
+    json.dump(stats, f, indent=2)
+sys.stdout.flush()
+os._exit(0)  # streaming dataset readers can leave non-daemon threads behind and hang a normal exit

scripts/build_image_data.py

@@ -0,0 +1,171 @@
+"""Build the image training set: resized JPEGs + records, mixed with replayed text rows.
+
+    python scripts/build_image_data.py --out data_img --n-train 5000 --text-replay data_a --text-rows 25000
+
+Every image is resized ONCE, here, to its final aspect-preserving size (multiples of 32 px, one
+token per 32x32 px) and its token grid is written into the record. The trainer can then bucket
+rows by length without decoding a single picture and only loads pixels per batch.
+
+Text rows are replayed alongside so the text skills the model already has do not fade.
+"""
+import argparse
+import collections
+import json
+import os
+import random
+import sys
+import traceback
+from pathlib import Path
+
+from zico.data import read_jsonl, write_jsonl
+from zico.image_sources import IMAGE_HELDOUT_SOURCES, IMAGE_TRAIN_SOURCES
+from zico.vision import smart_resize
+
+p = argparse.ArgumentParser()
+p.add_argument("--out", default="data_img")
+p.add_argument("--n-train", type=int, default=5000, help="images per training source")
+p.add_argument("--doc-mult", type=float, default=2.0, help="documents are the weak spot: take this many times more")
+p.add_argument("--n-val", type=int, default=60)
+p.add_argument("--n-test", type=int, default=150)
+p.add_argument("--n-heldout", type=int, default=300)
+p.add_argument("--text-replay", default=None, help="a text data dir (train/val/test/heldout .jsonl) to mix in")
+p.add_argument("--text-rows", type=int, default=25000)
+p.add_argument("--only", default=None)
+p.add_argument("--quick", action="store_true", help="stream every source (adapter smoke test; samples may be class-skewed)")
+p.add_argument("--min-pixels", type=int, default=256 * 256)
+p.add_argument("--seed", type=int, default=0)
+args = p.parse_args()
+
+from datasets import load_dataset  # noqa: E402
+from PIL import Image  # noqa: E402
+
+rng = random.Random(args.seed)
+out = Path(args.out)
+splits = collections.defaultdict(list)
+only = set(args.only.split(",")) if args.only else None
+
+
+def take(hf_id, cfg, split, n, stream):
+    if args.quick:  # adapter smoke test only: the sample is class-skewed (see below)
+        ds = load_dataset(hf_id, cfg, split=split, streaming=True)
+        return list(ds.shuffle(seed=args.seed, buffer_size=min(4000, max(500, n))).take(n)), ds.features
+    if stream:
+        return take_uniform_stream(hf_id, cfg, split, n)
+    ds = load_dataset(hf_id, cfg, split=split).shuffle(seed=args.seed)
+    return [ds[i] for i in range(min(n, len(ds)))], ds.features
+
+
+def take_uniform_stream(hf_id, cfg, split, n):
+    """A uniform sample of a large dataset without storing it.
+
+    `streaming.shuffle(buffer).take(n)` reads only the first few files, and image datasets are usually
+    grouped by class -- both across files and inside them. The first image run learned 9 of RVL-CDIP's 16
+    document types that way (and 17 of ImageNet-100's classes) and got *worse* on documents. Here every
+    file contributes equally, and within a file the shuffle buffer is larger than the file, so the draw is
+    uniform. Files are read in parallel threads.
+    """
+    import math
+    from concurrent.futures import ThreadPoolExecutor
+
+    ds = load_dataset(hf_id, cfg, split=split, streaming=True)
+    k = ds.n_shards
+    per = math.ceil(n / k) + 2
+
+    def one(i):
+        part = ds.shard(num_shards=k, index=i) if k > 1 else ds
+        return list(part.shuffle(seed=args.seed + i, buffer_size=30_000).take(per))
+
+    with ThreadPoolExecutor(min(16, k)) as pool:
+        rows = [r for part in pool.map(one, range(k)) for r in part]
+    random.Random(args.seed).shuffle(rows)
+    return rows[:n], ds.features
+
+
+def gold_classes(records) -> set:
+    """Distinct correct answers among the Choice questions: a cheap check that a sample covers its classes."""
+    return {list(r["question"]["criteria"])[r["label"]] for r in records if r["question"]["type"] == "choice"}
+
+
+def label_names(features, rows):
+    names = {k: f.names for k, f in features.items() if getattr(f, "names", None)}
+    for col in ("masterCategory", "subCategory", "articleType", "baseColour", "gender", "usage"):  # string attributes
+        if rows and col in rows[0]:
+            names[col] = sorted({r[col] for r in rows if r.get(col)})
+    return names
+
+
+def emit(name, spec, rows, names, split):
+    _, _, _, _, _, image_col, adapter, max_pixels, keys = spec
+    records = []
+    for i, ex in enumerate(rows):
+        img = ex[image_col]
+        if img is None:
+            continue
+        img = img.convert("RGB")
+        w, h = smart_resize(img.width, img.height, 32, min(args.min_pixels, max_pixels), max_pixels)
+        rel = f"images/{name}/{split}-{i:06d}.jpg"
+        (out / rel).parent.mkdir(parents=True, exist_ok=True)
+        img.resize((w, h), Image.BICUBIC).save(out / rel, quality=92)
+        part = {"type": "image", "path": rel}
+        state = rng.choice([{rng.choice(keys): part}, part, {"input": {rng.choice(keys): part}}])
+        for question, label in adapter(ex, rng, names):
+            records.append({"source": name, "sid": f"{split}-{name}-{i}", "state": state, "question": question, "label": label,
+                            "image_grids": [[h // 32, w // 32]]})  # fmt: skip
+    return records
+
+
+for name, spec in IMAGE_HELDOUT_SOURCES.items():
+    if only and name not in only:
+        continue
+    try:
+        rows, feats = take(spec[0], spec[1], spec[3], args.n_heldout, spec[4])
+        splits["heldout"] += emit(name, spec, rows, label_names(feats, rows), "heldout")
+        print(f"ok   {name} (held out)", flush=True)
+    except Exception:  # noqa: BLE001
+        print(f"FAIL {name}\n{traceback.format_exc()}", flush=True)
+
+for name, spec in IMAGE_TRAIN_SOURCES.items():
+    if only and name not in only:
+        continue
+    try:
+        n = int(args.n_train * (args.doc_mult if name == "rvl_cdip" else 1))
+        pool, feats = take(spec[0], spec[1], spec[2], n + args.n_val, spec[4])
+        names = label_names(feats, pool)
+        splits["val"] += emit(name, spec, pool[: args.n_val], names, "val")
+        train = emit(name, spec, pool[args.n_val :], names, "train")
+        splits["train"] += train
+        n_test_classes = 0
+        if spec[3]:
+            test, _ = take(spec[0], spec[1], spec[3], args.n_test, spec[4])
+            test_recs = emit(name, spec, test, names, "test")
+            splits["test"] += test_recs
+            n_test_classes = len(gold_classes(test_recs))
+        n_all = max((len(v) for v in names.values()), default=0)
+        covered = len(gold_classes(train))
+        flag = "   <-- WARNING: skewed sample" if n_all and covered < 0.6 * min(n_all, len(train) / 4) else ""
+        print(f"ok   {name}: {len(train)} train rows, {covered} distinct answer classes in train (largest label set {n_all}), "
+              f"{n_test_classes} in test{flag}", flush=True)  # fmt: skip
+    except Exception:  # noqa: BLE001
+        print(f"FAIL {name}\n{traceback.format_exc()}", flush=True)
+
+if args.text_replay:
+    n_img = len(splits["train"])
+    text = read_jsonl(f"{args.text_replay}/train.jsonl")
+    splits["train"] += rng.sample(text, min(args.text_rows, len(text)))
+    for split, n in (("val", 1500), ("test", 3000)):
+        rows = read_jsonl(f"{args.text_replay}/{split}.jsonl")
+        splits[split] += rng.sample(rows, min(n, len(rows)))
+    splits["heldout"] += read_jsonl(f"{args.text_replay}/heldout.jsonl")  # text held-out sets: did anything fade?
+    print(f"text replay: {len(splits['train']) - n_img} text rows added to {n_img} image rows")
+
+stats = {}
+for split, records in splits.items():
+    if split == "train":
+        rng.shuffle(records)
+    write_jsonl(out / f"{split}.jsonl", records)
+    by = collections.Counter(r["source"] for r in records)
+    stats[split] = {"rows": len(records), "image_rows": sum("image_grids" in r for r in records), "by_source": dict(sorted(by.items()))}
+    print(f"{split:8s} {len(records):7d} rows ({stats[split]['image_rows']} with images)")
+(out / "stats.json").write_text(json.dumps(stats, indent=2))
+sys.stdout.flush()
+os._exit(0)  # streaming readers leave threads behind that hang a normal exit

scripts/check_kernels.py

@@ -0,0 +1,79 @@
+"""Numerically compare the fused Gated-DeltaNet kernel (fla) against an independent reference,
+forward AND backward. Run once on any new GPU / Triton combination before training:
+
+    python scripts/check_kernels.py
+
+Motivation: fla refuses to run its backward on Hopper GPUs with Triton 3.4-3.7.0 because that
+combination silently produces wrong gradients. Whatever workaround is installed (newer Triton,
+tilelang backend) should be verified rather than trusted.
+
+The reference is a naive token-by-token recurrence in float64 written here. (HF's
+`torch_chunk_gated_delta_rule` is NOT independent: it is decorated to dispatch to fla when fla is
+installed, so comparing against it compares fla with itself.)
+
+    S_t = exp(g_t) * S_{t-1}
+    S_t = S_t + k_t (beta_t * (v_t - S_t^T k_t))^T
+    o_t = S_t^T q_t                      with q, k L2-normalised and q scaled by Dk^-1/2
+"""
+import torch
+import torch.nn.functional as F
+from fla.ops.gated_delta_rule import chunk_gated_delta_rule
+
+dev = "cuda"
+B, L, H, DK, DV = 2, 200, 4, 64, 64  # L deliberately not a multiple of the 64-token chunk
+
+
+def naive(q, k, v, g, beta):
+    q, k = F.normalize(q, dim=-1), F.normalize(k, dim=-1)
+    q = q * DK**-0.5
+    S = q.new_zeros(B, H, DK, DV)
+    outs = []
+    for t in range(L):
+        S = S * g[:, t].exp()[..., None, None]
+        pred = (S * k[:, t][..., None]).sum(-2)  # S^T k  -> (B, H, DV)
+        S = S + k[:, t][..., None] * ((v[:, t] - pred) * beta[:, t][..., None])[..., None, :]
+        outs.append((S * q[:, t][..., None]).sum(-2))
+    return torch.stack(outs, 1)
+
+
+def make(dtype):
+    gen = torch.Generator(device=dev).manual_seed(1)
+    q = torch.randn(B, L, H, DK, device=dev, generator=gen)
+    k = torch.randn(B, L, H, DK, device=dev, generator=gen)
+    v = torch.randn(B, L, H, DV, device=dev, generator=gen)
+    g = -torch.rand(B, L, H, device=dev, generator=gen) * 0.5  # log-space decay, <= 0
+    beta = torch.rand(B, L, H, device=dev, generator=gen)
+    w = torch.randn(B, L, H, DV, device=dev, generator=gen)  # fixed projection: same loss on both paths
+    leaves = [t.to(dtype).requires_grad_(True) for t in (q, k, v, g, beta)]
+    return leaves, w
+
+
+def grads_of(out, leaves, w):
+    (out.double() * w.double()).sum().backward()
+    return out.double().detach(), [t.grad.double() for t in leaves]
+
+
+leaves, w = make(torch.float64)
+ref_out, ref_grads = grads_of(naive(*leaves), leaves, w)
+
+ok = True
+rel = lambda a, b: ((a - b).norm() / b.norm().clamp_min(1e-12)).item()  # noqa: E731
+for dtype in (torch.float32, torch.bfloat16):
+    (q, k, v, g, beta), w = make(dtype)
+    try:
+        out = chunk_gated_delta_rule(q, k, v, g=g.float() if dtype == torch.bfloat16 else g, beta=beta, use_qk_l2norm_in_kernel=True)
+        out = out[0] if isinstance(out, tuple) else out
+        out, grads = grads_of(out, (q, k, v, g, beta), w)
+    except Exception as e:  # noqa: BLE001
+        print(f"fla {dtype}: FAILED to run: {type(e).__name__}: {str(e)[:300]}")
+        ok = False
+        continue
+    errs = {"out": rel(out, ref_out), **{n: rel(a, b) for n, a, b in zip(("dq", "dk", "dv", "dg", "dbeta"), grads, ref_grads)}}
+    # Triton's tl.dot runs in TF32 on Ampere+ GPUs, so even "fp32" agrees only to ~1e-3. A genuinely
+    # wrong backward shows up as O(0.1-1) error on specific gradients, not a uniform small offset
+    # shared with the forward output.
+    tol = 1e-2 if dtype == torch.float32 else 5e-2
+    bad = {n: e for n, e in errs.items() if not e < tol}
+    ok &= not bad
+    print(f"fla {str(dtype):15s} rel-err vs naive fp64: " + "  ".join(f"{n}={e:.1e}" for n, e in errs.items()) + ("   <-- MISMATCH" if bad else "   ok"))
+print("KERNEL CHECK PASSED" if ok else "KERNEL CHECK FAILED")

scripts/check_mlx_parity.py

@@ -0,0 +1,88 @@
+"""Compare the MLX backend with the PyTorch backend on identical requests, then time both.
+
+    python scripts/check_mlx_parity.py runs/qwen35-4b-best-bf16 --data data --states 40
+"""
+import argparse
+import collections
+import json
+import statistics
+import time
+
+import torch
+
+from zico import DecisionModel, question_from_dict
+from zico.backbone import pick_device
+from zico.data import read_jsonl
+from zico.mlx_backend import MLXDecisionModel
+
+p = argparse.ArgumentParser()
+p.add_argument("ckpt")
+p.add_argument("--data", default="data")
+p.add_argument("--states", type=int, default=40)
+p.add_argument("--skip-torch-timing", action="store_true")
+args = p.parse_args()
+
+# requests: real rows grouped by state (so multi-question states are exercised) + the example ticket
+by_state = collections.OrderedDict()
+for split in ("heldout", "test"):
+    for r in read_jsonl(f"{args.data}/{split}.jsonl"):
+        by_state.setdefault(r["sid"], {"state": r["state"], "questions": {}})["questions"][f"q{len(by_state[r['sid']]['questions'])}"] = r["question"]
+multi = [v for v in by_state.values() if len(v["questions"]) > 1][: args.states // 2]
+single = [v for v in list(by_state.values())[:: max(1, len(by_state) // args.states)] if len(v["questions"]) == 1][: args.states - len(multi)]
+example = json.load(open("examples/systemone_request.json"))
+requests = [{"state": example["state"], "questions": example["questions"]}] + multi + single
+print(f"{len(requests)} requests, {sum(len(r['questions']) for r in requests)} questions")
+
+dev = pick_device()
+torch_model, renderer, _ = DecisionModel.load(args.ckpt, device=dev, dtype=torch.bfloat16)
+mlx_model, mlx_renderer, _ = MLXDecisionModel.load(args.ckpt)
+
+
+def flat(answer):
+    a = answer
+    if a["type"] == "noul":
+        return [1 - a["noul"], a["noul"]]
+    return list(a["probabilities"].values())
+
+
+def top(answer):
+    return answer.get("choice", round(answer.get("noul", answer.get("score", 0))))
+
+
+diffs, agree, total, t_torch, t_mlx = [], 0, 0, [], []
+for i, req in enumerate(requests):
+    qs = {k: question_from_dict(v) for k, v in req["questions"].items()}
+    t0 = time.perf_counter()
+    a = torch_model.answer(renderer, req["state"], qs)
+    if dev == "mps":
+        torch.mps.synchronize()
+    t1 = time.perf_counter()
+    b = mlx_model.answer(mlx_renderer, req["state"], qs)
+    t2 = time.perf_counter()
+    if i:  # first request warms both up
+        t_torch.append(t1 - t0)
+        t_mlx.append(t2 - t1)
+    assert a["usage"] == b["usage"], (a["usage"], b["usage"])
+    for name in qs:
+        pa, pb = flat(a["answers"][name]), flat(b["answers"][name])
+        diffs += [abs(x - y) for x, y in zip(pa, pb)]
+        same = top(a["answers"][name]) == top(b["answers"][name])
+        agree += same
+        total += 1
+        if not same:  # expected only for near-ties, where bf16 rounding decides
+            two = lambda p: [round(v, 4) for v in sorted(p, reverse=True)[:2]]  # noqa: E731
+            print(f"  top answer differs on a {a['answers'][name]['type']} question: torch top-2 {two(pa)} vs mlx top-2 {two(pb)} "
+                  f"(torch says {top(a['answers'][name])!r}, mlx says {top(b['answers'][name])!r})")  # fmt: skip
+
+print(f"probabilities compared: {len(diffs)} | max |dp| = {max(diffs):.4f} | mean |dp| = {statistics.mean(diffs):.5f} | "
+      f"same top answer: {agree}/{total}")  # fmt: skip
+print(f"median latency per request  PyTorch ({dev}, bf16): {statistics.median(t_torch)*1000:.0f} ms   MLX: {statistics.median(t_mlx)*1000:.0f} ms   "
+      f"-> {statistics.median(t_torch)/statistics.median(t_mlx):.1f}x")  # fmt: skip
+ex = mlx_model.answer(mlx_renderer, example["state"], {k: question_from_dict(v) for k, v in example["questions"].items()})
+print("example ticket via MLX:", {k: (v.get("choice") or v.get("noul") or v.get("score")) for k, v in ex["answers"].items()})
+reps = []
+for _ in range(5):
+    t0 = time.perf_counter()
+    mlx_model.answer(mlx_renderer, example["state"], {k: question_from_dict(v) for k, v in example["questions"].items()})
+    reps.append(time.perf_counter() - t0)
+print(f"example ticket (5 questions, {ex['usage']['input_tokens']} tokens) via MLX: best {min(reps)*1000:.0f} ms, median {statistics.median(reps)*1000:.0f} ms")

scripts/eval.py

@@ -0,0 +1,126 @@
+"""Evaluate a checkpoint: accuracy + calibration on in-distribution test data and on
+held-out schemas, post-hoc temperature scaling, and the fork==single invariance / latency check.
+
+    python scripts/eval.py runs/qwen35-0.8b/best --save-calibrated
+"""
+import argparse
+import json
+import time
+from pathlib import Path
+
+import torch
+
+from zico import DecisionModel, question_from_dict
+from zico.backbone import pick_device
+from zico.calibration import fit_temperature
+from zico.data import read_jsonl, render_records, subsample
+from zico.evaluate import format_report, predict, report, strip_tables
+from zico.schema import QTYPES
+
+p = argparse.ArgumentParser()
+p.add_argument("ckpt")
+p.add_argument("--data", default="data")
+p.add_argument("--max-len", type=int, default=768)
+p.add_argument("--limit", type=int, default=None, help="cap rows per split (quick runs)")
+p.add_argument("--precision", default="bf16", choices=["bf16", "fp32"])
+p.add_argument("--save-calibrated", action="store_true", help="write fitted temperatures into the checkpoint")
+args = p.parse_args()
+
+dev = pick_device()
+model, renderer, meta = DecisionModel.load(args.ckpt, device=dev)
+model.autocast_dtype = torch.bfloat16 if args.precision == "bf16" else None
+results = {"ckpt": args.ckpt, "meta": meta}
+
+
+def run(split):
+    recs = subsample(read_jsonl(f"{args.data}/{split}.jsonl"), args.limit)
+    image_meta = model.image_pre.meta if model.image_pre is not None else None
+    rows, recs = render_records(recs, renderer, args.max_len, image_meta, args.data)
+    return predict(model, rows, recs, renderer.slot_id), recs
+
+
+def rescale(preds, temps):
+    out = []
+    for r in preds:
+        lg = r["logits"] / temps[r["qtype"]]
+        out.append({**r, "p": torch.softmax(lg, -1)})
+    return out
+
+
+# ------------------------------------------------------------ temperature from val
+model.log_temp.zero_()
+val_preds, _ = run("val")
+temps = []
+for qi, name in enumerate(QTYPES):
+    sel = [r for r in val_preds if r["qtype"] == qi]
+    temps.append(fit_temperature([r["logits"] for r in sel], [r["label"] for r in sel]) if len(sel) >= 50 else 1.0)
+print("fitted temperatures:", {n: round(t, 3) for n, t in zip(QTYPES, temps)})
+results["temperatures"] = dict(zip(QTYPES, temps))
+
+# ------------------------------------------------------------ test + held-out
+for split in ("test", "heldout"):
+    preds, _ = run(split)
+    raw, cal = report(preds), report(rescale(preds, temps))
+    print("\n" + format_report(raw, f"{split} (raw)"))
+    print("\n" + format_report(cal, f"{split} (temp-scaled)"))
+    results[split] = {"raw": strip_tables(raw), "calibrated": strip_tables(cal),
+                      "reliability_raw": raw["overall"]["reliability"], "reliability_calibrated": cal["overall"]["reliability"]}  # fmt: skip
+
+# ------------------------------------------------------------ fork invariance + fan-out latency
+recs = [r for r in read_jsonl(f"{args.data}/test.jsonl") if "image_grids" not in r]  # a text state, text questions
+if not recs:
+    (Path(args.ckpt) / "eval.json").write_text(json.dumps(results, indent=2))
+    raise SystemExit("no text rows in the test split: skipping the fork / latency section")
+state = next((r["state"] for r in recs if r["source"] == "yelp"), recs[0]["state"])
+pool = [question_from_dict(r["question"]) for r in recs if r["source"] in ("yelp", "go_emotions", "civil_comments", "ag_news")] \
+    or [question_from_dict(r["question"]) for r in recs]  # fmt: skip
+seen, questions = set(), []
+for q in pool:  # distinct questions only
+    key = json.dumps([q.type, q.instructions, q.criteria], sort_keys=True)
+    if key not in seen:
+        seen.add(key)
+        questions.append(q)
+questions = questions[:64]
+
+
+def sync():
+    if dev == "mps":
+        torch.mps.synchronize()
+    elif dev == "cuda":
+        torch.cuda.synchronize()
+
+
+def timed(fn, reps=3):
+    best = float("inf")
+    for _ in range(reps):
+        sync(); t = time.time(); out = fn(); sync()  # noqa: E702
+        best = min(best, time.time() - t)
+    return out, best * 1000
+
+
+model.autocast_dtype = None  # compare the two paths in fp32 so differences mean logic, not rounding
+single = [model.question_logits(renderer, state, [q], fork=False)[0][0] for q in questions[:16]]
+forked = model.question_logits(renderer, state, questions[:16], fork=True)[0]
+gap = max((torch.softmax(a, -1) - torch.softmax(b, -1)).abs().max().item() for a, b in zip(single, forked))
+print(f"\nfork vs single-question, max |delta p| over 16 questions (fp32): {gap:.2e}")
+results["fork_invariance_max_dp"] = gap
+
+model.autocast_dtype = torch.bfloat16 if args.precision == "bf16" else None
+print(f"fan-out latency on one state ({len(renderer.state(state))} state tokens):")
+results["latency_ms"] = {}
+for n in (1, 4, 16, 64):
+    if n > len(questions):
+        break
+    qs = questions[:n]
+    _, t_fork = timed(lambda: model.question_logits(renderer, state, qs, fork=True, max_batch=64))
+    _, t_seq = timed(lambda: [model.question_logits(renderer, state, [q], fork=False) for q in qs], reps=1)
+    print(f"  {n:3d} questions: forked batch {t_fork:7.0f} ms | {n} separate calls {t_seq:7.0f} ms  ({t_seq/t_fork:.1f}x)")
+    results["latency_ms"][n] = {"fork": t_fork, "separate": t_seq}
+
+if args.save_calibrated:
+    with torch.no_grad():
+        model.log_temp.copy_(torch.tensor(temps).log())
+    head = {k: v for k, v in model.state_dict().items() if not k.startswith("backbone.")}
+    torch.save(head, Path(args.ckpt) / "head.pt")
+    print("saved temperatures into", Path(args.ckpt) / "head.pt")
+(Path(args.ckpt) / "eval.json").write_text(json.dumps(results, indent=2))

scripts/eval_images.py

@@ -0,0 +1,82 @@
+"""Zero-shot image decisions through the production path: base64 image parts inside `state`.
+
+    python scripts/eval_images.py runs/qwen35-4b-best-bf16 --tasks cifar10,documents --detail auto,high --n 100
+
+(scripts/probe_vision.py was the original feasibility probe; this goes through DecisionModel.answer,
+i.e. exactly what the server runs, with aspect-preserving resize and the per-image "detail" hint.)
+"""
+import argparse
+import base64
+import io
+import json
+import random
+import time
+from pathlib import Path
+
+import torch
+from datasets import load_dataset
+
+from zico import Choice, DecisionModel
+from zico.backbone import pick_device
+from zico.calibration import ece
+
+p = argparse.ArgumentParser()
+p.add_argument("ckpt")
+p.add_argument("--tasks", default="cifar10,documents")
+p.add_argument("--detail", default="auto", help="comma-separated: auto, low, high")
+p.add_argument("--n", type=int, default=100)
+p.add_argument("--out", default="logs/eval_images.json")
+args = p.parse_args()
+
+SPECS = {  # name: (dataset, split, image column, label column, instructions, streamed?)
+    "cifar10": ("uoft-cs/cifar10", "test", "img", "label", "What does the image show?", False),
+    "mnist": ("ylecun/mnist", "test", "image", "label", "Which digit is written in the image?", False),
+    "pets": ("timm/oxford-iiit-pet", "test", "image", "label", "Which breed is the animal in the photo?", True),
+    "food101": ("ethz/food101", "validation", "image", "label", "Which dish is shown in the photo?", True),
+    "documents": ("jordyvl/rvl_cdip_100_examples_per_class", "test", "image", "label", "What type of document is this scan?", False),
+}
+
+dev = pick_device()
+model, renderer, _ = DecisionModel.load(args.ckpt, device=dev, dtype=torch.bfloat16 if dev != "cpu" else torch.float32,
+                                        renderer_overrides={"max_state_tokens": 8192})  # serving limit, not the training-time one
+assert model.supports_images, "run scripts/add_vision.py on this checkpoint first"
+
+
+def encode(img) -> str:
+    buf = io.BytesIO()
+    img.convert("RGB").save(buf, format="PNG")
+    return base64.b64encode(buf.getvalue()).decode()
+
+
+results = {}
+for task in args.tasks.split(","):
+    name, split, img_col, lab_col, instructions, stream = SPECS[task]
+    if stream:
+        ds = load_dataset(name, split=split, streaming=True).shuffle(seed=0, buffer_size=3000)
+        rows, names = list(ds.take(args.n)), load_dataset(name, split=split, streaming=True).features[lab_col].names
+    else:
+        ds = load_dataset(name, split=split).shuffle(seed=0)
+        rows, names = [ds[i] for i in range(min(args.n, len(ds)))], ds.features[lab_col].names
+    keys = [str(x).replace(" ", "_").lower() for x in names]
+    for detail in args.detail.split(","):
+        rng = random.Random(0)
+        conf, correct, tokens, t0 = [], [], [], time.time()
+        for r in rows:
+            order = keys[:]
+            rng.shuffle(order)
+            part = {"type": "image", "data": encode(r[img_col])}
+            if detail != "auto":
+                part["detail"] = detail
+            out = model.answer(renderer, {"image": part}, {"q": Choice(instructions, order)})
+            a = out["answers"]["q"]
+            conf.append(max(a["probabilities"].values()))
+            correct.append(a["choice"] == keys[r[lab_col]])
+            tokens.append(out["usage"]["input_tokens"])
+        e, _ = ece(torch.tensor(conf), torch.tensor(correct))
+        m = {"n": len(rows), "options": len(keys), "acc": sum(correct) / len(correct), "ece": e,
+             "mean_tokens": sum(tokens) / len(tokens), "sec_per_image": (time.time() - t0) / len(rows)}  # fmt: skip
+        results[f"{task}/{detail}"] = m
+        print(f"{task:10s} detail={detail:5s} n={m['n']:4d} options={m['options']:3d}  acc={m['acc']:.3f}  ece={m['ece']:.3f}  "
+              f"tokens/request={m['mean_tokens']:.0f}  {m['sec_per_image']:.2f}s/img", flush=True)  # fmt: skip
+Path(args.out).parent.mkdir(parents=True, exist_ok=True)
+Path(args.out).write_text(json.dumps(results, indent=2))

scripts/infer.py

@@ -0,0 +1,37 @@
+"""Answer a System One request (state + typed questions) with a trained checkpoint.
+
+    python scripts/infer.py runs/qwen35-0.8b/best examples/request.json
+
+Request:  {"state": <str | object>, "questions": {"name": {"type": "choice|score|noul", "instructions": ..., "criteria": ...}}}
+Response: {"answers": {"name": {...}}, "usage": {...}, "model": ..., "latency_ms": ...}
+"""
+import argparse
+import json
+import sys
+import time
+
+import torch
+
+from zico import DecisionModel, question_from_dict
+from zico.backbone import pick_device
+
+p = argparse.ArgumentParser()
+p.add_argument("ckpt")
+p.add_argument("request", nargs="?", default="-", help="path to a JSON request, or - for stdin")
+p.add_argument("--precision", default="bf16", choices=["bf16", "fp32"])
+args = p.parse_args()
+
+req = json.load(sys.stdin if args.request == "-" else open(args.request))
+dev = pick_device()
+model, renderer, meta = DecisionModel.load(args.ckpt, device=dev)
+model.autocast_dtype = torch.bfloat16 if args.precision == "bf16" else None
+questions = {name: question_from_dict(q) for name, q in req["questions"].items()}
+
+model.answer(renderer, req["state"], questions)  # warm-up (kernel compilation)
+t0 = time.time()
+resp = model.answer(renderer, req["state"], questions)
+if dev == "mps":
+    torch.mps.synchronize()
+resp["latency_ms"] = round((time.time() - t0) * 1000, 1)
+resp["model"] = f"zico/{meta['base_model']}"
+print(json.dumps(resp, indent=2))

scripts/make_soft_targets.py

@@ -0,0 +1,50 @@
+"""Label rows with a model's own probabilities, to use as soft training targets ("anchor" rows).
+
+    python scripts/make_soft_targets.py runs/qwen35-4b-best-bf16 --rows data_b/train.jsonl --n 80000 --out data_img3/anchor.jsonl
+
+Why: every further full fine-tune we tried made the model more overconfident on tasks it had never
+seen (the second text pass, both image runs). Training the replayed text rows toward the *current*
+model's distribution instead of toward hard labels keeps its text behaviour -- including how unsure it
+is -- where it was, while image rows still learn from hard labels. Soft targets are just another
+distribution for the proper scoring losses, so nothing else in the trainer changes. Rows from tasks
+the teacher never trained on are especially useful: they pin down its zero-shot behaviour.
+"""
+import argparse
+import json
+import random
+
+import torch
+
+from zico import DecisionModel
+from zico.backbone import pick_device
+from zico.data import read_jsonl, render_records
+from zico.evaluate import predict
+
+p = argparse.ArgumentParser()
+p.add_argument("teacher")
+p.add_argument("--rows", required=True, help="jsonl of text records")
+p.add_argument("--n", type=int, default=80000)
+p.add_argument("--out", required=True)
+p.add_argument("--max-len", type=int, default=1024)
+p.add_argument("--seed", type=int, default=0)
+args = p.parse_args()
+
+records = [r for r in read_jsonl(args.rows) if "image_grids" not in r]
+random.Random(args.seed).shuffle(records)
+records = records[: args.n]
+
+dev = pick_device()
+model, renderer, _ = DecisionModel.load(args.teacher, device=dev)
+model.autocast_dtype = torch.bfloat16 if dev != "cpu" else None
+rows, kept = render_records(records, renderer, args.max_len)
+preds = predict(model, rows, kept, renderer.slot_id, max_tokens=16384, max_rows=256)
+
+agree = conf = 0.0
+with open(args.out, "w") as f:
+    for rec, pr in zip(kept, preds):
+        probs = [round(float(x), 6) for x in pr["p"]]
+        agree += int(pr["p"].argmax()) == rec["label"]
+        conf += max(probs)
+        f.write(json.dumps({**rec, "soft": probs, "anchor": True}, ensure_ascii=False) + "\n")
+n = len(kept)
+print(f"wrote {n} anchor rows to {args.out}: teacher agrees with the gold label on {agree / n:.3f}, mean confidence {conf / n:.3f}")

scripts/run_4b.sh

@@ -0,0 +1,34 @@
+#!/usr/bin/env bash
+# Runbook for the 4B run. Each stage is restartable; long stages should be launched under nohup/tmux.
+#   bash scripts/run_4b.sh data     # build the broad dataset            (~1 h, CPU + network)
+#   bash scripts/run_4b.sh speed    # measure tokens/s and memory        (~5 min)
+#   bash scripts/run_4b.sh train    # train; re-running resumes          (hours)
+#   bash scripts/run_4b.sh eval     # test + held-out + calibration
+# Tunables (env): MODEL, DATA, RUN, N_TRAIN, EPOCHS, MAX_TOKENS, ACCUM, LR
+set -euo pipefail
+cd "$(dirname "$0")/.."
+# shellcheck disable=SC1091
+if [ -f /venv/main/bin/activate ]; then source /venv/main/bin/activate; else source .venv/bin/activate; fi
+export TOKENIZERS_PARALLELISM=false HF_HUB_DISABLE_PROGRESS_BARS=1
+
+MODEL=${MODEL:-Qwen/Qwen3.5-4B-Base}
+DATA=${DATA:-data_a}
+RUN=${RUN:-runs/qwen35-4b}
+mkdir -p logs
+
+case "${1:-}" in
+  data)
+    python scripts/build_data.py --out "$DATA" --sources all --include-ni \
+      --n-train "${N_TRAIN:-15000}" --n-val 300 --n-test 500 ;;
+  speed)
+    python scripts/spike_backbone.py "$MODEL" --dtype bf16 ;;
+  train)
+    python scripts/train.py --model "$MODEL" --data "$DATA" --out "$RUN" --resume \
+      --grad-checkpointing --max-tokens "${MAX_TOKENS:-16384}" --max-rows 256 --grad-accum "${ACCUM:-4}" \
+      --max-len 1024 --max-state-tokens 768 --lr "${LR:-1e-5}" --head-lr 5e-4 \
+      --epochs "${EPOCHS:-1}" --eval-every 250 --val-rows 3000 --log-every 10 ;;
+  eval)
+    python scripts/eval.py "$RUN/best" --data "$DATA" --max-len 1024 --save-calibrated ;;
+  *)
+    echo "usage: $0 {data|speed|train|eval}"; exit 1 ;;
+esac

scripts/run_img.sh

@@ -0,0 +1,34 @@
+#!/usr/bin/env bash
+# Runbook for the image fine-tune on a rented CUDA box. Each stage is restartable.
+#   bash scripts/run_img.sh vision   # package the base ViT into the starting checkpoint      (~2 min)
+#   bash scripts/run_img.sh data     # image datasets + replayed text -> data_img/            (~45 min, CPU + network)
+#   bash scripts/run_img.sh speed    # 40 optimizer steps without saving: tokens/s and memory (~5 min)
+#   bash scripts/run_img.sh train    # train; re-running resumes                              (~1-1.5 h)
+#   bash scripts/run_img.sh eval     # image + text test/held-out, calibration, text same-row benchmark
+# Needs on the box first: the starting checkpoint (INIT), the text data to replay (TEXT), data_ref/.
+set -euo pipefail
+cd "$(dirname "$0")/.."
+# shellcheck disable=SC1091
+if [ -f /venv/main/bin/activate ]; then source /venv/main/bin/activate; else source .venv/bin/activate; fi
+export TOKENIZERS_PARALLELISM=false HF_HUB_DISABLE_PROGRESS_BARS=1
+
+INIT=${INIT:-runs/qwen35-4b-best-bf16}
+TEXT=${TEXT:-data_a}
+DATA=${DATA:-data_img}
+RUN=${RUN:-runs/qwen35-4b-img}
+# lower peak learning rate than the text runs: the model is already good at images zero-shot, and the
+# second text pass showed how quickly calibration degrades when we push
+TRAIN_ARGS="--init $INIT --data $DATA --grad-checkpointing --max-tokens ${MAX_TOKENS:-16384} --max-rows 128 \
+  --grad-accum ${ACCUM:-4} --max-len 1024 --lr ${LR:-5e-6} --head-lr 2e-4 --warmup 0.05 --epochs ${EPOCHS:-1} --log-every 10"
+mkdir -p logs
+
+case "${1:-}" in
+  vision) python scripts/add_vision.py "$INIT" ;;
+  data)   python scripts/build_image_data.py --out "$DATA" --n-train "${N_TRAIN:-5000}" --text-replay "$TEXT" --text-rows "${TEXT_ROWS:-25000}" ;;
+  speed)  python scripts/train.py $TRAIN_ARGS --out runs/_speed --train-rows 9000 --val-rows 100 --eval-every 100000 --no-save ;;
+  train)  python scripts/train.py $TRAIN_ARGS --out "$RUN" --resume --eval-every 150 --val-rows 2500 ;;
+  eval)
+    python scripts/eval.py "$RUN/best" --data "$DATA" --max-len 1024 --save-calibrated
+    [ -d data_ref ] && python -u scripts/bench_ours.py "$RUN/best" --data data_ref --permutations 1 --out logs/bench_img_x1.json ;;
+  *) echo "usage: $0 {vision|data|speed|train|eval}"; exit 1 ;;
+esac

scripts/serve.sh

@@ -0,0 +1,63 @@
+#!/usr/bin/env bash
+# Run the System-One-compatible server for a zico checkpoint.
+#   scripts/serve.sh start | stop | restart | keep-awake | status | log | key-path
+# Config (env): ZICO_CKPT, ZICO_HOST, ZICO_PORT, ZICO_MAX_STATE_TOKENS, ZICO_MODEL_NAME
+# By default it listens on this machine's Tailscale address if there is one (reachable from your
+# other devices, not from the local network or the internet), otherwise on 127.0.0.1.
+set -euo pipefail
+cd "$(dirname "$0")/.."
+if [ -f /venv/main/bin/activate ]; then source /venv/main/bin/activate; else source .venv/bin/activate; fi
+mkdir -p logs
+PIDFILE=logs/server.pid
+LOG=logs/server.log
+
+tailscale_ip() {
+  for t in tailscale /Applications/Tailscale.app/Contents/MacOS/Tailscale; do
+    if command -v "$t" >/dev/null 2>&1 || [ -x "$t" ]; then "$t" ip -4 2>/dev/null | head -1 && return; fi
+  done
+}
+# Optional serve.env (KEY=VALUE lines) holds this machine's choices, e.g. which checkpoint is live.
+# Variables already set in the environment win, so a one-off override still works.
+if [ -f serve.env ]; then
+  while IFS='=' read -r key value; do
+    case "$key" in ''|\#*) continue ;; esac
+    if [ -z "${!key:-}" ]; then export "$key=$value"; fi
+  done < serve.env
+fi
+export ZICO_CKPT=${ZICO_CKPT:-.}
+export ZICO_PORT=${ZICO_PORT:-8787}  # 8080 is commonly taken (llama.cpp's server defaults to it)
+if [ -z "${ZICO_HOST:-}" ]; then ZICO_HOST=$(tailscale_ip || true); fi
+export ZICO_HOST=${ZICO_HOST:-127.0.0.1}
+export TOKENIZERS_PARALLELISM=false HF_HUB_OFFLINE=1
+
+running() { [ -f "$PIDFILE" ] && kill -0 "$(cat "$PIDFILE")" 2>/dev/null; }
+
+# Keep a Mac awake while the server lives (no effect elsewhere). -s matters as much as -i: a Mac that has
+# already dozed off only surfaces for ~45 s "dark wakes", and those ignore the idle assertion alone.
+awake() { if command -v caffeinate >/dev/null 2>&1; then (nohup caffeinate -i -s "$@" >/dev/null 2>&1 < /dev/null &); fi; }
+
+start() {
+  if running; then echo "already running (pid $(cat "$PIDFILE"))"; return; fi
+  nohup python -m zico.server > "$LOG" 2>&1 < /dev/null &
+  echo $! > "$PIDFILE"
+  awake -w "$(cat "$PIDFILE")"
+  echo "starting on http://$ZICO_HOST:$ZICO_PORT (pid $(cat "$PIDFILE")); follow with: scripts/serve.sh log"
+}
+
+case "${1:-}" in
+  start) start ;;
+  restart)
+    # Hold the machine awake across the gap: an idle Mac sleeps within seconds of the old assertion ending.
+    awake -t 300
+    if running; then kill "$(cat "$PIDFILE")"; rm -f "$PIDFILE"; sleep 2; fi
+    start ;;
+  keep-awake)  # re-assert for a server that is already running (e.g. after the Mac slept anyway)
+    if running; then command -v caffeinate >/dev/null 2>&1 && caffeinate -u -t 5 || true; awake -w "$(cat "$PIDFILE")"; echo "asserted for pid $(cat "$PIDFILE")"; else echo "not running"; fi ;;
+  stop)
+    if running; then kill "$(cat "$PIDFILE")"; rm -f "$PIDFILE"; echo stopped; else echo "not running"; fi ;;
+  status)
+    if running; then echo "running (pid $(cat "$PIDFILE")) on http://$ZICO_HOST:$ZICO_PORT"; else echo "not running"; fi ;;
+  log) tail -n 30 "$LOG" ;;
+  key-path) echo "${ZICO_KEY_FILE:-$HOME/.config/zico/api_key}" ;;
+  *) echo "usage: $0 {start|stop|restart|keep-awake|status|log|key-path}"; exit 1 ;;
+esac

scripts/setup_box.sh

@@ -0,0 +1,56 @@
+#!/usr/bin/env bash
+# One-shot environment setup for a rented CUDA box (Vast.ai PyTorch "devel" template or similar).
+# Usage:  bash scripts/setup_box.sh
+# Puts nothing secret on the machine: only public models and datasets are downloaded.
+set -uo pipefail
+cd "$(dirname "$0")/.."
+
+echo "=== hardware"
+nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader
+echo "cpus: $(nproc) | ram: $(free -g | awk '/^Mem:/{print $2}') GB | disk free here: $(df -h . | awk 'NR==2{print $4}')"
+
+echo "=== python env (reuses the image's CUDA-matched torch if there is one)"
+if [ -f /venv/main/bin/activate ]; then
+  # Vast PyTorch images ship torch in /venv/main
+  # shellcheck disable=SC1091
+  source /venv/main/bin/activate
+else
+  python3 -m venv --system-site-packages .venv
+  # shellcheck disable=SC1091
+  source .venv/bin/activate
+fi
+python -m pip --version >/dev/null 2>&1 || python -m ensurepip -q
+pip install -q --upgrade pip
+python -c "import torch" 2>/dev/null || pip install -q torch
+export MAX_JOBS=${MAX_JOBS:-32}
+pip install -q "transformers>=5.2" datasets peft safetensors numpy pytest bitsandbytes kernels
+pip install -q -e .
+
+echo "=== fused kernels for the Gated DeltaNet layers (the slow PyTorch fallback is 4x+ slower)"
+pip install -q flash-linear-attention || echo "WARN: flash-linear-attention did not install"
+pip install -q causal-conv1d --no-build-isolation || echo "WARN: causal-conv1d did not install (conv falls back to PyTorch)"
+# fla's DeltaNet backward is numerically WRONG on Hopper GPUs with Triton 3.4-3.7.0 (fla refuses to run it);
+# the tilelang backend is the fix that leaves PyTorch's pinned Triton alone. Verified below, not assumed.
+pip install -q tilelang || echo "WARN: tilelang did not install"
+pip install -q pillow torchvision fastapi uvicorn httpx || true
+
+echo "=== sanity"
+python - <<'PY'
+import torch, transformers
+print("torch", torch.__version__, "| cuda", torch.version.cuda, "| transformers", transformers.__version__)
+assert torch.cuda.is_available(), "CUDA not available"
+print("gpu:", torch.cuda.get_device_name(0), f"{torch.cuda.get_device_properties(0).total_memory/1e9:.0f} GB",
+      "| bf16:", torch.cuda.is_bf16_supported())
+for mod in ("bitsandbytes", "fla", "causal_conv1d", "kernels"):
+    try:
+        __import__(mod); print(f"  {mod}: ok")
+    except Exception as e:  # noqa: BLE001
+        print(f"  {mod}: MISSING ({type(e).__name__})")
+from transformers import AutoTokenizer
+for m in ("Qwen/Qwen3-0.6B", "Qwen/Qwen3.5-0.8B-Base"):
+    AutoTokenizer.from_pretrained(m)  # tokenizers only, so the unit tests can run
+PY
+HF_HUB_OFFLINE=1 python -m pytest -q 2>&1 | tail -3
+echo "=== fused-kernel gradients vs an independent float64 reference"
+python scripts/check_kernels.py 2>&1 | grep -E "rel-err|KERNEL CHECK|FAILED"
+echo "=== setup finished"

scripts/train.py

@@ -0,0 +1,279 @@
+"""Fine-tune a pretrained decoder into a typed-decision model.
+
+    python scripts/train.py --model Qwen/Qwen3.5-4B-Base --data data_a --out runs/qwen35-4b --grad-checkpointing
+    python scripts/train.py ... --resume          # continue an interrupted run from <out>/last
+
+Stage 1 (default) is full-feedback training with proper scoring rules. Stage 2 is the
+speculative bandit-feedback "RLCD" objective: continue from a checkpoint with
+    --init runs/... --objective bandit
+"""
+import argparse
+import json
+import math
+import random
+import time
+from pathlib import Path
+
+import torch
+from transformers import AutoTokenizer
+
+from zico import DecisionModel, Renderer
+from zico.backbone import load_text_backbone, pick_device
+from zico.batching import bucketed_batches, collate
+from zico.data import read_jsonl, render_records, subsample
+from zico.evaluate import format_report, predict, report, strip_tables
+from zico.losses import bandit_proper_score_loss, proper_score_loss
+
+p = argparse.ArgumentParser()
+p.add_argument("--model", default="Qwen/Qwen3.5-0.8B-Base")
+p.add_argument("--init", default=None, help="start from a saved zico checkpoint instead of --model")
+p.add_argument("--resume", action="store_true", help="continue from <out>/last (weights, optimizer, data position)")
+p.add_argument("--data", default="data")
+p.add_argument("--out", required=True)
+p.add_argument("--epochs", type=float, default=1.0)
+p.add_argument("--max-tokens", type=int, default=4096, help="padded tokens per micro-batch")
+p.add_argument("--max-rows", type=int, default=32)
+p.add_argument("--grad-accum", type=int, default=4)
+p.add_argument("--max-len", type=int, default=768, help="rows longer than this are dropped")
+p.add_argument("--max-state-tokens", type=int, default=512)
+p.add_argument("--lr", type=float, default=2e-5, help="backbone lr (use ~2e-4 with --lora-rank)")
+p.add_argument("--head-lr", type=float, default=1e-3)
+p.add_argument("--weight-decay", type=float, default=0.01)
+p.add_argument("--warmup", type=float, default=0.05)
+p.add_argument("--optim", default="auto", choices=["auto", "adamw", "adamw8bit"], help="auto = 8-bit AdamW on CUDA if available")
+p.add_argument("--lora-rank", type=int, default=0, help="0 = full fine-tuning")
+p.add_argument("--train-embeddings", action="store_true")
+p.add_argument("--grad-checkpointing", action="store_true")
+p.add_argument("--precision", default="bf16", choices=["bf16", "fp32"], help="bf16 = autocast over fp32 weights")
+p.add_argument("--objective", default="supervised", choices=["supervised", "bandit"])
+p.add_argument("--brier-weight", type=float, default=0.5)
+p.add_argument("--rps-weight", type=float, default=1.0)
+p.add_argument("--bandit-eps", type=float, default=0.1)
+p.add_argument("--train-rows", type=int, default=None)
+p.add_argument("--val-rows", type=int, default=1200)
+p.add_argument("--eval-every", type=int, default=200, help="optimizer steps; also writes the resumable <out>/last")
+p.add_argument("--log-every", type=int, default=10)
+p.add_argument("--no-save", action="store_true", help="skip checkpoints (speed / memory probing)")
+p.add_argument("--seed", type=int, default=0)
+args = p.parse_args()
+
+torch.manual_seed(args.seed)
+dev = pick_device()
+out = Path(args.out)
+out.mkdir(parents=True, exist_ok=True)
+(out / "args.json").write_text(json.dumps(vars(args), indent=2))
+log_f = open(out / "log.jsonl", "a")
+
+
+def log(**kw):
+    log_f.write(json.dumps(kw) + "\n")
+    log_f.flush()
+
+
+# ------------------------------------------------------------------ model
+state_path = out / "last" / "trainer_state.pt"
+resuming = args.resume and state_path.exists()
+if args.resume and not resuming:
+    print(f"--resume given but {state_path} not found; starting fresh")
+lora_attached = False
+if resuming or args.init:
+    src = out / "last" if resuming else Path(args.init)
+    model, renderer, meta = DecisionModel.load(src, merge_lora=not resuming, lora_trainable=resuming)
+    lora_attached = resuming and meta["lora"]
+    if not resuming:
+        with torch.no_grad():
+            model.log_temp.zero_()  # post-hoc temperatures belong to the old weights; refit after training
+    print(f"{'resuming' if resuming else 'starting'} from {src} (base {meta['base_model']})")
+else:
+    tok = AutoTokenizer.from_pretrained(args.model)
+    renderer = Renderer(tok, max_state_tokens=args.max_state_tokens)
+    model = DecisionModel(load_text_backbone(args.model, torch.float32), base_model=args.model)
+
+if args.lora_rank and not lora_attached:
+    from peft import LoraConfig, get_peft_model
+
+    model.backbone = get_peft_model(
+        model.backbone, LoraConfig(r=args.lora_rank, lora_alpha=2 * args.lora_rank, target_modules="all-linear", lora_dropout=0.0)
+    )
+if not args.train_embeddings:
+    model.backbone.get_input_embeddings().weight.requires_grad_(False)
+if args.grad_checkpointing:
+    model.backbone.gradient_checkpointing_enable()
+    model.backbone.enable_input_require_grads()
+model.autocast_dtype = torch.bfloat16 if args.precision == "bf16" else None
+model.to(dev).train()
+if model.image_pre is not None:  # frozen ViT: never in the optimizer, runs without gradients, bf16 on GPUs
+    model.vision.model.to(device=dev, dtype=torch.bfloat16 if dev != "cpu" and args.precision == "bf16" else torch.float32)
+    print("vision tower attached (frozen): image rows are supported")
+
+backbone_params = [q for q in model.backbone.parameters() if q.requires_grad]
+head_params = model.head_parameters()
+groups = [
+    {"params": backbone_params, "lr": args.lr, "weight_decay": args.weight_decay},
+    {"params": head_params, "lr": args.head_lr, "weight_decay": 0.0},
+]
+opt, opt_name = None, "torch AdamW"
+if args.optim == "adamw8bit" or (args.optim == "auto" and dev == "cuda"):
+    try:
+        import bitsandbytes as bnb
+
+        opt, opt_name = bnb.optim.AdamW8bit(groups, betas=(0.9, 0.98)), "bitsandbytes AdamW8bit"
+    except ImportError:
+        print("bitsandbytes not installed; falling back to torch AdamW (needs far more memory)")
+if opt is None:
+    opt = torch.optim.AdamW(groups, betas=(0.9, 0.98))
+n_bb, n_head = sum(q.numel() for q in backbone_params), sum(q.numel() for q in head_params)
+print(f"device={dev} precision={args.precision} optimizer={opt_name} trainable: backbone {n_bb/1e6:.1f}M + head {n_head/1e6:.2f}M")
+
+# ------------------------------------------------------------------ data
+t0 = time.time()
+train_recs = subsample(read_jsonl(f"{args.data}/train.jsonl"), args.train_rows, args.seed)
+val_recs = subsample(read_jsonl(f"{args.data}/val.jsonl"), args.val_rows, args.seed)
+image_meta = model.image_pre.meta if model.image_pre is not None else None
+train_rows, train_recs = render_records(train_recs, renderer, args.max_len, image_meta, args.data)
+val_rows, val_recs = render_records(val_recs, renderer, args.max_len, image_meta, args.data)
+lengths = [len(r) for r in train_rows]
+n_img = sum(bool(r.image_paths) for r in train_rows)
+if n_img:
+    print(f"{n_img} of {len(train_rows)} train rows contain images")
+load_image = model.image_pre.load_prepared if model.image_pre is not None else None
+print(f"rendered {len(train_rows)} train / {len(val_rows)} val rows in {time.time()-t0:.0f}s; "
+      f"mean len {sum(lengths)/len(lengths):.0f}, max {max(lengths)}, total {sum(lengths)/1e6:.2f}M tokens/epoch")  # fmt: skip
+
+
+def epoch_batches(epoch: int):
+    """Deterministic per epoch, so a resumed run sees exactly the batches it has not done yet.
+
+    Static (rows, length) shapes: see batching.bucketed_batches for why this matters on GPUs.
+    """
+    return bucketed_batches(lengths, args.max_tokens, args.max_rows, random.Random(args.seed * 1000 + epoch), drop_last=True)
+
+
+batches_per_epoch = len(epoch_batches(0))
+total_steps = max(1, int(batches_per_epoch * args.epochs) // args.grad_accum)
+warmup_steps = max(1, int(args.warmup * total_steps))
+print(f"{batches_per_epoch} micro-batches/epoch -> {total_steps} optimizer steps")
+
+
+def prefetch(items, build, workers: int = 12, ahead: int = 24):
+    """Yield (item[0], item[1], build(item)) in order, building up to `ahead` items concurrently."""
+    import collections
+    from concurrent.futures import ThreadPoolExecutor
+
+    if not any(r.image_paths for r in train_rows):  # text-only data: nothing worth parallelising
+        for it in items:
+            yield it[0], it[1], build(it)
+        return
+    with ThreadPoolExecutor(workers) as pool:
+        queue, it = collections.deque(), iter(items)
+        for item in it:
+            queue.append((item, pool.submit(build, item)))
+            if len(queue) >= ahead:
+                first, fut = queue.popleft()
+                yield first[0], first[1], fut.result()
+        while queue:
+            first, fut = queue.popleft()
+            yield first[0], first[1], fut.result()
+
+
+def lr_scale(step):
+    if step < warmup_steps:
+        return (step + 1) / warmup_steps
+    t = (step - warmup_steps) / max(1, total_steps - warmup_steps)
+    return 0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * t))
+
+
+def run_eval(step):
+    preds = predict(model, val_rows, val_recs, renderer.slot_id)
+    rep = report(preds)
+    print(format_report(rep, f"val @ step {step}"), flush=True)
+    log(kind="eval", step=step, **strip_tables(rep))
+    return rep["overall"]["nll"]
+
+
+def mem() -> str:
+    if dev == "cuda":
+        return f" mem {torch.cuda.max_memory_allocated()/1e9:.0f}G"
+    if dev == "mps":
+        return f" mem {torch.mps.driver_allocated_memory()/1e9:.0f}G"
+    return ""
+
+
+# ------------------------------------------------------------------ loop
+best_nll, step, epoch, batches_done, tokens_seen = float("inf"), 0, 0, 0, 0
+if resuming:
+    st = torch.load(state_path, map_location="cpu", weights_only=False)
+    opt.load_state_dict(st["opt"])
+    best_nll, step, epoch, batches_done, tokens_seen = st["best_nll"], st["step"], st["epoch"], st["batches_done"], st["tokens_seen"]
+    torch.set_rng_state(st["torch_rng"])
+    print(f"resumed at step {step} (epoch {epoch}, {batches_done} micro-batches into it), best val nll {best_nll:.4f}")
+
+
+def snapshot():
+    model.save(out / "last", renderer, {"step": step})
+    torch.save({"opt": opt.state_dict(), "best_nll": best_nll, "step": step, "epoch": epoch, "batches_done": batches_done,
+                "tokens_seen": tokens_seen, "torch_rng": torch.get_rng_state()}, state_path)  # fmt: skip
+
+
+run_loss, run_acc, run_n, micro = 0.0, 0.0, 0, 0
+step_start, tokens_start = step, tokens_seen
+t_start = t_log = time.time()
+tokens_log = 0
+done = step >= total_steps
+while not done:
+    todo = [(bi, idx, pad_len) for bi, (idx, pad_len) in enumerate(epoch_batches(epoch)) if bi >= batches_done]
+    # JPEG decoding + patchifying happens in background threads, a few batches ahead of the GPU
+    for bi, idx, cpu_batch in prefetch(todo, lambda t: collate([train_rows[i] for i in t[1]], renderer.slot_id, t[2], load_image)):
+        b = {k: v.to(dev) for k, v in cpu_batch.items()}
+        logits = model(**b)
+        if args.objective == "supervised":
+            loss = proper_score_loss(logits, b["opt_mask"], b["target"], b["qtype"], args.brier_weight, args.rps_weight)
+        else:
+            loss, _ = bandit_proper_score_loss(logits, b["opt_mask"], b["label"], args.bandit_eps, args.brier_weight)
+        (loss.mean() / args.grad_accum).backward()
+        with torch.no_grad():
+            run_loss += loss.sum().item()
+            run_acc += (logits.argmax(-1) == b["label"]).sum().item()
+            run_n += len(idx)
+        n_tok = int(b["attention_mask"].sum())
+        tokens_seen += n_tok
+        tokens_log += n_tok
+        micro += 1
+        batches_done = bi + 1
+        if micro % args.grad_accum:
+            continue
+
+        scale = lr_scale(step)
+        for g, base in zip(opt.param_groups, (args.lr, args.head_lr)):
+            g["lr"] = base * scale
+        gnorm = torch.nn.utils.clip_grad_norm_(backbone_params + head_params, 1.0).item()
+        opt.step()
+        opt.zero_grad(set_to_none=True)
+        step += 1
+
+        if step % args.log_every == 0:
+            now = time.time()
+            eta = (total_steps - step) * (now - t_start) / max(1, step - step_start)
+            print(f"step {step:5d}/{total_steps} loss {run_loss/run_n:.4f} acc {run_acc/run_n:.3f} gnorm {gnorm:.2f} "
+                  f"lr {args.lr*scale:.2e} {tokens_log/(now-t_log):,.0f} tok/s eta {eta/60:.0f}m{mem()}", flush=True)  # fmt: skip
+            log(kind="train", step=step, loss=run_loss / run_n, acc=run_acc / run_n, gnorm=gnorm, tokens=tokens_seen)
+            run_loss, run_acc, run_n, t_log, tokens_log = 0.0, 0.0, 0, now, 0
+        if dev == "mps" and step % 25 == 0:
+            torch.mps.empty_cache()  # variable shapes make the MPS allocator cache balloon
+        if step % args.eval_every == 0 or step == total_steps:
+            nll = run_eval(step)
+            if nll < best_nll:
+                best_nll = nll
+                if not args.no_save:
+                    model.save(out / "best", renderer, {"step": step, "val_nll": nll})
+                    print(f"  saved best (val nll {nll:.4f})", flush=True)
+            if not args.no_save:
+                snapshot()
+        if step >= total_steps:
+            done = True
+            break
+    else:
+        epoch, batches_done = epoch + 1, 0
+
+trained = tokens_seen - tokens_start
+print(f"done: {step} steps, {trained/1e6:.2f}M tokens this session, {(time.time()-t_start)/60:.1f} min, best val nll {best_nll:.4f}")

tests/test_core.py

@@ -0,0 +1,217 @@
+"""Core invariants, checked on tiny randomly-initialised backbones (CPU, fp32, no training).
+
+Tokenizers are loaded from the local HF cache; tests skip if a tokenizer is unavailable.
+"""
+import math
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from zico import Choice, DecisionModel, Noul, Renderer, Score
+from zico.batching import collate, make_row, token_budget_batches
+from zico.losses import bandit_proper_score_loss, proper_score_loss
+from zico.render import M_OPT
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+STATE = {"ticket": {"id": 4412, "body": "I was charged twice this month and want my money back.", "plan": "pro"}}
+QUESTIONS = {
+    "refund": Noul("The customer is asking for a refund."),
+    "team": Choice(
+        "Which team should handle this?",
+        {"billing": "Charges, invoices, refunds", "technical": "Bugs, outages", "sales": "Upgrades and pricing"},
+    ),
+    "severity": Score("How severe is the issue?", ["Cosmetic", "Degraded but usable", "Blocking"]),
+    "lang": Choice("Which language is the message in?", ["english", "spanish", "german", "french", "other"]),
+}
+
+
+def _tiny(kind):
+    from transformers import AutoTokenizer
+
+    tok_id = {"qwen3": "Qwen/Qwen3-0.6B", "qwen3_5": "Qwen/Qwen3.5-0.8B-Base"}[kind]
+    try:
+        tok = AutoTokenizer.from_pretrained(tok_id, local_files_only=True)
+    except Exception:  # noqa: BLE001
+        pytest.skip(f"tokenizer {tok_id} not in local cache")
+    torch.manual_seed(0)
+    if kind == "qwen3":
+        from transformers import Qwen3Config, Qwen3Model
+
+        cfg = Qwen3Config(
+            vocab_size=len(tok), hidden_size=64, intermediate_size=128, num_hidden_layers=3,
+            num_attention_heads=4, num_key_value_heads=2, head_dim=16, max_position_embeddings=2048,
+        )
+        backbone = Qwen3Model(cfg)
+    else:
+        from transformers.models.qwen3_5 import Qwen3_5TextConfig, Qwen3_5TextModel
+
+        cfg = Qwen3_5TextConfig(
+            vocab_size=len(tok), hidden_size=64, intermediate_size=128, num_hidden_layers=4,
+            num_attention_heads=4, num_key_value_heads=2, head_dim=16, full_attention_interval=4,
+            # 32-dim DeltaNet heads: mlx-lm's Metal kernel cannot be built for smaller ones
+            linear_num_key_heads=4, linear_num_value_heads=4, linear_key_head_dim=32, linear_value_head_dim=32,
+            max_position_embeddings=2048,
+        )
+        backbone = Qwen3_5TextModel(cfg)
+    model = DecisionModel(backbone, d_head=32, base_model=tok_id).eval()
+    with torch.no_grad():  # random head so logits are not trivially equal
+        model.unary.weight.normal_(0, 0.5)
+    # with fused DeltaNet kernels installed, HF routes to CUDA-only code, so test on the GPU if there is one
+    return model.to(DEVICE), Renderer(tok)
+
+
+@pytest.fixture(scope="module", params=["qwen3", "qwen3_5"])
+def tiny(request):
+    return _tiny(request.param)
+
+
+def test_render_slots_and_shared_prefix(tiny):
+    _, r = tiny
+    s1, s2 = r.state(STATE), r.state(STATE)
+    assert s1 == s2  # the state segment never depends on the question
+    for q in QUESTIONS.values():
+        rq = r.question(q)
+        assert len(rq.opt_pos) == len(q.options()) == len(rq.keys)
+        assert all(rq.marker_ids[p] == M_OPT for p in rq.opt_pos)
+        assert rq.marker_ids[rq.ans_pos] > M_OPT and rq.ans_pos == len(rq.input_ids) - 1
+        assert sum(m >= 0 for m in rq.marker_ids) == len(rq.opt_pos) + 1
+
+
+def test_schema_limits():
+    with pytest.raises(ValueError):
+        Choice("x", ["only-one"])
+    with pytest.raises(ValueError):
+        Choice("x", [str(i) for i in range(256)])
+    with pytest.raises(ValueError):
+        Score("x", [str(i) for i in range(11)])
+    assert len(Choice("x", [str(i) for i in range(255)]).options()) == 255
+
+
+def test_fork_equals_rows_equals_single(tiny):
+    """Batched questions over a forked state cache == each question run alone with the state."""
+    model, r = tiny
+    qs = list(QUESTIONS.values())
+    forked, _, _ = model.question_logits(r, STATE, qs, fork=True)
+    rows, _, _ = model.question_logits(r, STATE, qs, fork=False)
+    # On CUDA the cached-continuation path runs fla's Triton kernels, whose tl.dot uses TF32: logits agree to
+    # ~3e-3 (probabilities to ~2e-4) rather than to float32 rounding. A logic error is orders of magnitude larger.
+    atol = 5e-3 if DEVICE == "cuda" else 2e-4
+    for i, q in enumerate(qs):
+        single, _, _ = model.question_logits(r, STATE, [q], fork=False)
+        assert torch.allclose(forked[i], single[0], atol=atol), (i, forked[i], single[0])
+        assert torch.allclose(rows[i], single[0], atol=2e-4), (i, rows[i], single[0])
+
+
+def test_answers_are_type_safe(tiny):
+    model, r = tiny
+    out = model.answer(r, STATE, QUESTIONS)["answers"]
+    assert set(out) == set(QUESTIONS)
+    assert 0.0 <= out["refund"]["noul"] <= 1.0
+    team = out["team"]
+    assert team["choice"] in QUESTIONS["team"].criteria
+    assert set(team["probabilities"]) == set(QUESTIONS["team"].criteria)
+    assert math.isclose(sum(team["probabilities"].values()), 1.0, abs_tol=1e-4)
+    sev = out["severity"]
+    assert 0.0 <= sev["score"] <= 2.0 and set(sev["probabilities"]) == {"0", "1", "2"}
+    assert 0.0 <= sev["confidence"] <= 1.0
+
+
+def test_gradients_reach_backbone_and_head(tiny):
+    model, r = tiny
+    s = r.state(STATE)
+    rows = [make_row(s, r.question(q), label=1) for q in QUESTIONS.values()]
+    b = {k: v.to(DEVICE) for k, v in collate(rows, r.slot_id).items()}
+    loss = proper_score_loss(model(**b), b["opt_mask"], b["target"], b["qtype"], 0.5, 0.5).mean()
+    loss.backward()
+    assert model.marker_emb.weight.grad.abs().sum() > 0
+    assert model.q_proj.weight.grad.abs().sum() > 0
+    assert any(p.grad is not None and p.grad.abs().sum() > 0 for p in model.backbone.layers[0].parameters())
+    model.zero_grad()
+
+
+def test_log_score_is_cross_entropy():
+    torch.manual_seed(0)
+    logits = torch.randn(6, 5)
+    mask = torch.ones(6, 5, dtype=torch.bool)
+    label = torch.randint(0, 5, (6,))
+    target = F.one_hot(label, 5).float()
+    ours = proper_score_loss(logits, mask, target, torch.zeros(6, dtype=torch.long))
+    assert torch.allclose(ours, F.cross_entropy(logits, label, reduction="none"), atol=1e-6)
+
+
+def test_rps_is_ordinal():
+    """Putting mass on a neighbouring level must cost less than putting it on a distant one."""
+    mask = torch.ones(2, 5, dtype=torch.bool)
+    target = F.one_hot(torch.tensor([0, 0]), 5).float()
+    logits = torch.full((2, 5), -20.0)
+    logits[0, 1] = 20.0  # near miss
+    logits[1, 4] = 20.0  # far miss
+    qtype = torch.ones(2, dtype=torch.long)  # score
+    ce_only = proper_score_loss(logits, mask, target, qtype)
+    with_rps = proper_score_loss(logits, mask, target, qtype, rps_weight=1.0)
+    rps = with_rps - ce_only
+    assert rps[0] < rps[1]
+    assert torch.allclose(rps, torch.tensor([0.25, 1.0]), atol=1e-4)
+
+
+def test_bandit_estimator_is_unbiased():
+    """E[bandit loss] == -log p_y, using only the right/wrong bit of a sampled action."""
+    g = torch.Generator().manual_seed(0)
+    logits = torch.tensor([[1.0, 0.2, -0.5, 0.0]]).repeat(200_000, 1)
+    mask = torch.ones_like(logits, dtype=torch.bool)
+    label = torch.full((logits.shape[0],), 2)
+    loss, _ = bandit_proper_score_loss(logits, mask, label, eps=0.2, clip=1e9, generator=g)
+    expected = -F.log_softmax(logits[0], -1)[2].item()
+    assert abs(loss.mean().item() - expected) < 0.03, (loss.mean().item(), expected)
+
+
+def test_temperature_scaling_recovers_known_miscalibration():
+    """Labels drawn from softmax(z); an overconfident model reporting 3z must get T ~= 3."""
+    from zico.calibration import fit_temperature
+
+    g = torch.Generator().manual_seed(0)
+    z = torch.randn(4000, 6, generator=g) * 1.5
+    labels = torch.multinomial(torch.softmax(z, -1), 1, generator=g).squeeze(-1).tolist()
+    assert abs(fit_temperature(list(3 * z), labels) - 3.0) < 0.25
+    assert abs(fit_temperature(list(z), labels) - 1.0) < 0.1
+    # ragged option counts and a hopeless model stay inside the bounds instead of diverging
+    ragged = [torch.randn(int(k), generator=g) * 50 for k in torch.randint(2, 9, (300,), generator=g)]
+    t = fit_temperature(ragged, [0] * 300)
+    assert 0.05 <= t <= 20.0
+
+
+def test_token_budget_batches_cover_everything():
+    import random
+
+    lengths = [random.Random(1).randint(10, 400) for _ in range(500)]
+    batches = token_budget_batches(lengths, max_tokens=2048, max_rows=16, shuffle_rng=random.Random(0))
+    assert sorted(i for b in batches for i in b) == list(range(500))
+    assert all(max(lengths[i] for i in b) * len(b) <= 2048 and len(b) <= 16 for b in batches)
+
+
+def test_bucketed_batches_use_few_static_shapes():
+    import random
+
+    from zico.batching import bucketed_batches
+
+    r = random.Random(1)
+    n = 60_000
+    lengths = [r.randint(10, 1000) for _ in range(n)]
+    batches = bucketed_batches(lengths, max_tokens=16384, max_rows=256, shuffle_rng=random.Random(0))
+    assert sorted(i for idx, _ in batches for i in idx) == list(range(n))  # nothing lost without drop_last
+    assert all(max(lengths[i] for i in idx) <= L and len(idx) * L <= 16384 for idx, L in batches)
+    static = bucketed_batches(lengths, 16384, 256, random.Random(0), drop_last=True)
+    assert len({(len(idx), L) for idx, L in static}) <= 9  # a handful of shapes, tuned once each
+    assert sum(len(idx) for idx, _ in static) > 0.98 * n  # only one ragged batch per bucket is dropped
+
+
+def test_save_load_roundtrip(tiny, tmp_path):
+    model, r = tiny
+    qs = list(QUESTIONS.values())
+    before, _, _ = model.question_logits(r, STATE, qs)
+    model.save(tmp_path / "ckpt", r)
+    loaded, r2, _ = DecisionModel.load(tmp_path / "ckpt", device=DEVICE)
+    after, _, _ = loaded.question_logits(r2, STATE, qs)
+    for a, b in zip(before, after):
+        assert torch.allclose(a, b, atol=1e-5)

tests/test_mlx.py

@@ -0,0 +1,80 @@
+"""MLX backend parity on a tiny random Qwen3.5 (Apple silicon only; skipped elsewhere).
+
+Guards the weight conversion (conv layout, the +1 shift on zero-centred norms, key nesting), the
+slot-embedding injection, the cache fork for both layer types, and the head port.
+"""
+import pytest
+
+pytest.importorskip("mlx_lm")
+
+import torch  # noqa: E402
+
+from zico.mlx_backend import MLXDecisionModel  # noqa: E402
+from tests.test_core import QUESTIONS, STATE, _tiny  # noqa: E402
+from tests.test_vision import tiny_mm  # noqa: E402, F401  (pytest fixture)
+
+
+@pytest.fixture(scope="module")
+def pair(tmp_path_factory):
+    model, renderer = _tiny("qwen3_5")
+    with torch.no_grad():  # make every ported parameter matter
+        model.readout_norm.weight.normal_(1, 0.2)
+        model.readout_norm.bias.normal_(0, 0.2)
+        model.log_temp.copy_(torch.tensor([0.1, -0.2, 0.3]))
+    path = tmp_path_factory.mktemp("ckpt") / "tiny"
+    model.save(path, renderer)
+    mlx_model, mlx_renderer, _ = MLXDecisionModel.load(path)
+    return model, renderer, mlx_model, mlx_renderer
+
+
+def test_mlx_matches_pytorch(pair):
+    model, renderer, mlx_model, mlx_renderer = pair
+    a = model.answer(renderer, STATE, QUESTIONS)
+    b = mlx_model.answer(mlx_renderer, STATE, QUESTIONS)
+    assert a["usage"] == b["usage"]
+    for name in QUESTIONS:
+        x, y = a["answers"][name], b["answers"][name]
+        assert x["type"] == y["type"]
+        # fp32 on both sides: agreement should be at rounding level. A looser tolerance would have hidden
+        # the q/k epsilon mismatch that _reference_qk_norm corrects (it showed up as 6e-2 here).
+        if x["type"] == "noul":
+            assert abs(x["noul"] - y["noul"]) < 1e-4
+        else:
+            assert all(abs(x["probabilities"][k] - y["probabilities"][k]) < 1e-4 for k in x["probabilities"])
+
+
+def test_mlx_image_path_matches_pytorch(tiny_mm, tmp_path):
+    """Explicit 3-part rope, ViT features scattered into the state, and the fork after an image."""
+    from tests.test_vision import QUESTIONS as IMAGE_QUESTIONS
+    from tests.test_vision import STATE as IMAGE_STATE
+
+    _, model, renderer = tiny_mm
+    path = tmp_path / "ckpt"
+    model.save(path, renderer)  # a checkpoint of a vision-enabled model carries its (frozen) ViT along
+    assert (path / "vision" / "model.safetensors").exists()
+    mlx_model, mlx_renderer, _ = MLXDecisionModel.load(path)
+    mlx_model.attach_vision(path / "vision", device="cpu", dtype=torch.float32, min_pixels=32 * 32, max_pixels=128 * 128)
+    qs = {f"q{i}": q for i, q in enumerate(IMAGE_QUESTIONS)}
+    a = model.answer(renderer, IMAGE_STATE, qs)
+    b = mlx_model.answer(mlx_renderer, IMAGE_STATE, qs)
+    assert a["usage"] == b["usage"]
+    for name in qs:
+        x, y = a["answers"][name], b["answers"][name]
+        pa = x["probabilities"] if "probabilities" in x else {"true": x["noul"]}
+        pb = y["probabilities"] if "probabilities" in y else {"true": y["noul"]}
+        assert all(abs(pa[k] - pb[k]) < 1e-4 for k in pa), (name, pa, pb)
+    # a different picture under identical token ids must not hit the cached prefix
+    other = {**IMAGE_STATE, "second": {"type": "image", "data": __import__("tests.test_vision", fromlist=["png"]).png(64, 128, seed=9)}}
+    assert mlx_model.answer(mlx_renderer, other, qs) != b
+
+
+def test_state_prefix_is_reused_and_not_corrupted(pair):
+    _, _, mlx_model, mlx_renderer = pair
+    mlx_model._prefixes.clear()
+    first = mlx_model.answer(mlx_renderer, STATE, QUESTIONS)
+    assert len(mlx_model._prefixes) == 1
+    again = mlx_model.answer(mlx_renderer, STATE, QUESTIONS)  # served from the cached prefix
+    assert len(mlx_model._prefixes) == 1 and first == again
+    for i in range(mlx_model.prefix_cache_size + 2):  # LRU stays bounded
+        mlx_model.answer(mlx_renderer, {"n": i}, {"q": QUESTIONS["refund"]})
+    assert len(mlx_model._prefixes) == mlx_model.prefix_cache_size

tests/test_server.py

@@ -0,0 +1,122 @@
+"""The HTTP contract, exercised against a tiny random model (no real checkpoint needed)."""
+import math
+
+import pytest
+
+pytest.importorskip("fastapi")
+pytest.importorskip("httpx")
+from fastapi.testclient import TestClient  # noqa: E402
+
+from zico.server import create_app  # noqa: E402
+from tests.test_core import _tiny  # noqa: E402
+
+KEY = "test-key"
+AUTH = {"Authorization": f"Bearer {KEY}"}
+REQUEST = {
+    "model": "zico-latest",
+    "state": {"ticket": {"body": "I was charged twice and want my money back."}},
+    "questions": {
+        "refund": {"type": "noul", "instructions": "The customer asks for a refund."},
+        "team": {"type": "choice", "instructions": "Which team?", "criteria": {"billing": "Charges, refunds", "technical": None}},
+        "severity": {"type": "score", "instructions": "How severe?", "criteria": ["Cosmetic", "Degraded", "Blocking"]},
+    },
+}
+
+
+@pytest.fixture(scope="module")
+def client():
+    model, renderer = _tiny("qwen3")
+    return TestClient(create_app(model, renderer, {KEY}, "tiny-test-model"))
+
+
+def test_requires_bearer_key(client):
+    assert client.post("/v1/systemone", json=REQUEST).status_code == 401
+    assert client.post("/v1/systemone", json=REQUEST, headers={"Authorization": "Bearer wrong"}).status_code == 401
+    assert client.get("/healthz").status_code == 200  # health check is open
+
+
+def test_response_matches_the_documented_shape(client):
+    r = client.post("/v1/systemone", json=REQUEST, headers=AUTH)
+    assert r.status_code == 200
+    assert r.headers["x-typesafe-request-id"] == r.headers["x-request-id"]
+    body = r.json()
+    assert set(body) == {"model", "answers", "usage"}
+    assert body["model"] == "tiny-test-model"  # always the real model, never the alias that was requested
+    assert body["usage"]["output_tokens"] == 0 and body["usage"]["input_tokens"] > 0
+    a = body["answers"]
+    assert set(a) == set(REQUEST["questions"])
+    assert a["refund"]["type"] == "noul" and 0 <= a["refund"]["noul"] <= 1 and set(a["refund"]) == {"type", "noul"}
+    team = a["team"]
+    assert set(team) == {"type", "choice", "probabilities", "confidence"} and team["choice"] in ("billing", "technical")
+    assert set(team["probabilities"]) == {"billing", "technical"} and math.isclose(sum(team["probabilities"].values()), 1, abs_tol=1e-4)
+    sev = a["severity"]
+    assert set(sev) == {"type", "score", "legend", "probabilities", "confidence"}
+    assert sev["legend"] == {"0": "Cosmetic", "1": "Degraded", "2": "Blocking"} and 0 <= sev["score"] <= 2
+
+
+@pytest.mark.parametrize("mutate, status", [
+    (lambda b: b.pop("state"), 422),
+    (lambda b: b.update(model="not-a-model"), 404),
+    (lambda b: b.update(questions={}), 422),
+    (lambda b: b["questions"]["team"].update(criteria={"only": "one"}), 422),
+    (lambda b: b["questions"]["severity"].update(criteria=[str(i) for i in range(11)]), 422),
+    (lambda b: b["questions"]["refund"].update(type="essay"), 422),
+    (lambda b: b["questions"]["refund"].update(criteria={"maybe": "x"}), 422),
+])  # fmt: skip
+def test_validation_errors(client, mutate, status):
+    import copy
+
+    body = copy.deepcopy(REQUEST)
+    mutate(body)
+    r = client.post("/v1/systemone", json=body, headers=AUTH)
+    assert r.status_code == status
+    detail = r.json()["detail"]
+    if status == 422:  # published HTTPValidationError shape
+        assert set(detail[0]) >= {"loc", "msg", "type"} and detail[0]["loc"][0] == "body"
+    else:
+        assert isinstance(detail, str)
+
+
+def test_instructions_are_optional(client):
+    body = {**REQUEST, "questions": {"sev": {"type": "score", "criteria": ["Low", "High"]}}}
+    assert client.post("/v1/systemone", json=body, headers=AUTH).status_code == 200
+
+
+def test_models_listing(client):
+    cards = client.get("/v1/models", headers=AUTH).json()["models"]
+    assert cards[0]["name"] == "tiny-test-model" and set(cards[0]) == {"name", "description", "release_date"}
+
+
+def test_test_page_is_public_but_the_api_stays_locked(client):
+    for path in ("/", "/ui"):
+        r = client.get(path)  # no key: the page holds no secrets, the key is typed into it
+        assert r.status_code == 200 and r.headers["content-type"].startswith("text/html") and "zico test bench" in r.text
+        csp = r.headers["content-security-policy"]
+        assert "connect-src 'self'" in csp and "default-src 'none'" in csp and "frame-ancestors 'none'" in csp  # it may only talk to this server
+    assert client.get("/v1/models").status_code == 401
+    assert client.post("/v1/systemone", json=REQUEST).status_code == 401
+
+
+def test_test_page_loads_nothing_from_elsewhere():
+    from zico.server import UI_FILE
+
+    html = UI_FILE.read_text(encoding="utf-8")
+    assert "innerHTML" not in html  # model output and option names only ever go in as text
+    for needle in ('src="http', "src='http", 'href="http', "@import", "url(http"):
+        assert needle not in html
+
+
+def test_test_page_can_be_switched_off():
+    model, renderer = _tiny("qwen3")
+    off = TestClient(create_app(model, renderer, {KEY}, "tiny-test-model", ui=False))
+    assert off.get("/").status_code == 401 and off.get("/", headers=AUTH).status_code == 404
+    assert off.get("/healthz").status_code == 200
+
+
+def test_extra_model_names_are_configuration():
+    model, renderer = _tiny("qwen3")
+    app = TestClient(create_app(model, renderer, {KEY}, "tiny-test-model", aliases=("the-name-my-client-sends",)))
+    ok = {**REQUEST, "model": "the-name-my-client-sends"}
+    assert app.post("/v1/systemone", json=ok, headers=AUTH).status_code == 200
+    assert app.post("/v1/systemone", json={**REQUEST, "model": "tiny-test-model"}, headers=AUTH).status_code == 200  # the real name always works
+    assert app.post("/v1/systemone", json=REQUEST, headers=AUTH).status_code == 404  # default aliases were replaced

tests/test_vision.py

@@ -0,0 +1,248 @@
+"""Images in the state: parsing, positions, and equivalence with Hugging Face's multimodal forward.
+
+Uses a tiny random multimodal Qwen3.5 (vision tower + hybrid language model), CPU / fp32.
+"""
+import base64
+import io
+import json
+
+import numpy as np
+import pytest
+import torch
+
+from zico import Choice, DecisionModel, Noul, Renderer
+from zico.vision import ImageError, mrope_positions, render_image_state, smart_resize, split_state
+from tests.test_core import DEVICE
+
+
+def png(width=96, height=64, seed=0) -> str:
+    from PIL import Image
+
+    rng = np.random.default_rng(seed)
+    buf = io.BytesIO()
+    Image.fromarray(rng.integers(0, 255, (height, width, 3), dtype=np.uint8)).save(buf, format="PNG")
+    return base64.b64encode(buf.getvalue()).decode()
+
+
+# --------------------------------------------------------------------------- parsing
+def test_split_state_finds_images_anywhere_and_in_three_dialects():
+    a, b, c = png(seed=1), png(seed=2), png(seed=3)
+    state = {
+        "listing": {"title": "Blue mug", "photo": {"type": "image", "data": a}},
+        "extra": [{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": b}},
+                  {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{c}"}}],
+        "note": {"type": "image"},  # no payload: ordinary data, not an image part
+    }  # fmt: skip
+    segments, images = split_state(state)
+    assert [base64.b64encode(i.raw).decode() for i in images] == [a, b, c]
+    assert [s for s in segments if isinstance(s, int)] == [0, 1, 2]
+    text = "".join(s if isinstance(s, str) else "<IMG>" for s in segments)
+    assert text == '{"listing": {"title": "Blue mug", "photo": <IMG>}, "extra": [<IMG>, <IMG>], "note": {"type": "image"}}'
+    assert a not in text  # the base64 payload never reaches the tokenizer
+
+
+def test_split_state_without_images_and_bad_images():
+    assert split_state({"a": 1, "b": ["x", {"type": "text"}]}) == ([], [])
+    assert split_state("plain text") == ([], [])
+    with pytest.raises(ImageError):
+        split_state({"type": "image", "data": "not base64!!"})
+    with pytest.raises(ImageError):
+        split_state({"type": "image", "url": "https://example.com/cat.png"})  # remote URLs are not fetched
+    with pytest.raises(ImageError):
+        split_state([{"type": "image", "data": png()}] * 9)
+
+
+def test_smart_resize_respects_budget_factor_and_aspect():
+    for w, h in [(4000, 3000), (640, 480), (20, 20), (1000, 50)]:
+        rw, rh = smart_resize(w, h, factor=32, min_pixels=64 * 64, max_pixels=512 * 512)
+        assert rw % 32 == 0 and rh % 32 == 0 and rw * rh <= 512 * 512
+        assert rw * rh >= 64 * 64 or (w, h) == (1000, 50)
+    rw, rh = smart_resize(4000, 3000, 32, 64 * 64, 512 * 512)
+    assert abs(rw / rh - 4 / 3) < 0.15
+
+
+# --------------------------------------------------------------------------- tiny multimodal model
+@pytest.fixture(scope="module")
+def tiny_mm(tmp_path_factory):
+    from safetensors.torch import save_file
+    from transformers import AutoImageProcessor, AutoTokenizer
+    from transformers.models.qwen3_5 import Qwen3_5Config, Qwen3_5ForConditionalGeneration
+
+    base = "Qwen/Qwen3.5-0.8B-Base"
+    try:
+        tok = AutoTokenizer.from_pretrained(base, local_files_only=True)
+        image_proc = AutoImageProcessor.from_pretrained(base, local_files_only=True)
+        real = __import__("transformers").AutoConfig.from_pretrained(base, local_files_only=True)
+    except Exception:  # noqa: BLE001
+        pytest.skip("base checkpoint files are not in the local cache")
+    torch.manual_seed(0)
+    cfg = Qwen3_5Config(
+        text_config=dict(vocab_size=len(tok), hidden_size=64, intermediate_size=128, num_hidden_layers=4, num_attention_heads=4,
+                         num_key_value_heads=2, head_dim=16, full_attention_interval=4, linear_num_key_heads=4,
+                         linear_num_value_heads=4, linear_key_head_dim=32, linear_value_head_dim=32, max_position_embeddings=2048),
+        vision_config=dict(depth=2, hidden_size=32, intermediate_size=64, num_heads=2, patch_size=16, spatial_merge_size=2,
+                           temporal_patch_size=2, in_channels=3, out_hidden_size=64, num_position_embeddings=256),
+        image_token_id=real.image_token_id, video_token_id=real.video_token_id,
+        vision_start_token_id=real.vision_start_token_id, vision_end_token_id=real.vision_end_token_id,
+    )  # fmt: skip
+    full = Qwen3_5ForConditionalGeneration(cfg).eval().to(DEVICE)
+    vdir = tmp_path_factory.mktemp("vision")
+    save_file({k: v.contiguous().cpu() for k, v in full.model.visual.state_dict().items()}, vdir / "model.safetensors")
+    (vdir / "config.json").write_text(json.dumps({k: v for k, v in cfg.vision_config.to_dict().items() if k not in ("transformers_version", "architectures")}))
+    image_proc.save_pretrained(vdir)
+    (vdir / "zico_vision.json").write_text(json.dumps({
+        "base_model": base, "image_token_id": cfg.image_token_id, "vision_start_token_id": cfg.vision_start_token_id,
+        "vision_end_token_id": cfg.vision_end_token_id, "patch_size": 16, "spatial_merge_size": 2}))  # fmt: skip
+    model = DecisionModel(full.model.language_model, d_head=32, base_model=base).eval().to(DEVICE)
+    with torch.no_grad():
+        model.unary.weight.normal_(0, 0.5)
+    model.attach_vision(vdir, min_pixels=32 * 32, max_pixels=128 * 128)
+    return full, model, Renderer(tok)
+
+
+STATE = {"listing": {"title": "Blue mug", "photo": {"type": "image", "data": png(96, 64)}, "price": 12}, "second": {"type": "image", "data": png(64, 128, seed=5)}}
+QUESTIONS = [Noul("The photo shows a mug."), Choice("What is in the photo?", {"mug": "A cup", "shoe": "Footwear", "lamp": None}),
+             Choice("Pick a colour.", ["red", "green", "blue", "black"])]  # fmt: skip
+
+
+def test_positions_match_hugging_face(tiny_mm):
+    full, model, renderer = tiny_mm
+    st = render_image_state(renderer, STATE, model.image_pre)
+    assert [(i.rows, i.cols) for i in st.images] == [(2, 3), (4, 2)]  # 96x64 and 64x128 at 32 px per token
+    ids = torch.tensor([st.input_ids], device=DEVICE)
+    grids = torch.tensor(np.stack([i.grid_thw for i in st.images]), device=DEVICE)
+    is_image = (ids == model.image_pre.meta["image_token_id"]).int()
+    ref, _ = full.model.get_rope_index(ids, is_image, image_grid_thw=grids, attention_mask=torch.ones_like(ids))
+    assert torch.equal(ref[:3, 0].cpu(), torch.from_numpy(st.positions))
+    assert st.next_pos == int(ref.max()) + 1 < len(st.input_ids)  # positions advance by grid size, not token count
+
+
+def test_image_state_matches_hugging_face_forward(tiny_mm):
+    """Our embedding scatter + positions reproduce HF's own multimodal forward pass."""
+    full, model, renderer = tiny_mm
+    st = render_image_state(renderer, STATE, model.image_pre)
+    ids = torch.tensor([st.input_ids], device=DEVICE)
+    px = torch.from_numpy(np.concatenate([i.pixel_values for i in st.images])).to(DEVICE)
+    grids = torch.tensor(np.stack([i.grid_thw for i in st.images]), device=DEVICE)
+    with torch.no_grad():
+        ref = full.model(input_ids=ids, pixel_values=px, image_grid_thw=grids, attention_mask=torch.ones_like(ids),
+                         mm_token_type_ids=(ids == model.image_pre.meta["image_token_id"]).int()).last_hidden_state  # fmt: skip
+        emb, pos = model._image_state_embeddings(st, torch.device(DEVICE))
+        ours = model.backbone(inputs_embeds=emb, position_ids=pos, attention_mask=torch.ones_like(ids), use_cache=False).last_hidden_state
+    assert torch.allclose(ours, ref, atol=1e-4), (ours - ref).abs().max()
+
+
+def test_fork_equals_single_row_with_images(tiny_mm):
+    _, model, renderer = tiny_mm
+    forked, _, n1 = model.question_logits(renderer, STATE, QUESTIONS, fork=True)
+    rows, _, n2 = model.question_logits(renderer, STATE, QUESTIONS, fork=False)
+    assert n1 == n2
+    for a, b in zip(forked, rows):
+        assert torch.allclose(a, b, atol=2e-4), (a, b)
+
+
+def test_answers_depend_on_the_image_and_text_states_are_untouched(tiny_mm):
+    _, model, renderer = tiny_mm
+    qs = {"q": QUESTIONS[1]}
+    a = model.answer(renderer, {"photo": {"type": "image", "data": png(seed=1)}}, qs)
+    b = model.answer(renderer, {"photo": {"type": "image", "data": png(seed=2)}}, qs)
+    assert a["answers"]["q"]["probabilities"] != b["answers"]["q"]["probabilities"]
+    assert a["usage"]["input_tokens"] < 200  # image tokens are counted; the base64 payload is not
+    text_state = {"note": "no pictures here"}
+    before = model.answer(renderer, text_state, qs)
+    model.vision, keep = None, model.vision  # a text-only model must give the identical answer
+    model.image_pre, keep_pre = None, model.image_pre
+    try:
+        assert model.answer(renderer, text_state, qs) == before
+    finally:
+        model.vision, model.image_pre = keep, keep_pre
+
+
+def test_detail_hint_changes_the_token_budget(tiny_mm):
+    _, model, renderer = tiny_mm
+    big = png(512, 384)
+    n = lambda part: render_image_state(renderer, {"img": part}, model.image_pre).images[0].n_tokens  # noqa: E731
+    default, low, high = n({"type": "image", "data": big}), n({"type": "image", "data": big, "detail": "low"}), \
+        n({"type": "image_url", "image_url": {"url": f"data:image/png;base64,{big}", "detail": "high"}})  # fmt: skip
+    assert default <= 16 < low <= 64 < high <= 1024  # fixture budget 128^2, "low" 256^2, "high" 1024^2 (capped by the image itself)
+    with pytest.raises(ImageError):
+        split_state({"type": "image", "data": big, "detail": "ultra"})
+
+
+def test_server_accepts_and_rejects_images(tiny_mm):
+    pytest.importorskip("fastapi")
+    from fastapi.testclient import TestClient
+
+    from zico.server import create_app
+    from tests.test_core import _tiny
+
+    _, model, renderer = tiny_mm
+    client = TestClient(create_app(model, renderer, {"k"}, "tiny-vision"))
+    auth = {"Authorization": "Bearer k"}
+    body = {"model": "zico-latest", "state": {"photo": {"type": "image", "data": png()}},
+            "questions": {"what": {"type": "choice", "instructions": "What is shown?", "criteria": {"mug": None, "shoe": None}}}}  # fmt: skip
+    r = client.post("/v1/systemone", json=body, headers=auth)
+    assert r.status_code == 200 and r.json()["answers"]["what"]["choice"] in ("mug", "shoe")
+    bad = {**body, "state": {"photo": {"type": "image", "data": "%%%"}}}
+    r = client.post("/v1/systemone", json=bad, headers=auth)
+    assert r.status_code == 422 and r.json()["detail"][0]["type"] == "invalid_image"
+    undecodable = {**body, "state": {"photo": {"type": "image", "data": base64.b64encode(b"not an image").decode()}}}
+    assert client.post("/v1/systemone", json=undecodable, headers=auth).status_code == 422
+    text_model, text_renderer = _tiny("qwen3")
+    text_client = TestClient(create_app(text_model, text_renderer, {"k"}, "tiny-text"))
+    r = text_client.post("/v1/systemone", json=body, headers=auth)
+    assert r.status_code == 422 and r.json()["detail"][0]["type"] == "images_not_supported"
+
+
+def test_training_batch_with_images_matches_inference(tiny_mm, tmp_path):
+    """Image rows + a text row in one padded batch == the single-row inference path; the ViT gets no gradient."""
+    from PIL import Image
+
+    from zico.batching import collate
+    from zico.data import render_records
+    from zico.losses import proper_score_loss
+
+    _, model, renderer = tiny_mm
+    rng = np.random.default_rng(3)
+    files = {"a.png": (96, 64), "b.png": (64, 128)}  # already multiples of 32 and inside the fixture's budget
+    for name, (w, h) in files.items():
+        Image.fromarray(rng.integers(0, 255, (h, w, 3), dtype=np.uint8)).save(tmp_path / name)
+    q0, q1, q2 = ({"type": "noul", "instructions": "The photo shows a mug."},
+                  {"type": "choice", "instructions": "What is in the photo?", "criteria": {"mug": "A cup", "shoe": None, "lamp": None}},
+                  {"type": "choice", "instructions": "Pick a colour.", "criteria": ["red", "green", "blue", "black"]})  # fmt: skip
+    records = [
+        {"source": "t", "sid": "s0", "state": {"listing": {"title": "Blue mug", "photo": {"type": "image", "path": "a.png"}}}, "question": q1, "label": 0, "image_grids": [[2, 3]]},
+        {"source": "t", "sid": "s0", "state": {"listing": {"title": "Blue mug", "photo": {"type": "image", "path": "a.png"}}}, "question": q0, "label": 1, "image_grids": [[2, 3]]},
+        {"source": "t", "sid": "s1", "state": {"note": "no picture here"}, "question": q2, "label": 2},
+        {"source": "t", "sid": "s2", "state": [{"type": "image", "path": "b.png"}, {"type": "image", "path": "a.png"}], "question": q2, "label": 3, "image_grids": [[4, 2], [2, 3]]},
+    ]  # fmt: skip
+    rows, kept = render_records(records, renderer, None, model.image_pre.meta, tmp_path)
+    assert len(rows) == 4 and [len(r.image_paths) for r in rows] == [1, 1, 0, 2]
+    batch = {k: v.to(DEVICE) for k, v in collate(rows, renderer.slot_id, 256, model.image_pre.load_prepared).items()}
+    assert batch["position_ids"].shape == (3, 4, 256) and batch["image_grid_thw"].shape == (4, 3)
+    logits = model(**batch)
+
+    def b64(name):
+        return base64.b64encode((tmp_path / name).read_bytes()).decode()
+
+    api_states = [{"listing": {"title": "Blue mug", "photo": {"type": "image", "data": b64("a.png")}}}] * 2 + \
+        [{"note": "no picture here"}, [{"type": "image", "data": b64("b.png")}, {"type": "image", "data": b64("a.png")}]]  # fmt: skip
+    from zico import question_from_dict
+
+    for i, (state, rec) in enumerate(zip(api_states, records)):
+        ref, _, _ = model.question_logits(renderer, state, [question_from_dict(rec["question"])], fork=False)
+        assert torch.allclose(logits[i, : len(ref[0])], ref[0], atol=2e-4), (i, logits[i], ref[0])
+
+    model.zero_grad()
+    proper_score_loss(logits, batch["opt_mask"], batch["target"], batch["qtype"], 0.5, 1.0).mean().backward()
+    assert any(p.grad is not None and p.grad.abs().sum() > 0 for p in model.backbone.layers[0].parameters())
+    assert all(p.grad is None for p in model.vision.model.parameters())  # the tower is never trained
+    model.zero_grad()
+
+
+def test_mrope_layout_arithmetic():
+    pos, nxt = mrope_positions([("text", 2), ("image", (2, 3)), ("text", 1)])
+    assert pos[:, :2].tolist() == [[0, 1]] * 3
+    assert pos[0, 2:8].tolist() == [2] * 6  # time is constant over an image
+    assert pos[1, 2:8].tolist() == [2, 2, 2, 3, 3, 3] and pos[2, 2:8].tolist() == [2, 3, 4, 2, 3, 4]
+    assert pos[:, 8].tolist() == [5, 5, 5] and nxt == 6  # continues from max(rows, cols), not from 6 image tokens

tokenizer/chat_template.jinja

@@ -0,0 +1,154 @@
+{%- set image_count = namespace(value=0) %}
+{%- set video_count = namespace(value=0) %}
+{%- macro render_content(content, do_vision_count, is_system_content=false) %}
+    {%- if content is string %}
+        {{- content }}
+    {%- elif content is iterable and content is not mapping %}
+        {%- for item in content %}
+            {%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain images.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set image_count.value = image_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Picture ' ~ image_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|image_pad|><|vision_end|>' }}
+            {%- elif 'video' in item or item.type == 'video' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain videos.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set video_count.value = video_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Video ' ~ video_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|video_pad|><|vision_end|>' }}
+            {%- elif 'text' in item %}
+                {{- item.text }}
+            {%- else %}
+                {{- raise_exception('Unexpected item type in content.') }}
+            {%- endif %}
+        {%- endfor %}
+    {%- elif content is none or content is undefined %}
+        {{- '' }}
+    {%- else %}
+        {{- raise_exception('Unexpected content type.') }}
+    {%- endif %}
+{%- endmacro %}
+{%- if not messages %}
+    {{- raise_exception('No messages provided.') }}
+{%- endif %}
+{%- if tools and tools is iterable and tools is not mapping %}
+    {{- '<|im_start|>system\n' }}
+    {{- "# Tools\n\nYou have access to the following functions:\n\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>" }}
+    {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {%- if content %}
+            {{- '\n\n' + content }}
+        {%- endif %}
+    {%- endif %}
+    {{- '<|im_end|>\n' }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {{- '<|im_start|>system\n' + content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" %}
+        {%- set content = render_content(message.content, false)|trim %}
+        {%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}
+            {%- set ns.multi_step_tool = false %}
+            {%- set ns.last_query_index = index %}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if ns.multi_step_tool %}
+    {{- raise_exception('No user query found in messages.') }}
+{%- endif %}
+{%- for message in messages %}
+    {%- set content = render_content(message.content, true)|trim %}
+    {%- if message.role == "system" %}
+        {%- if not loop.first %}
+            {{- raise_exception('System message must be at the beginning.') }}
+        {%- endif %}
+    {%- elif message.role == "user" %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- set reasoning_content = reasoning_content|trim %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\n\n' + content }}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if tool_call.function is defined %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {%- if loop.first %}
+                    {%- if content|trim %}
+                        {{- '\n\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- else %}
+                        {{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- endif %}
+                {%- else %}
+                    {{- '\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                {%- endif %}
+                {%- if tool_call.arguments is defined %}
+                    {%- for args_name, args_value in tool_call.arguments|items %}
+                        {{- '<parameter=' + args_name + '>\n' }}
+                        {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
+                        {{- args_value }}
+                        {{- '\n</parameter>\n' }}
+                    {%- endfor %}
+                {%- endif %}
+                {{- '</function>\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.previtem and loop.previtem.role != "tool" %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if not loop.last and loop.nextitem.role != "tool" %}
+            {{- '<|im_end|>\n' }}
+        {%- elif loop.last %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- else %}
+        {{- raise_exception('Unexpected message role.') }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- else %}
+        {{- '<think>\n' }}
+    {%- endif %}
+{%- endif %}

tokenizer/tokenizer.json

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

tokenizer/tokenizer_config.json

@@ -0,0 +1,32 @@
+{
+  "add_prefix_space": false,
+  "audio_bos_token": "<|audio_start|>",
+  "audio_eos_token": "<|audio_end|>",
+  "audio_token": "<|audio_pad|>",
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "image_token": "<|image_pad|>",
+  "is_local": false,
+  "local_files_only": false,
+  "model_max_length": 262144,
+  "model_specific_special_tokens": {
+    "audio_bos_token": "<|audio_start|>",
+    "audio_eos_token": "<|audio_end|>",
+    "audio_token": "<|audio_pad|>",
+    "image_token": "<|image_pad|>",
+    "video_token": "<|video_pad|>",
+    "vision_bos_token": "<|vision_start|>",
+    "vision_eos_token": "<|vision_end|>"
+  },
+  "pad_token": "<|endoftext|>",
+  "pretokenize_regex": "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?[\\p{L}\\p{M}]+|\\p{N}| ?[^\\s\\p{L}\\p{M}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null,
+  "video_token": "<|video_pad|>",
+  "vision_bos_token": "<|vision_start|>",
+  "vision_eos_token": "<|vision_end|>"
+}

vision/config.json

@@ -0,0 +1,36 @@
+{
+  "output_hidden_states": false,
+  "return_dict": true,
+  "dtype": null,
+  "chunk_size_feed_forward": 0,
+  "is_encoder_decoder": false,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1"
+  },
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1
+  },
+  "problem_type": null,
+  "depth": 24,
+  "hidden_size": 1024,
+  "hidden_act": "gelu_pytorch_tanh",
+  "intermediate_size": 4096,
+  "num_heads": 16,
+  "in_channels": 3,
+  "patch_size": 16,
+  "spatial_merge_size": 2,
+  "temporal_patch_size": 2,
+  "out_hidden_size": 2560,
+  "num_position_embeddings": 2304,
+  "initializer_range": 0.02,
+  "rope_parameters": {
+    "rope_theta": 10000.0,
+    "rope_type": "axial"
+  },
+  "_name_or_path": "",
+  "deepstack_visual_indexes": [],
+  "model_type": "qwen3_5_vision",
+  "output_attentions": false
+}

vision/model.safetensors

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

vision/preprocessor_config.json

@@ -0,0 +1,26 @@
+{
+  "do_convert_rgb": true,
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "image_processor_type": "Qwen2VLImageProcessor",
+  "image_std": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "merge_size": 2,
+  "patch_size": 16,
+  "resample": 3,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "longest_edge": 16777216,
+    "shortest_edge": 65536
+  },
+  "temporal_patch_size": 2
+}

vision/zico_vision.json

@@ -0,0 +1,8 @@
+{
+  "base_model": "Qwen/Qwen3.5-4B-Base",
+  "image_token_id": 248056,
+  "vision_start_token_id": 248053,
+  "vision_end_token_id": 248054,
+  "patch_size": 16,
+  "spatial_merge_size": 2
+}

zico.json

@@ -0,0 +1,12 @@
+{
+  "base_model": "Qwen/Qwen3.5-4B-Base",
+  "d_head": 256,
+  "lora": false,
+  "renderer": {
+    "max_state_tokens": 768,
+    "max_instruction_tokens": 256,
+    "max_option_tokens": 64
+  },
+  "step": 734,
+  "val_nll": 0.27321311831474304
+}

zico/__init__.py

@@ -0,0 +1,13 @@
+from .render import Renderer
+from .schema import Choice, Noul, Score, question_from_dict, question_to_dict
+
+__all__ = ["DecisionModel", "Renderer", "Choice", "Score", "Noul", "question_from_dict", "question_to_dict"]
+
+
+def __getattr__(name):
+    # Lazy, so the MLX serving backend can use the package without importing PyTorch.
+    if name == "DecisionModel":
+        from .model import DecisionModel
+
+        return DecisionModel
+    raise AttributeError(name)

zico/answers.py

@@ -0,0 +1,24 @@
+"""Turn per-question probability vectors into the typed answer objects (framework-free)."""
+from __future__ import annotations
+
+import math
+
+
+def format_answers(names: list[str], probs: list[list[float]], rendered: list, questions: dict, n_tokens: int) -> dict:
+    """rendered[i] is the RenderedQuestion for names[i]; probs[i] its distribution over rendered[i].keys."""
+    answers = {}
+    for name, p, rq in zip(names, probs, rendered):
+        by_key = {k: round(v, 6) for k, v in zip(rq.keys, p)}
+        if rq.qtype == "noul":
+            answers[name] = {"type": "noul", "noul": by_key["true"]}
+            continue
+        entropy = -sum(v * math.log(max(v, 1e-12)) for v in p)
+        conf = round(1.0 - entropy / math.log(len(p)), 6)
+        if rq.qtype == "choice":
+            best = max(range(len(p)), key=p.__getitem__)
+            answers[name] = {"type": "choice", "choice": rq.keys[best], "probabilities": by_key, "confidence": conf}
+        else:
+            score = sum(i * v for i, v in enumerate(p))
+            legend = dict(questions[name].options())
+            answers[name] = {"type": "score", "score": round(score, 4), "legend": legend, "probabilities": by_key, "confidence": conf}
+    return {"answers": answers, "usage": {"input_tokens": n_tokens, "output_tokens": 0}}

zico/backbone.py

@@ -0,0 +1,53 @@
+"""Backbone loading and prefix-cache forking.
+
+Any HF causal decoder stack works: pure softmax attention (Qwen3, Ministral, ...) or
+hybrids with recurrent layers (Qwen3.5 Gated DeltaNet). We never edit attention masks,
+so nothing here depends on the token-mixing layer type.
+"""
+from __future__ import annotations
+
+import copy
+
+import torch
+from transformers import AutoConfig, AutoModel, AutoModelForImageTextToText
+
+
+def pick_device() -> str:
+    if torch.cuda.is_available():
+        return "cuda"
+    if torch.backends.mps.is_available():
+        return "mps"
+    return "cpu"
+
+
+def load_text_backbone(model_id: str, dtype=torch.float32):
+    """Return the text decoder stack only (no LM head, no vision tower)."""
+    cfg = AutoConfig.from_pretrained(model_id)
+    if hasattr(cfg, "text_config") and hasattr(cfg, "vision_config"):
+        full = AutoModelForImageTextToText.from_pretrained(model_id, dtype=dtype)
+        return full.model.language_model
+    return AutoModel.from_pretrained(model_id, dtype=dtype)
+
+
+def fork_cache(cache, n: int):
+    """Copy a batch-1 prefix cache into a batch-n cache.
+
+    Handles attention KV tensors as well as the {state_idx: tensor} dicts that hybrid
+    layers use for conv / recurrent state (HF's batch_repeat_interleave does not).
+    """
+
+    def rep(v):
+        if torch.is_tensor(v) and v.dim() > 0 and v.shape[0] == 1:
+            return v.repeat_interleave(n, dim=0)
+        return v
+
+    forked = copy.deepcopy(cache)
+    for layer in forked.layers:
+        for name, val in list(vars(layer).items()):
+            if torch.is_tensor(val):
+                setattr(layer, name, rep(val))
+            elif isinstance(val, dict):
+                setattr(layer, name, {k: rep(v) for k, v in val.items()})
+            elif isinstance(val, list):
+                setattr(layer, name, [rep(v) for v in val])
+    return forked

zico/batching.py

@@ -0,0 +1,182 @@
+"""Rows and collation. One row = one (state, question) pair, right-padded.
+
+Right padding is deliberate: padded positions come after every readout slot, so with a
+causal backbone (attention or recurrent) they cannot influence any value we read.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+from .render import RenderedQuestion
+from .schema import QTYPE_INDEX
+
+
+@dataclass
+class Row:
+    input_ids: list[int]
+    marker_ids: list[int]
+    opt_pos: list[int]  # absolute positions
+    ans_pos: int
+    qtype: int
+    label: int = -1
+    soft: list[float] | None = None  # optional soft target over options
+    # image rows only: 3-part rotary positions [3, len] and the prepared image files, in order of appearance
+    positions: object = None
+    image_paths: tuple = ()
+
+    def __len__(self):
+        return len(self.input_ids)
+
+
+def make_row(state_ids: list[int], rq: RenderedQuestion, label: int = -1, soft=None) -> Row:
+    s = len(state_ids)
+    return Row(
+        input_ids=state_ids + rq.input_ids,
+        marker_ids=[-1] * s + rq.marker_ids,
+        opt_pos=[s + p for p in rq.opt_pos],
+        ans_pos=s + rq.ans_pos,
+        qtype=QTYPE_INDEX[rq.qtype],
+        label=label,
+        soft=soft,
+    )
+
+
+def make_image_row(st, rq: RenderedQuestion, label: int = -1, soft=None) -> Row:
+    """st: vision.RenderedImageState laid out from recorded grids (its `images` are ImageParts with paths)."""
+    import numpy as np
+
+    s = len(st.input_ids)
+    q_pos = st.next_pos + np.arange(len(rq.input_ids))  # text after an image continues from the grid size
+    return Row(
+        input_ids=st.input_ids + rq.input_ids,
+        marker_ids=[-1] * s + rq.marker_ids,
+        opt_pos=[s + p for p in rq.opt_pos],
+        ans_pos=s + rq.ans_pos,
+        qtype=QTYPE_INDEX[rq.qtype],
+        label=label,
+        soft=soft,
+        positions=np.concatenate([st.positions, np.stack([q_pos] * 3)], axis=1),
+        image_paths=tuple(p.path for p in st.images),
+    )
+
+
+def collate(rows: list[Row], pad_id: int = 0, pad_len: int | None = None, load_image=None) -> dict[str, torch.Tensor]:
+    """load_image(path) -> vision.EncodedImage is required when any row carries images."""
+    batch = _collate_text(rows, pad_id, pad_len)
+    if not any(r.image_paths for r in rows):
+        return batch
+    import numpy as np
+
+    B, L = batch["input_ids"].shape
+    pos = torch.arange(L)[None, None].repeat(3, B, 1)  # text rows (and padding): all three parts advance together
+    pixels, grids = [], []
+    for i, r in enumerate(rows):
+        if r.positions is not None:
+            pos[:, i, : len(r)] = torch.from_numpy(r.positions)
+        for path in r.image_paths:  # row order == the order of <|image_pad|> runs in the flattened batch
+            img = load_image(path)
+            pixels.append(img.pixel_values)
+            grids.append(img.grid_thw)
+    batch["position_ids"] = pos
+    batch["pixel_values"] = torch.from_numpy(np.concatenate(pixels))
+    batch["image_grid_thw"] = torch.from_numpy(np.stack(grids))
+    return batch
+
+
+def _collate_text(rows: list[Row], pad_id: int = 0, pad_len: int | None = None) -> dict[str, torch.Tensor]:
+    B = len(rows)
+    L = max(pad_len or 0, max(len(r) for r in rows))
+    K = max(len(r.opt_pos) for r in rows)
+    input_ids = torch.full((B, L), pad_id, dtype=torch.long)
+    marker_ids = torch.full((B, L), -1, dtype=torch.long)
+    attention_mask = torch.zeros((B, L), dtype=torch.long)
+    opt_pos = torch.zeros((B, K), dtype=torch.long)
+    opt_mask = torch.zeros((B, K), dtype=torch.bool)
+    target = torch.zeros((B, K), dtype=torch.float32)
+    for i, r in enumerate(rows):
+        n, k = len(r), len(r.opt_pos)
+        input_ids[i, :n] = torch.tensor(r.input_ids)
+        marker_ids[i, :n] = torch.tensor(r.marker_ids)
+        attention_mask[i, :n] = 1
+        opt_pos[i, :k] = torch.tensor(r.opt_pos)
+        opt_mask[i, :k] = True
+        if r.soft is not None:
+            target[i, :k] = torch.tensor(r.soft, dtype=torch.float32)
+        elif r.label >= 0:
+            target[i, r.label] = 1.0
+    return {
+        "input_ids": input_ids,
+        "marker_ids": marker_ids,
+        "attention_mask": attention_mask,
+        "opt_pos": opt_pos,
+        "opt_mask": opt_mask,
+        "ans_pos": torch.tensor([r.ans_pos for r in rows], dtype=torch.long),
+        "qtype": torch.tensor([r.qtype for r in rows], dtype=torch.long),
+        "label": torch.tensor([r.label for r in rows], dtype=torch.long),
+        "target": target,  # one-hot or soft distribution over options
+    }
+
+
+LENGTH_BUCKETS = (64, 96, 128, 192, 256, 384, 512, 768, 1024, 1536, 2048, 3072, 4096, 6144, 8192)
+
+
+def bucket_len(n: int) -> int:
+    for b in LENGTH_BUCKETS:
+        if n <= b:
+            return b
+    return -(-n // 2048) * 2048
+
+
+def bucketed_batches(lengths: list[int], max_tokens: int, max_rows: int, shuffle_rng=None, drop_last=False):
+    """Batches with a small, fixed set of (rows, padded length) shapes -> list of (indices, padded length).
+
+    Fused GPU kernels (Triton autotune, fla) re-tune for every new tensor shape, which costs seconds.
+    Free-form token-budget batching makes nearly every batch a new shape and leaves the GPU idle
+    most of the time; a dozen static shapes are tuned once and then run at full speed. The price
+    is some padding up to the bucket boundary. `drop_last` skips the one ragged batch per bucket.
+    """
+    by_bucket: dict[int, list[int]] = {}
+    for i, n in enumerate(lengths):
+        by_bucket.setdefault(bucket_len(n), []).append(i)
+    batches = []
+    for L, idx in sorted(by_bucket.items()):
+        if shuffle_rng is not None:
+            shuffle_rng.shuffle(idx)
+        B = max(1, min(max_rows, max_tokens // L))
+        for s in range(0, len(idx), B):
+            chunk = idx[s : s + B]
+            if len(chunk) == B or not drop_last:
+                batches.append((chunk, L))
+    if shuffle_rng is not None:
+        shuffle_rng.shuffle(batches)
+    return batches
+
+
+def token_budget_batches(lengths: list[int], max_tokens: int, max_rows: int, shuffle_rng=None) -> list[list[int]]:
+    """Group row indices into batches whose padded size stays under max_tokens.
+
+    Rows are bucketed by length (so little compute is wasted on padding) and the batch
+    order is shuffled; within-bucket membership is reshuffled each call via the rng.
+    """
+    idx = list(range(len(lengths)))
+    if shuffle_rng is not None:
+        shuffle_rng.shuffle(idx)
+        chunk = max_rows * 50
+        idx = [j for c in range(0, len(idx), chunk) for j in sorted(idx[c : c + chunk], key=lengths.__getitem__)]
+    else:
+        idx.sort(key=lengths.__getitem__)
+    batches, cur, cur_max = [], [], 0
+    for j in idx:
+        new_max = max(cur_max, lengths[j])
+        if cur and (new_max * (len(cur) + 1) > max_tokens or len(cur) >= max_rows):
+            batches.append(cur)
+            cur, new_max = [], lengths[j]
+        cur.append(j)
+        cur_max = new_max
+    if cur:
+        batches.append(cur)
+    if shuffle_rng is not None:
+        shuffle_rng.shuffle(batches)
+    return batches

zico/benchmark.py

@@ -0,0 +1,51 @@
+"""Shared pieces for the head-to-head benchmarks (same rows, same scoring for every system)."""
+from __future__ import annotations
+
+import statistics
+
+import torch
+
+from .calibration import ece
+from .data import read_jsonl
+
+# (split, source): heldout = zero-shot for every system; test = in-distribution for zico checkpoints
+BENCH_SETS = [("heldout", "massive_intent"), ("heldout", "emotion"), ("heldout", "imdb"), ("heldout", "arc_easy"),
+              ("test", "ag_news"), ("test", "banking77")]  # fmt: skip
+
+
+def load_rows(data_dir: str, split: str, source: str, n: int) -> list[dict]:
+    return [r for r in read_jsonl(f"{data_dir}/{split}.jsonl") if r["source"] == source][:n]
+
+
+def auroc(conf, correct) -> float:
+    """P(confidence of a correct answer > confidence of a wrong one); 0.5 = uninformative."""
+    pos = [c for c, ok in zip(conf, correct) if ok]
+    neg = [c for c, ok in zip(conf, correct) if not ok]
+    if not pos or not neg:
+        return float("nan")
+    return sum((a > b) + 0.5 * (a == b) for a in pos for b in neg) / (len(pos) * len(neg))
+
+
+def summarise(preds, labels, confs, times) -> dict:
+    """preds may contain None (abstention / error), which is scored as wrong."""
+    answered = [p is not None for p in preds]
+    correct = [p == y for p, y in zip(preds, labels)]
+    conf = [c if c is not None else 0.0 for c in confs]
+    n_ans = sum(answered)
+    top = sorted(range(len(conf)), key=lambda i: -conf[i])[: max(1, len(conf) // 2)]
+    return {
+        "n": len(preds),
+        "acc": sum(correct) / len(correct),
+        "abstain": 1 - n_ans / len(preds),
+        "acc_answered": sum(c for c, a in zip(correct, answered) if a) / max(1, n_ans),
+        "ece": ece(torch.tensor(conf), torch.tensor(correct))[0],
+        "auroc": auroc(conf, correct),
+        "acc_top50_conf": sum(correct[i] for i in top) / len(top),
+        "median_ms": statistics.median(times) * 1000 if times else float("nan"),
+    }
+
+
+def format_line(name: str, m: dict) -> str:
+    return (f"  {name:13s} acc {m['acc']:.3f} | abstain {m['abstain']:.3f} | acc(answered) {m['acc_answered']:.3f} | "
+            f"ece {m['ece']:.3f} | auroc {m['auroc']:.3f} | acc@top50%conf {m['acc_top50_conf']:.3f} | "
+            f"median {m['median_ms']:.0f} ms")  # fmt: skip

zico/calibration.py

@@ -0,0 +1,86 @@
+"""Calibration metrics and post-hoc temperature scaling."""
+from __future__ import annotations
+
+import math
+
+import torch
+
+from .schema import QTYPE_INDEX, QTYPES
+
+
+def ece(conf: torch.Tensor, correct: torch.Tensor, n_bins: int = 15):
+    """Expected calibration error of the top-label confidence, plus the reliability table."""
+    edges = torch.linspace(0, 1, n_bins + 1)
+    total, table = 0.0, []
+    for lo, hi in zip(edges[:-1], edges[1:]):
+        m = (conf > lo) & (conf <= hi)
+        if m.any():
+            c, a, w = conf[m].mean().item(), correct[m].float().mean().item(), m.float().mean().item()
+            total += w * abs(c - a)
+            table.append({"lo": round(lo.item(), 3), "hi": round(hi.item(), 3), "conf": c, "acc": a, "frac": w})
+    return total, table
+
+
+def metrics(records: list[dict]) -> dict:
+    """records: {"p": 1-D tensor over options, "label": int, "qtype": int, "target": optional soft tensor}."""
+    if not records:
+        return {}
+    conf = torch.tensor([r["p"].max().item() for r in records])
+    pred = torch.tensor([int(r["p"].argmax()) for r in records])
+    label = torch.tensor([r["label"] for r in records])
+    correct = pred == label
+    nll = torch.tensor([-math.log(max(r["p"][r["label"]].item(), 1e-12)) for r in records])
+    brier = torch.tensor(
+        [((r["p"] - torch.nn.functional.one_hot(torch.tensor(r["label"]), len(r["p"]))) ** 2).sum().item() for r in records]
+    )
+    e, table = ece(conf, correct)
+    out = {
+        "n": len(records),
+        "acc": correct.float().mean().item(),
+        "nll": nll.mean().item(),
+        "brier": brier.mean().item(),
+        "ece": e,
+        "mean_conf": conf.mean().item(),
+        "reliability": table,
+    }
+    score_rows = [r for r in records if r["qtype"] == QTYPE_INDEX["score"]]
+    if score_rows:
+        exp = torch.tensor([(r["p"] * torch.arange(len(r["p"]))).sum().item() for r in score_rows])
+        gold = torch.tensor([float(r["label"]) for r in score_rows])
+        out["score_mae"] = (exp - gold).abs().mean().item()
+    return out
+
+
+def metrics_by(records: list[dict], key: str) -> dict:
+    groups: dict = {}
+    for r in records:
+        groups.setdefault(r[key], []).append(r)
+    return {(QTYPES[k] if key == "qtype" else k): metrics(v) for k, v in sorted(groups.items(), key=lambda kv: str(kv[0]))}
+
+
+def fit_temperature(logits: list[torch.Tensor], labels: list[int], lo: float = 0.05, hi: float = 20.0) -> float:
+    """Single scalar T in [lo, hi] minimising NLL of softmax(logits / T).
+
+    Accuracy is unchanged by construction. The problem is one-dimensional, so we use a
+    bounded grid + golden-section search rather than a gradient method that can diverge.
+    """
+    k = max(len(l) for l in logits)
+    padded = torch.full((len(logits), k), -1e9)
+    for i, l in enumerate(logits):
+        padded[i, : len(l)] = l.detach().float().cpu()
+    y = torch.tensor(labels)
+
+    def nll(log_t: float) -> float:
+        return torch.nn.functional.cross_entropy(padded / math.exp(log_t), y).item()
+
+    grid = torch.linspace(math.log(lo), math.log(hi), 61).tolist()
+    best = min(range(len(grid)), key=lambda i: nll(grid[i]))
+    a, b = grid[max(best - 1, 0)], grid[min(best + 1, len(grid) - 1)]
+    phi = (math.sqrt(5) - 1) / 2
+    for _ in range(40):
+        c, d = b - phi * (b - a), a + phi * (b - a)
+        if nll(c) < nll(d):
+            b = d
+        else:
+            a = c
+    return min(max(math.exp((a + b) / 2), lo), hi)

zico/data.py

@@ -0,0 +1,298 @@
+"""Public datasets recast as (state, typed question, label) records.
+
+The aim is a model that *reads the schema* rather than memorising label sets, so records
+vary what a fixed classifier would hold constant: option order, option subsets, presence
+of descriptions, instruction wording, and whether the state is a string or a JSON object.
+
+Record format (one JSON object per line):
+    {"source", "sid", "state", "question": {"type", "instructions", "criteria"}, "label", "soft"?}
+"""
+from __future__ import annotations
+
+import json
+import random
+from pathlib import Path
+
+from .batching import Row, make_row
+from .render import Renderer
+from .schema import question_from_dict
+
+GO_EMOTIONS = [
+    "admiration", "amusement", "anger", "annoyance", "approval", "caring", "confusion", "curiosity", "desire",
+    "disappointment", "disapproval", "disgust", "embarrassment", "excitement", "fear", "gratitude", "grief", "joy",
+    "love", "nervousness", "optimism", "pride", "realization", "relief", "remorse", "sadness", "surprise", "neutral",
+]  # fmt: skip
+
+
+def _q(qtype, instructions, criteria=None):
+    q = {"type": qtype, "instructions": instructions}
+    if criteria is not None:
+        q["criteria"] = criteria
+    return q
+
+
+def _shuffled_choice(rng, instructions, options: list[tuple[str, str]], gold_key: str, describe=True):
+    """options: [(key, description)]. Returns (question, label index) with a random option order."""
+    options = options[:]
+    rng.shuffle(options)
+    keys = [k for k, _ in options]
+    criteria = {k: d for k, d in options} if describe and any(d for _, d in options) else keys
+    return _q("choice", instructions, criteria), keys.index(gold_key)
+
+
+def _wrap(rng, field, text):
+    """States arrive as plain strings or as JSON objects; train on both."""
+    return {field: text} if rng.random() < 0.6 else text
+
+
+# --------------------------------------------------------------------------- training sources
+def boolq(ex, rng, ctx):
+    question = ex["question"].strip().rstrip("?")
+    question = question[0].upper() + question[1:] + "?"
+    yield _wrap(rng, "passage", ex["passage"]), _q("noul", question), int(ex["answer"]), None
+
+
+_AG = [
+    ("world", "International news, politics, conflicts and diplomacy"),
+    ("sports", "Sports events, teams, athletes and results"),
+    ("business", "Companies, markets, the economy and finance"),
+    ("science_technology", "Science, technology, computing, the internet and gadgets"),
+]
+_AG_T = ["Which news category does this article belong to?", "What is the main topic of the article?",
+         "Route this article to the right news desk."]  # fmt: skip
+
+
+def ag_news(ex, rng, ctx):
+    q, y = _shuffled_choice(rng, rng.choice(_AG_T), _AG, _AG[ex["label"]][0], describe=rng.random() < 0.7)
+    yield _wrap(rng, "article", ex["text"]), q, y, None
+
+
+_SAT5 = ["Very dissatisfied: a bad experience, would warn others", "Dissatisfied: more negatives than positives",
+         "Mixed: an average or so-so experience", "Satisfied: a good experience with minor issues",
+         "Delighted: an excellent experience, highly recommended"]  # fmt: skip
+_SAT3 = ["Negative overall", "Mixed or neutral overall", "Positive overall"]
+
+
+def yelp(ex, rng, ctx):
+    stars, state = ex["label"], _wrap(rng, "review", ex["text"])
+    if rng.random() < 0.6:
+        yield state, _q("score", "How satisfied is the reviewer with the business?", _SAT5), stars, None
+    else:
+        yield state, _q("score", "What is the overall sentiment of the review?", _SAT3), (0, 0, 1, 2, 2)[stars], None
+    if stars != 2 and rng.random() < 0.5:
+        if rng.random() < 0.5:
+            yield state, _q("noul", "The reviewer would recommend this business."), int(stars >= 3), None
+        else:
+            yield state, _q("noul", "The reviewer had a negative experience."), int(stars <= 1), None
+
+
+def go_emotions(ex, rng, ctx):
+    gold = [GO_EMOTIONS[i] for i in ex["labels"]]
+    state = _wrap(rng, "comment", ex["text"])
+    if len(gold) == 1:
+        k = rng.randint(4, 10)
+        distract = rng.sample([e for e in GO_EMOTIONS if e != gold[0]], k - 1)
+        q, y = _shuffled_choice(rng, "Which emotion does the comment primarily express?",
+                                [(e, "") for e in [gold[0]] + distract], gold[0])  # fmt: skip
+        yield state, q, y, None
+    non_neutral = [e for e in gold if e != "neutral"]
+    if non_neutral and rng.random() < 0.5:
+        yield state, _q("noul", f"The author expresses {rng.choice(non_neutral)}."), 1, None
+    else:
+        neg = rng.choice([e for e in GO_EMOTIONS if e not in gold and e != "neutral"])
+        yield state, _q("noul", f"The author expresses {neg}."), 0, None
+
+
+_INTENT_T = ["Which intent best matches the customer's message?", "What is the customer trying to do?",
+             "Classify the support request."]  # fmt: skip
+
+
+def banking77(ex, rng, ctx):
+    names = ctx["banking77_labels"]
+    gold = ex["label_text"]
+    k = rng.choice([16, 32, len(names), len(names)])  # favour the full 77-way list
+    keys = [gold] + rng.sample([n for n in names if n != gold], k - 1)
+    q, y = _shuffled_choice(rng, rng.choice(_INTENT_T), [(n, "") for n in keys], gold)
+    yield _wrap(rng, "message", ex["text"]), q, y, None
+
+
+_NLI = [("entailment", "The hypothesis must be true if the premise is true"),
+        ("neutral", "The hypothesis may or may not be true given the premise"),
+        ("contradiction", "The hypothesis cannot be true if the premise is true")]  # fmt: skip
+
+
+def mnli(ex, rng, ctx):
+    state = {"premise": ex["premise"], "hypothesis": ex["hypothesis"]}
+    if rng.random() < 0.7:
+        q, y = _shuffled_choice(rng, "How does the hypothesis relate to the premise?", _NLI, _NLI[ex["label"]][0],
+                                describe=rng.random() < 0.7)  # fmt: skip
+        yield state, q, y, None
+    else:
+        yield state, _q("noul", "The hypothesis necessarily follows from the premise."), int(ex["label"] == 0), None
+
+
+def race(ex, rng, ctx):
+    opts = list(ex["options"])
+    gold = opts["ABCD".index(ex["answer"])]
+    rng.shuffle(opts)
+    criteria = {"ABCD"[i]: o for i, o in enumerate(opts)}
+    yield {"passage": ex["article"]}, _q("choice", ex["question"], criteria), opts.index(gold), None
+
+
+_TOX = {
+    "toxicity": "This comment is toxic: rude, disrespectful, or likely to make someone leave a discussion.",
+    "insult": "This comment insults a person or group.",
+    "obscene": "This comment contains obscene or vulgar language.",
+    "threat": "This comment contains a threat.",
+    "identity_attack": "This comment attacks someone because of their identity.",
+}
+
+
+def civil_comments(ex, rng, ctx):
+    """Targets are annotator fractions: a natural source of *soft*, calibrated labels."""
+    if ex["toxicity"] == 0 and rng.random() > 0.25:
+        return
+    state = _wrap(rng, "comment", ex["text"])
+    attrs = ["toxicity"] + ([rng.choice(list(_TOX)[1:])] if rng.random() < 0.5 else [])
+    for a in attrs:
+        t = float(ex[a])
+        yield state, _q("noul", _TOX[a]), int(t >= 0.5), [1 - t, t]
+
+
+# --------------------------------------------------------------------------- held-out sources
+def emotion(ex, rng, ctx):
+    names = ["sadness", "joy", "love", "anger", "fear", "surprise"]
+    q, y = _shuffled_choice(rng, "Which emotion does the author express?", [(n, "") for n in names], names[ex["label"]])
+    yield ex["text"], q, y, None
+
+
+def imdb(ex, rng, ctx):
+    yield {"review": ex["text"]}, _q("noul", "The reviewer liked the movie."), int(ex["label"]), None
+
+
+def massive_intent(ex, rng, ctx):
+    names = ctx["massive_labels"]
+    q, y = _shuffled_choice(rng, "Which intent does the voice assistant request express?",
+                            [(n, "") for n in names], ex["label_text"])  # fmt: skip
+    yield {"utterance": ex["text"]}, q, y, None
+
+
+def arc_easy(ex, rng, ctx):
+    texts, labels = ex["choices"]["text"], ex["choices"]["label"]
+    if ex["answerKey"] not in labels or len(texts) < 2:
+        return
+    criteria = {"ABCDE"[i]: t for i, t in enumerate(texts)}
+    yield {"question": ex["question"]}, _q("choice", "Select the correct answer to the question.", criteria), \
+        labels.index(ex["answerKey"]), None  # fmt: skip
+
+
+# name -> (hf id, config, train split, test split, stream?, adapter)
+TRAIN_SOURCES = {
+    "boolq": ("google/boolq", None, "train", "validation", False, boolq),
+    "ag_news": ("fancyzhx/ag_news", None, "train", "test", False, ag_news),
+    "yelp": ("Yelp/yelp_review_full", None, "train", "test", True, yelp),
+    "go_emotions": ("google-research-datasets/go_emotions", "simplified", "train", "test", False, go_emotions),
+    "banking77": ("mteb/banking77", None, "train", "test", False, banking77),
+    "mnli": ("nyu-mll/glue", "mnli", "train", "validation_matched", False, mnli),
+    "race": ("ehovy/race", "all", "train", "test", False, race),
+    "civil_comments": ("google/civil_comments", None, "train", "test", True, civil_comments),
+}
+HELDOUT_SOURCES = {
+    "emotion": ("dair-ai/emotion", "split", None, "test", False, emotion),
+    "imdb": ("stanfordnlp/imdb", "plain_text", None, "test", False, imdb),
+    "massive_intent": ("mteb/amazon_massive_intent", "en", None, "test", False, massive_intent),
+    "arc_easy": ("allenai/ai2_arc", "ARC-Easy", None, "test", False, arc_easy),
+}
+
+
+def take(hf_id, config, split, n, stream, seed, buffer_size=20_000):
+    from datasets import load_dataset
+
+    if stream:
+        ds = load_dataset(hf_id, config, split=split, streaming=True).shuffle(seed=seed, buffer_size=buffer_size)
+        return list(ds.take(n))
+    ds = load_dataset(hf_id, config, split=split).shuffle(seed=seed)
+    return [ds[i] for i in range(min(n, len(ds)))]
+
+
+def label_context() -> dict:
+    from datasets import load_dataset
+
+    return {
+        "banking77_labels": sorted(set(load_dataset("mteb/banking77", split="train")["label_text"])),
+        "massive_labels": sorted(set(load_dataset("mteb/amazon_massive_intent", "en", split="train")["label_text"])),
+    }
+
+
+def to_records(name, examples, adapter, rng, ctx, prefix):
+    out = []
+    for i, ex in enumerate(examples):
+        for state, question, label, soft in adapter(ex, rng, ctx):
+            rec = {"source": name, "sid": f"{prefix}-{name}-{i}", "state": state, "question": question, "label": label}
+            if soft is not None:
+                rec["soft"] = soft
+            out.append(rec)
+    return out
+
+
+# --------------------------------------------------------------------------- loading
+def read_jsonl(path) -> list[dict]:
+    with open(path) as f:
+        return [json.loads(line) for line in f if line.strip()]
+
+
+def write_jsonl(path, records):
+    Path(path).parent.mkdir(parents=True, exist_ok=True)
+    with open(path, "w") as f:
+        for r in records:
+            f.write(json.dumps(r, ensure_ascii=False) + "\n")
+
+
+def render_records(records: list[dict], renderer: Renderer, max_len: int | None = None, image_meta=None,
+                   image_root=None) -> tuple[list[Row], list[dict]]:  # fmt: skip
+    """Tokenise records into rows; drops rows longer than max_len (returns the kept records too).
+
+    Image records carry "image_grids" (token grid per image, recorded when the data was built), so their
+    length and positions are known without decoding any picture; pixels are loaded per batch by collate().
+    image_meta = the vision tower's zico_vision.json; image_root = directory the image paths are relative to.
+    """
+    from .batching import make_image_row
+    from .vision import ImageError, render_image_state
+
+    rows, kept, state_cache = [], [], {}
+    for rec in records:
+        grids = rec.get("image_grids")
+        if grids:
+            if image_meta is None:
+                raise ValueError("the data contains images but the checkpoint has no vision tower (scripts/add_vision.py)")
+            if rec["sid"] not in state_cache:
+                try:
+                    st = render_image_state(renderer, rec["state"], image_meta, grids=grids)
+                except ImageError:
+                    st = None  # state longer than the renderer's limit
+                if st is not None and image_root is not None:
+                    for part in st.images:
+                        part.path = str(Path(image_root) / part.path)
+                state_cache[rec["sid"]] = st
+            if state_cache[rec["sid"]] is None:
+                continue
+            row = make_image_row(state_cache[rec["sid"]], renderer.question(question_from_dict(rec["question"])),
+                                 label=rec["label"], soft=rec.get("soft"))  # fmt: skip
+            if max_len is None or len(row) <= max_len:
+                rows.append(row)
+                kept.append(rec)
+            continue
+        if rec["sid"] not in state_cache:
+            state_cache[rec["sid"]] = renderer.state(rec["state"])
+        row = make_row(state_cache[rec["sid"]], renderer.question(question_from_dict(rec["question"])),
+                       label=rec["label"], soft=rec.get("soft"))  # fmt: skip
+        if max_len is None or len(row) <= max_len:
+            rows.append(row)
+            kept.append(rec)
+    return rows, kept
+
+
+def subsample(records, n, seed=0):
+    if n is None or len(records) <= n:
+        return records
+    return random.Random(seed).sample(records, n)

zico/evaluate.py

@@ -0,0 +1,49 @@
+"""Batched prediction over rendered rows (the training-time, row-per-question path)."""
+from __future__ import annotations
+
+import torch
+
+from .batching import Row, bucketed_batches, collate
+from .calibration import metrics, metrics_by
+
+
+@torch.no_grad()
+def predict(model, rows: list[Row], records: list[dict], pad_id: int, max_tokens=8192, max_rows=64) -> list[dict]:
+    dev = next(model.parameters()).device
+    was_training = model.training
+    model.eval()
+    load_image = model.image_pre.load_prepared if getattr(model, "image_pre", None) is not None else None
+    out: list[dict | None] = [None] * len(rows)
+    for idx, pad_len in bucketed_batches([len(r) for r in rows], max_tokens, max_rows):
+        b = collate([rows[i] for i in idx], pad_id, pad_len, load_image)
+        logits = model(**{k: v.to(dev) for k, v in b.items()}).float().cpu()
+        for j, i in enumerate(idx):
+            k = len(rows[i].opt_pos)
+            lg = logits[j, :k]
+            out[i] = {"logits": lg, "p": torch.softmax(lg, -1), "label": rows[i].label, "qtype": rows[i].qtype,
+                      "source": records[i]["source"]}  # fmt: skip
+    model.train(was_training)
+    return out
+
+
+def report(preds: list[dict]) -> dict:
+    return {"overall": metrics(preds), "by_type": metrics_by(preds, "qtype"), "by_source": metrics_by(preds, "source")}
+
+
+def format_report(rep: dict, title: str = "") -> str:
+    lines = [f"{title:<22s} {'n':>6s} {'acc':>7s} {'nll':>7s} {'brier':>7s} {'ece':>7s} {'conf':>7s}"]
+
+    def row(name, m):
+        return f"{name:<22s} {m['n']:6d} {m['acc']:7.3f} {m['nll']:7.3f} {m['brier']:7.3f} {m['ece']:7.3f} {m['mean_conf']:7.3f}"
+
+    lines.append(row("overall", rep["overall"]))
+    lines += [row(f"  type:{k}", m) for k, m in rep["by_type"].items()]
+    lines += [row(f"  {k}", m) for k, m in rep["by_source"].items()]
+    return "\n".join(lines)
+
+
+def strip_tables(rep: dict) -> dict:
+    """Drop reliability tables for compact logging."""
+    if "reliability" in rep:
+        return {k: v for k, v in rep.items() if k != "reliability"}
+    return {k: strip_tables(v) if isinstance(v, dict) else v for k, v in rep.items()}

zico/image_sources.py

@@ -0,0 +1,99 @@
+"""Public image datasets recast as typed questions about an image in the state.
+
+Same philosophy as data.py: vary what a fixed classifier would hold constant (option subsets and
+order, instruction wording, where the image sits in the state) so the model reads the schema.
+Held-out sources are never trained on, as with text.
+
+An adapter yields (question, label index) pairs for one example; the builder supplies the state.
+"""
+from __future__ import annotations
+
+import re
+
+from .data import _q, _shuffled_choice
+
+
+def _clean(name) -> str:
+    name = re.sub(r"^\d+\.", "", str(name)).split(",")[0]  # "001.Black_footed_Albatross" / "bonnet, poke bonnet"
+    return re.sub(r"[\s_]+", " ", name).strip()
+
+
+def _article(name: str) -> str:
+    return ("an " if name[:1].lower() in "aeiou" else "a ") + name
+
+
+def classify(label_col: str, instructions: list[str], sizes: list[int], statement: str | None = "The image shows {a}."):
+    """Single-label classification -> a Choice over a sampled option subset, sometimes plus a yes/no statement."""
+
+    def adapter(ex, rng, names):
+        names = [_clean(n) for n in names[label_col]]
+        gold = names[ex[label_col]]
+        others = sorted(set(names) - {gold})
+        k = min(len(others) + 1, rng.choice(sizes))
+        opts = [gold] + rng.sample(others, k - 1)
+        yield _shuffled_choice(rng, rng.choice(instructions), [(o, "") for o in opts], gold)
+        if statement and rng.random() < 0.35:
+            truth = rng.random() < 0.5
+            shown = gold if truth else rng.choice(others)
+            yield _q("noul", statement.format(a=_article(shown), x=shown)), int(truth)
+
+    return adapter
+
+
+def cifar100(ex, rng, names):
+    yield from classify("fine_label", ["What does the image show?", "Which object is pictured?"], [10, 20, 100])(ex, rng, names)
+    if rng.random() < 0.5:
+        yield from classify("coarse_label", ["Which broad category does the pictured thing belong to?"], [20], None)(ex, rng, names)
+
+
+def aircraft(ex, rng, names):
+    yield from classify("label", ["Which aircraft model variant is this?"], [8, 16, 32], "The aircraft is {a}.")(ex, rng, names)
+    which = rng.choice(["family", "manufacturer"])
+    ask = "Which aircraft family is this?" if which == "family" else "Which company built this aircraft?"
+    yield from classify(which, [ask], [8, 16, 30], None)(ex, rng, names)
+
+
+_FASHION = {"masterCategory": "Which top-level category does the product belong to?", "subCategory": "Which sub-category is the product?",
+            "articleType": "What type of article is this product?", "baseColour": "What is the main colour of the product?",
+            "gender": "Who is the product intended for?", "usage": "What occasion is the product meant for?"}  # fmt: skip
+
+
+def fashion(ex, rng, names):
+    """Several attributes per product photo: natural multi-question supervision."""
+    for col in rng.sample(list(_FASHION), 2):
+        gold, pool = ex.get(col), names.get(col) or []
+        if not gold or gold not in pool or len(pool) < 3:
+            continue
+        others = [p for p in pool if p != gold]
+        opts = [gold] + rng.sample(others, min(len(others), rng.choice([3, 7, 15, 31])))
+        yield _shuffled_choice(rng, _FASHION[col], [(_clean(o), "") for o in opts], _clean(gold))
+
+
+def pope(ex, rng, names):
+    yield _q("noul", ex["question"]), int(str(ex["answer"]).strip().lower() == "yes")
+
+
+_DOC = ["What type of document is this scan?", "Classify the scanned document.", "Which kind of document is shown?"]
+
+# name: (hf id, config, train split, test split, stream?, image column, adapter, max pixels, state keys)
+IMAGE_TRAIN_SOURCES = {
+    "rvl_cdip": ("chainyo/rvl-cdip", None, "train", "test", True, "image", classify("label", _DOC, [6, 16, 16], "This is {a}."), 640 * 640, ["document", "scan", "page"]),
+    "imagenet100": ("clane9/imagenet-100", None, "train", "validation", True, "image", classify("label", ["What does the image show?", "What is the main subject of the photo?"], [8, 16, 32, 100]), 448 * 448, ["image", "photo"]),
+    "cifar100": ("uoft-cs/cifar100", "cifar100", "train", "test", False, "img", cifar100, 256 * 256, ["image", "picture"]),
+    "stanford_cars": ("tanganke/stanford_cars", None, "train", "test", False, "image", classify("label", ["Which car model is this?", "Identify the make, model and year of the car."], [8, 16, 32, 64], "The car is {a}."), 512 * 512, ["photo", "image"]),
+    "cub_birds": ("bentrevett/caltech-ucsd-birds-200-2011", None, "train", "test", False, "image", classify("label", ["Which bird species is this?"], [8, 16, 32, 64], "The bird is {a}."), 512 * 512, ["photo", "image"]),
+    "aircraft": ("Multimodal-Fatima/FGVC_Aircraft_train", None, "train", None, False, "image", aircraft, 512 * 512, ["photo", "image"]),
+    "dtd": ("tanganke/dtd", None, "train", "test", False, "image", classify("label", ["Which texture does the surface have?"], [8, 16, 47], "The texture is {x}."), 448 * 448, ["image", "surface"]),
+    "eurosat": ("tanganke/eurosat", None, "train", "test", False, "image", classify("label", ["What kind of land use does this satellite image show?"], [5, 10], "The satellite image shows {x}."), 256 * 256, ["satellite_image", "image"]),
+    "gtsrb": ("tanganke/gtsrb", None, "train", "test", False, "image", classify("label", ["Which traffic sign is this?"], [8, 16, 43], None), 256 * 256, ["image", "sign"]),
+    "resisc45": ("tanganke/resisc45", None, "train", "test", False, "image", classify("label", ["What does this aerial image show?"], [8, 16, 45], "The aerial image shows {a}."), 384 * 384, ["aerial_image", "image"]),
+    "sun397": ("tanganke/sun397", None, "train", "test", True, "image", classify("label", ["What kind of place is this?", "Which scene category fits the photo?"], [8, 16, 32, 64], "The photo was taken in {a}."), 448 * 448, ["photo", "image"]),
+    "fashion": ("ashraq/fashion-product-images-small", None, "train", None, False, "image", fashion, 256 * 256, ["product_photo", "image"]),
+}  # fmt: skip
+IMAGE_HELDOUT_SOURCES = {
+    "pets": ("timm/oxford-iiit-pet", None, None, "test", True, "image", classify("label", ["Which breed is the animal in the photo?"], [37], None), 448 * 448, ["photo"]),
+    "food101": ("ethz/food101", None, None, "validation", True, "image", classify("label", ["Which dish is shown in the photo?"], [101], None), 448 * 448, ["photo"]),
+    "cifar10": ("uoft-cs/cifar10", None, None, "test", False, "img", classify("label", ["What does the image show?"], [10], None), 256 * 256, ["image"]),
+    "mnist": ("ylecun/mnist", None, None, "test", False, "image", classify("label", ["Which digit is written in the image?"], [10], None), 256 * 256, ["image"]),
+    "pope": ("lmms-lab/POPE", "Full", None, "random", True, "image", pope, 448 * 448, ["image"]),
+}  # fmt: skip

zico/losses.py

@@ -0,0 +1,67 @@
+"""Training objectives: strictly proper scoring rules.
+
+Our reading of "Reinforcement Learning for Calibrated Decisions": the reward is a strictly
+proper scoring rule S(p, y) of the predicted distribution p against the realised outcome y.
+A proper rule is maximised in expectation only by the true conditional distribution, so
+optimising it targets calibration and accuracy together.
+
+  * Full feedback (the label y is observed). The model outputs the whole distribution, so
+    the expected reward and its gradient are available in closed form -- no sampling, no
+    policy-gradient variance. With the log score this *is* cross-entropy; Brier and the
+    ranked probability score (ordinal, for Score questions) are the other two we use.
+    Soft targets (annotator fractions, teacher vote shares) slot in unchanged.
+
+  * Bandit feedback (the model commits to an answer and only learns "right" / "wrong",
+    e.g. a verifier, a downstream outcome, a game). Then RL is actually needed. We use an
+    importance-weighted estimator that is unbiased for the same proper score:
+        a ~ mu,  c = 1[a == y]      E[ c * log p_a / mu_a ] = log p_y
+                                     E[ sum p^2 - 2 c p_a / mu_a ] = Brier(p, y) - 1
+    so the calibrated distribution remains the optimum although y is never revealed.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+from .schema import QTYPE_INDEX
+
+
+def log_probs(logits: torch.Tensor) -> torch.Tensor:
+    """Masked options already carry a large negative logit (model.NEG)."""
+    return F.log_softmax(logits.float(), dim=-1)
+
+
+def proper_score_loss(logits, opt_mask, target, qtype, brier_weight=0.0, rps_weight=0.0):
+    """Per-row loss. `target` is one-hot or a soft distribution over the row's options."""
+    logp = log_probs(logits)
+    p = logp.exp()
+    mask = opt_mask.float()
+    loss = -(target * logp * mask).sum(-1)
+    if brier_weight:
+        loss = loss + brier_weight * (((p - target) ** 2) * mask).sum(-1)
+    if rps_weight:
+        is_score = (qtype == QTYPE_INDEX["score"]).float()
+        k = mask.sum(-1).clamp(min=2)
+        cdf_gap = (torch.cumsum(p * mask, -1) - torch.cumsum(target * mask, -1)) ** 2
+        loss = loss + rps_weight * is_score * (cdf_gap * mask).sum(-1) / (k - 1)
+    return loss
+
+
+def bandit_proper_score_loss(logits, opt_mask, label, eps=0.1, brier_weight=0.0, clip=50.0, generator=None):
+    """Proper-score objective from correctness-only feedback (see module docstring).
+
+    Returns (per-row loss, per-row correctness of the sampled action).
+    """
+    logp = log_probs(logits)
+    p = logp.exp()
+    mask = opt_mask.float()
+    uniform = mask / mask.sum(-1, keepdim=True)
+    mu = (1 - eps) * p.detach() * mask + eps * uniform
+    mu = mu / mu.sum(-1, keepdim=True)
+    a = torch.multinomial(mu, 1, generator=generator)
+    correct = (a.squeeze(-1) == label).float()
+    w = (correct / mu.gather(1, a).squeeze(-1)).clamp(max=clip)  # clipping trades a little bias for variance
+    loss = -w * logp.gather(1, a).squeeze(-1)
+    if brier_weight:
+        loss = loss + brier_weight * (((p**2) * mask).sum(-1) - 2 * w * p.gather(1, a).squeeze(-1))
+    return loss, correct