Standard One 8B SH: adapter, schema head and server code

LICENSE

@@ -0,0 +1,201 @@
+                                 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 2023-2024 SGLang Team
+
+   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.

README.md

@@ -0,0 +1,106 @@
+---
+base_model: StandardThinking/StandardOne-8B
+base_model_relation: adapter
+library_name: peft
+license: apache-2.0
+pipeline_tag: text-classification
+language:
+  - en
+tags:
+  - mistral3
+  - decision-model
+  - typed-decisions
+  - schema-head
+  - jev
+  - calibration
+  - decode-free
+---
+
+# Standard One 8B SH
+
+**Standard One 8B SH** is Standard One 8B with a **joint schema head** (SH). It reads the scenario once and
+scores every option of every question in a single forward pass, then returns probabilities through the same
+`POST /v1/systemone` contract as Standard One. It does not generate text.
+
+| | |
+|---|---|
+| Backbone | [StandardOne-8B](https://huggingface.co/StandardThinking/StandardOne-8B), revision `v2.2` (`ac495d1c221297f0ae951941eff27bcf58b5ce52`) |
+| This repository | LoRA adapter (rank 64) for the backbone, the schema head (70.6M parameters) and the server code |
+| Question types | `choice` (2–255 options), `noul` (yes/no), `score` (ordinal scale) |
+| Context | up to 32,768 tokens per request |
+| License | Apache-2.0 |
+
+## Why a schema head
+
+Standard One reads one answer letter per option, so a question is limited by the label alphabet and every
+question in a request needs its own pass over the scenario. The schema head instead attends over the hidden states
+of the whole request (the final layer and one intermediate layer), places every question and option in one joint
+graph, and returns a calibrated distribution for each question at once. In practice:
+
+- **Many options.** Intent and routing questions with 100–255 options keep nearly the accuracy of small option sets.
+- **Many questions per scenario.** All questions of a request share one pass over the scenario.
+- **Long scenarios.** Decisions buried in long documents are found far more reliably than with Standard One v2.2.
+
+## Results
+
+All numbers below compare Standard One 8B SH with Standard One 8B v2.2 on identical items, scored the same way.
+
+| Benchmark | Standard One 8B v2.2 | Standard One 8B SH |
+|---|---:|---:|
+| [Decision Index](https://huggingface.co/spaces/multimodalart/jev-decision-index) 0.2.1 (self-run, full suite) | 41.14 | **46.17** |
+| JevBench public — hard tier (111 items) | 56.76 | **69.37** |
+| JevBench public — original tier (72 items) | 98.61 | 98.61 |
+| JevBench public — easy tier (48 items) | 100.00 | 100.00 |
+| Intent labels with 27–151 options (CLINC150, BANKING77, MASSIVE; 18,000 items) | 82.67 | **92.09** |
+| Intent labels with 152–255 options (CLINC150 utterances, 4,400 items) | 85.34 | **96.91** |
+| Decision hidden in a long document (150 items) | 40.00 | **88.00** |
+| WinoGrande (dev, 2,010 items) | 87.86 | **94.98** |
+| Policy-judgement repair (1,200 items) | 60.58 | **74.25** |
+| Game positions: dots and boxes / snake and 2048 / tetris | 38.37 / 58.63 / 54.86 | **51.96 / 68.37 / 64.19** |
+| Held-out decision suite (600 items) | 82.50 | **88.33** |
+| Held-out realistic decision suite (600 items) | 91.50 | **91.67** |
+
+Accuracy in percent, except the Decision Index (its own 0–100 index). The Decision Index run covers all 150,759
+in-edition requests of edition 0.2.1; it is our own run with the kit's scorer, not a leaderboard entry.
+
+## Quick start
+
+```bash
+pip install "torch>=2.4" "transformers==5.12.1" "peft==0.21.0" safetensors
+hf download StandardThinking/StandardOne-8B --revision v2.2 --local-dir StandardOne-8B
+hf download StandardThinking/StandardOne-8B-SH --local-dir StandardOne-8B-SH
+
+python StandardOne-8B-SH/code/serve_head.py \
+  --snapshot StandardOne-8B --adapter StandardOne-8B-SH/adapter --head StandardOne-8B-SH/head \
+  --port 30171
+```
+
+```bash
+curl -s localhost:30171/v1/systemone -H 'content-type: application/json' -d '{
+  "state": {"message": "Please close my card, I lost it yesterday."},
+  "questions": {
+    "intent": {"type": "choice", "instructions": "What does the customer want?",
+               "criteria": {"cancel": "close the card", "limit": "change the limit", "other": null}},
+    "urgent": {"type": "noul", "instructions": "The request is urgent."},
+    "mood":   {"type": "score", "instructions": "How upset is the customer?", "criteria": ["calm", "annoyed", "angry"]}
+  }
+}'
+```
+
+The server batches concurrent requests (`--max-batch-requests`, `--max-batch-tokens`) and binds `127.0.0.1` by
+default. It runs on one GPU with the backbone in bfloat16 and the head in float32.
+
+## Limitations
+
+- Requests longer than 32,768 tokens are refused, not truncated.
+- On our hard held-out decision suite (600 items; long policies, multi-step lookups, date and number reasoning)
+  it scores 43.50 against 53.50 for Standard One 8B v2.2. Prefer Standard One 8B for that kind of workload.
+- Text input only has been validated for this release.
+- The server is a PyTorch reference server; it does not yet run on an inference engine, so throughput under heavy
+  load is lower than the engine-served Standard One models.
+- Probabilities are calibrated on our held-out data; refit thresholds on your own distribution before relying on them.
+
+## License
+
+Apache-2.0. The backbone is [Standard One 8B](https://huggingface.co/StandardThinking/StandardOne-8B) (Apache-2.0),
+built on [Ministral 3 8B](https://huggingface.co/mistralai/Ministral-3-8B-Instruct-2512-BF16) (Apache-2.0).

adapter/adapter_config.json

@@ -0,0 +1,54 @@
+{
+  "alora_invocation_tokens": null,
+  "alpha_pattern": {},
+  "arrow_config": null,
+  "auto_mapping": {
+    "base_model_class": "Mistral3ForConditionalGeneration",
+    "parent_library": "transformers.models.mistral3.modeling_mistral3"
+  },
+  "base_model_name_or_path": "StandardThinking/StandardOne-8B",
+  "bias": "none",
+  "corda_config": null,
+  "ensure_weight_tying": false,
+  "eva_config": null,
+  "exclude_modules": ".*\\.(vision_tower|multi_modal_projector)\\..*",
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "init_lora_weights": true,
+  "kasa_config": null,
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 128.0,
+  "lora_bias": false,
+  "lora_dropout": 0.0,
+  "lora_ga_config": null,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "monteclora_config": null,
+  "peft_type": "LORA",
+  "peft_version": "0.21.0",
+  "qalora_group_size": 16,
+  "r": 64,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": [
+    "v_proj",
+    "k_proj",
+    "q_proj",
+    "gate_proj",
+    "down_proj",
+    "o_proj",
+    "up_proj"
+  ],
+  "target_parameters": null,
+  "task_type": null,
+  "trainable_token_indices": null,
+  "use_bdlora": null,
+  "use_dora": false,
+  "use_qalora": false,
+  "use_rslora": false,
+  "velora_config": null
+}

adapter/adapter_model.safetensors

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

code/head.py

@@ -0,0 +1,325 @@
+"""Joint schema head (DESIGN.md section 2.2). Own design; provenance in DESIGN.md section 0.
+
+Reads tapped backbone hidden states for ONE request at a time (a micro-batch is a Python loop over requests; the
+head is small) and returns one logit vector per question over that question's options in CANONICAL order.
+
+Nodes: one per question and one per option. No positional encoding over nodes, attention pooling inside spans and
+relation-type attention biases that depend only on structure make the head exactly equivariant to the order of
+choice/noul options and of questions (tests/test_equivariance.py). Score levels carry an ordinal encoding and are
+deliberately not permutation-equivariant.
+"""
+import math
+from dataclasses import asdict, dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint
+
+QTYPE_ID = {"choice": 0, "noul": 1, "score": 2}
+# relation types for node pairs (i attends to j)
+REL_SELF, REL_SIBLING, REL_OPT_TO_OWN_Q, REL_Q_TO_OWN_OPT, REL_OPT_TO_OTHER_OPT, REL_OPT_TO_OTHER_Q, \
+    REL_Q_TO_OTHER_Q, REL_Q_TO_OTHER_OPT = range(8)
+N_REL = 8
+ORD_FEATS = 33
+
+
+@dataclass
+class HeadConfig:
+    hidden_size: int = 4096
+    n_taps: int = 2
+    d: int = 1024
+    heads: int = 16
+    ffn: int = 2816
+    route_blocks: int = 2
+    interact_blocks: int = 2
+    dropout: float = 0.1
+    n_roles: int = 5
+    ord_hidden: int = 256
+    max_nodes: int = 4096
+    grad_checkpoint: bool = True
+    taps: list = field(default_factory=lambda: ["final", 25])
+    template: str = "schema-v1"
+
+    def to_dict(self):
+        return asdict(self)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, d, hidden, dropout):
+        super().__init__()
+        self.w1 = nn.Linear(d, hidden, bias=False)
+        self.w3 = nn.Linear(d, hidden, bias=False)
+        self.w2 = nn.Linear(hidden, d, bias=False)
+        self.drop = nn.Dropout(dropout)
+
+    def forward(self, x):
+        return self.drop(self.w2(F.silu(self.w1(x)) * self.w3(x)))
+
+
+class Attention(nn.Module):
+    """Multi-head attention (queries from nodes, keys/values from `kv`), SDPA with optional additive bias."""
+
+    def __init__(self, d, heads, dropout):
+        super().__init__()
+        self.h, self.dh = heads, d // heads
+        self.q = nn.Linear(d, d, bias=False)
+        self.k = nn.Linear(d, d, bias=False)
+        self.v = nn.Linear(d, d, bias=False)
+        self.o = nn.Linear(d, d, bias=False)
+        self.p = dropout
+
+    def _chunk(self, q, part, p):
+        k = self.k(part).view(-1, self.h, self.dh).transpose(0, 1).unsqueeze(0)
+        v = self.v(part).view(-1, self.h, self.dh).transpose(0, 1).unsqueeze(0)
+        scores = (q @ k.transpose(-1, -2)) / math.sqrt(self.dh)
+        lse = torch.logsumexp(scores, -1, keepdim=True)
+        probs = F.dropout(torch.softmax(scores, -1), p, self.training)
+        return probs @ v, lse
+
+    def forward(self, x, kv, bias=None, kv_chunk=0):
+        n, m = x.shape[0], kv.shape[0]
+        q = self.q(x).view(n, self.h, self.dh).transpose(0, 1).unsqueeze(0)
+        p = self.p if self.training else 0.0
+        if kv_chunk and m > kv_chunk and bias is None:
+            # exact chunked attention over keys (log-sum-exp merge) to bound K/V memory at long contexts.
+            # Dropout is applied to each chunk's attention PROBABILITIES (as SDPA's dropout_p): the global weight of
+            # key j is softmax_chunk(j) * exp(lse_chunk - lse_all), so dropping the chunk-local probability with
+            # scale 1/(1-p) is the same random variable as dropping the global one. When training, every chunk is
+            # recomputed in backward (checkpoint), so its scores/probabilities are not kept for the whole sequence.
+            outs, lses = [], []
+            use_ckpt = self.training and torch.is_grad_enabled()
+            for s in range(0, m, kv_chunk):
+                part = kv[s:s + kv_chunk]
+                if use_ckpt:
+                    o, lse = checkpoint(self._chunk, q, part, p, use_reentrant=False)
+                else:
+                    o, lse = self._chunk(q, part, p)
+                outs.append(o)
+                lses.append(lse)
+            lse_all = torch.logsumexp(torch.cat(lses, -1), -1, keepdim=True)
+            out = sum(o * torch.exp(l - lse_all) for o, l in zip(outs, lses))
+        else:
+            k = self.k(kv).view(m, self.h, self.dh).transpose(0, 1).unsqueeze(0)
+            v = self.v(kv).view(m, self.h, self.dh).transpose(0, 1).unsqueeze(0)
+            mask = None if bias is None else bias.unsqueeze(0)
+            out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=p)
+        return self.o(out.squeeze(0).transpose(0, 1).reshape(n, -1))
+
+
+class RouteBlock(nn.Module):
+    """Stage 1 (evidence routing): nodes cross-attend over the state tokens, then FFN."""
+
+    def __init__(self, cfg):
+        super().__init__()
+        self.n1, self.n2, self.nkv = RMSNorm(cfg.d), RMSNorm(cfg.d), RMSNorm(cfg.d)
+        self.cross = Attention(cfg.d, cfg.heads, cfg.dropout)
+        self.ffn = SwiGLU(cfg.d, cfg.ffn, cfg.dropout)
+        self.drop = nn.Dropout(cfg.dropout)
+
+    def forward(self, x, state, kv_chunk=0):
+        x = x + self.drop(self.cross(self.n1(x), self.nkv(state), kv_chunk=kv_chunk))
+        return x + self.ffn(self.n2(x))
+
+
+class InteractBlock(nn.Module):
+    """Stage 2 (field interaction): node self-attention with relation-type biases, cross-attention back to the
+    state tokens, FFN."""
+
+    def __init__(self, cfg):
+        super().__init__()
+        self.n1, self.n2, self.n3, self.nkv = RMSNorm(cfg.d), RMSNorm(cfg.d), RMSNorm(cfg.d), RMSNorm(cfg.d)
+        self.self_attn = Attention(cfg.d, cfg.heads, cfg.dropout)
+        self.cross = Attention(cfg.d, cfg.heads, cfg.dropout)
+        self.rel_bias = nn.Parameter(torch.zeros(N_REL, cfg.heads))
+        self.ffn = SwiGLU(cfg.d, cfg.ffn, cfg.dropout)
+        self.drop = nn.Dropout(cfg.dropout)
+
+    def forward(self, x, state, rel, kv_chunk=0):
+        bias = self.rel_bias[rel].permute(2, 0, 1)  # [heads, N, N]
+        y = self.n1(x)
+        x = x + self.drop(self.self_attn(y, y, bias=bias))
+        x = x + self.drop(self.cross(self.n2(x), self.nkv(state), kv_chunk=kv_chunk))
+        return x + self.ffn(self.n3(x))
+
+
+def ordinal_features(k, K):
+    """33 features of score level k of K (2 <= K <= 10): position, scale size, one-hots, sin/cos."""
+    pos = k / (K - 1)
+    feats = [pos, K / 10.0]
+    feats += [1.0 if i == k else 0.0 for i in range(10)]
+    feats += [1.0 if K == i else 0.0 for i in range(2, 11)]
+    for f in range(1, 7):
+        feats += [math.sin(math.pi * f * pos), math.cos(math.pi * f * pos)]
+    assert len(feats) == ORD_FEATS
+    return feats
+
+
+class JointSchemaHead(nn.Module):
+    def __init__(self, cfg: HeadConfig):
+        super().__init__()
+        self.cfg = cfg
+        d = cfg.d
+        self.tap_norms = nn.ModuleList([RMSNorm(cfg.hidden_size) for _ in range(cfg.n_taps)])
+        self.tap_mix = nn.Parameter(torch.zeros(cfg.n_taps))
+        self.tap_gain = nn.Parameter(torch.ones(1))
+        self.in_proj = nn.Linear(cfg.hidden_size, d, bias=False)
+        self.role_emb = nn.Embedding(cfg.n_roles, d)
+        self.qtype_emb = nn.Embedding(3, d)
+        self.kind_emb = nn.Embedding(2, d)  # 0 question node, 1 option node
+        self.pool_query = nn.Parameter(torch.randn(2, d) / math.sqrt(d))
+        self.pool_key = nn.Linear(d, d, bias=False)
+        self.pool_out = nn.Linear(2 * d, d, bias=False)
+        self.ord_mlp = nn.Sequential(nn.Linear(ORD_FEATS, cfg.ord_hidden), nn.GELU(), nn.Linear(cfg.ord_hidden, d))
+        self.route = nn.ModuleList([RouteBlock(cfg) for _ in range(cfg.route_blocks)])
+        self.interact = nn.ModuleList([InteractBlock(cfg) for _ in range(cfg.interact_blocks)])
+        self.final_norm = RMSNorm(d)
+        self.scorer = nn.Sequential(nn.Linear(3 * d, d), nn.GELU(), nn.Linear(d, 1))
+        self.log_temp = nn.Parameter(torch.zeros(3))
+        self.kv_chunk = 0
+        self._init()
+
+    def _init(self):
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, std=0.02)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.Embedding):
+                nn.init.normal_(m.weight, std=0.02)
+
+    # ------------------------------------------------------------------ structure
+    @staticmethod
+    def plan(enc, device):
+        """Index tensors for one encoded request (see render.Encoder.encode)."""
+        node_kind, node_q, node_qtype, node_ord = [], [], [], []
+        pool_node, pool_tok, last_tok = [], [], []
+        q_node, opt_nodes = [], []
+        nq = len(enc["keys"])
+        for i in range(nq):
+            qt = enc["qtype"][i]
+            idx = len(node_kind)
+            q_node.append(idx)
+            node_kind.append(0); node_q.append(i); node_qtype.append(qt); node_ord.append(None)
+            toks = list(range(*enc["q_span"][i]))
+            if enc["p_span"][i]:
+                toks = list(range(*enc["p_span"][i])) + toks
+            pool_node += [idx] * len(toks); pool_tok += toks
+            last_tok.append(enc["q_span"][i][1] - 1)
+            ids = []
+            K = enc["n_opt"][i]
+            for j in range(K):
+                idx = len(node_kind)
+                ids.append(idx)
+                node_kind.append(1); node_q.append(i); node_qtype.append(qt)
+                node_ord.append((j, K) if qt == "score" else None)
+                a, b = enc["o_span"][i][j]
+                pool_node += [idx] * (b - a); pool_tok += list(range(a, b))
+                last_tok.append(b - 1)
+            opt_nodes.append(ids)
+        n = len(node_kind)
+        kind = torch.tensor(node_kind)
+        qn = torch.tensor(node_q)
+        same_q = qn[:, None] == qn[None, :]
+        ki, kj = kind[:, None], kind[None, :]
+        rel = torch.empty(n, n, dtype=torch.long)
+        rel[(ki == 1) & (kj == 1) & same_q] = REL_SIBLING
+        rel[(ki == 1) & (kj == 0) & same_q] = REL_OPT_TO_OWN_Q
+        rel[(ki == 0) & (kj == 1) & same_q] = REL_Q_TO_OWN_OPT
+        rel[(ki == 1) & (kj == 1) & ~same_q] = REL_OPT_TO_OTHER_OPT
+        rel[(ki == 1) & (kj == 0) & ~same_q] = REL_OPT_TO_OTHER_Q
+        rel[(ki == 0) & (kj == 0)] = REL_Q_TO_OTHER_Q
+        rel[(ki == 0) & (kj == 1) & ~same_q] = REL_Q_TO_OTHER_OPT
+        rel.fill_diagonal_(REL_SELF)
+        ords = [ordinal_features(*o) if o else [0.0] * ORD_FEATS for o in node_ord]
+        state_tok = [t for a, b in enc["state_spans"] for t in range(a, b)]
+        p = {"n": n, "kind": kind.to(device), "qtype": torch.tensor([QTYPE_ID[t] for t in node_qtype]).to(device),
+             "is_score": torch.tensor([o is not None for o in node_ord]).to(device),
+             "ord": torch.tensor(ords, dtype=torch.float32).to(device), "rel": rel.to(device),
+             "pool_node": torch.tensor(pool_node).to(device), "pool_tok": torch.tensor(pool_tok).to(device),
+             "last_tok": torch.tensor(last_tok).to(device), "q_node": q_node, "opt_nodes": opt_nodes,
+             "state_tok": torch.tensor(state_tok).to(device), "qtypes": list(enc["qtype"])}
+        return p
+
+    # ------------------------------------------------------------------ forward
+    def mix_taps(self, taps):
+        """taps: list of [T, hidden] tensors (any float dtype) -> [T, hidden] in the head's dtype (fp32)."""
+        w = torch.softmax(self.tap_mix, 0)
+        dtype = self.in_proj.weight.dtype
+        out = 0
+        for i, t in enumerate(taps):
+            out = out + w[i] * self.tap_norms[i](t.to(dtype))
+        return out * self.tap_gain
+
+    def _token_proj(self, roles, *taps):
+        return self.in_proj(self.mix_taps(list(taps))) + self.role_emb(roles)
+
+    def project_tokens(self, taps, roles, use_ckpt, chunk=8192):
+        """Per-token tap mix + in-proj + role embedding, in token chunks (each checkpointed when training), so the
+        fp32 intermediates of only one chunk exist at a time (forward and recompute)."""
+        outs = []
+        for a in range(0, roles.shape[0], chunk):
+            args = (roles[a:a + chunk], *[t[a:a + chunk] for t in taps])
+            outs.append(checkpoint(self._token_proj, *args, use_reentrant=False) if use_ckpt else self._token_proj(*args))
+        return outs[0] if len(outs) == 1 else torch.cat(outs, 0)
+
+    def forward_one(self, taps, enc, roles):
+        """taps: list of [T, hidden]; enc: encoded request; roles: LongTensor [T]. Returns list of logit vectors
+        (one per question in enc['keys'] order), each over CANONICAL option order, already divided by the
+        per-type temperature."""
+        cfg = self.cfg
+        dev = taps[0].device
+        p = self.plan(enc, dev)
+        require_nodes = p["n"] <= cfg.max_nodes
+        if not require_nodes:
+            raise ValueError(f"too many nodes: {p['n']} > {cfg.max_nodes}")
+        use_ckpt = cfg.grad_checkpoint and self.training and torch.is_grad_enabled()
+        # token-level projection: chunked + checkpointed so only the (already alive, bf16) tapped states are kept for
+        # backward, not ~4 fp32 [T, hidden] intermediates per tap (about 2 GB per tap at 32k tokens)
+        x_tok = self.project_tokens(taps, roles, use_ckpt)
+        state = x_tok.index_select(0, p["state_tok"])
+        # attention pooling inside spans (segment softmax via scatter)
+        toks = x_tok.index_select(0, p["pool_tok"])
+        node_of = p["pool_node"]
+        kind_of = p["kind"].index_select(0, node_of)
+        query = self.pool_query.index_select(0, kind_of)
+        scores = (self.pool_key(toks) * query).sum(-1) / math.sqrt(cfg.d)
+        mx = torch.full((p["n"],), -1e30, device=dev, dtype=scores.dtype)
+        mx = mx.scatter_reduce(0, node_of, scores, reduce="amax", include_self=True)
+        e = torch.exp(scores - mx.index_select(0, node_of))
+        den = torch.zeros(p["n"], device=dev, dtype=e.dtype).index_add(0, node_of, e)
+        w = e / den.index_select(0, node_of)
+        pooled = torch.zeros(p["n"], cfg.d, device=dev, dtype=toks.dtype).index_add(0, node_of, toks * w[:, None])
+        last = x_tok.index_select(0, p["last_tok"])
+        x = self.pool_out(torch.cat([pooled, last], -1))
+        x = x + self.qtype_emb(p["qtype"]) + self.kind_emb(p["kind"])
+        x = x + self.ord_mlp(p["ord"].to(x.dtype)) * p["is_score"][:, None].to(x.dtype)
+        for blk in self.route:
+            x = checkpoint(blk, x, state, self.kv_chunk, use_reentrant=False) if use_ckpt else blk(x, state, self.kv_chunk)
+        for blk in self.interact:
+            x = (checkpoint(blk, x, state, p["rel"], self.kv_chunk, use_reentrant=False) if use_ckpt
+                 else blk(x, state, p["rel"], self.kv_chunk))
+        x = self.final_norm(x)
+        out = []
+        for i, qn in enumerate(p["q_node"]):
+            opts = torch.tensor(p["opt_nodes"][i], device=dev)
+            o = x.index_select(0, opts)
+            q = x[qn].expand_as(o)
+            z = self.scorer(torch.cat([o, q, o * q], -1)).squeeze(-1)
+            t = torch.exp(self.log_temp[QTYPE_ID[p["qtypes"][i]]])
+            out.append(z / t)
+        return out
+
+
+def count_parameters(module):
+    return sum(p.numel() for p in module.parameters())

code/model.py

@@ -0,0 +1,239 @@
+"""Backbone (Ministral-3-8B, Mistral3ForConditionalGeneration) + LoRA + joint schema head (DESIGN.md 2.1, 3.1).
+
+- The inner model (`.model`, Mistral3Model) is called; its `last_hidden_state` is the post-norm final state.
+  `lm_head` is never called and is deleted after load (131,072 x 4,096 bf16 = 1.07e9 B).
+- Extra taps: outputs of decoder layers by index, captured with forward hooks.
+- LoRA: language-model linear layers only (q/k/v/o/gate/up/down), vision tower + projector frozen, alpha = 2r.
+  Same leaf discovery and vision-exclusion rule as the Standard One LoRA trainer.
+"""
+import json
+import os
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+
+from head import HeadConfig, JointSchemaHead
+from render import token_roles
+
+VISION_EXCLUDE_REGEX = r".*\.(vision_tower|multi_modal_projector)\..*"
+
+
+def lora_leaves(model):
+    """Leaf names of the language-model attention + MLP linears (as the trainer's `attention_mlp` targets)."""
+    attention, mlp = set(), set()
+    for name, module in model.named_modules():
+        if not isinstance(module, nn.Linear) or ".language_model." not in name:
+            continue
+        leaf = name.rsplit(".", 1)[-1]
+        if ".self_attn." in name:
+            attention.add(leaf)
+        elif ".mlp." in name and leaf in {"gate_proj", "up_proj", "down_proj"}:
+            mlp.add(leaf)
+    if not attention:
+        raise ValueError("no attention projections found")
+    return sorted(attention | mlp)
+
+
+def decoder_layers(model):
+    lm = model.model.language_model
+    return lm.layers
+
+
+def tiny_config(vocab_size=131072, layers=3, hidden=64, heads=4, kv_heads=2, intermediate=128):
+    """A tiny, randomly initialised Mistral3 (Ministral3 text + Pixtral vision) config for CPU smoke tests."""
+    from transformers import Mistral3Config
+    text = {"model_type": "ministral3", "hidden_size": hidden, "intermediate_size": intermediate,
+            "num_hidden_layers": layers, "num_attention_heads": heads, "num_key_value_heads": kv_heads,
+            "head_dim": hidden // heads, "vocab_size": vocab_size, "max_position_embeddings": 262144,
+            "rms_norm_eps": 1e-5, "hidden_act": "silu", "tie_word_embeddings": False, "sliding_window": None,
+            "rope_parameters": {"beta_fast": 32.0, "beta_slow": 1.0, "factor": 16.0, "llama_4_scaling_beta": 0.1,
+                                "mscale": 1.0, "mscale_all_dim": 1.0, "original_max_position_embeddings": 16384,
+                                "rope_theta": 1000000.0, "rope_type": "yarn", "type": "yarn"}}
+    vision = {"model_type": "pixtral", "hidden_size": 32, "intermediate_size": 64, "num_hidden_layers": 1,
+              "num_attention_heads": 2, "head_dim": 16, "image_size": 1540, "patch_size": 14, "num_channels": 3,
+              "rope_parameters": {"rope_theta": 10000.0, "rope_type": "default"}, "rope_theta": 10000.0}
+    return Mistral3Config(text_config=text, vision_config=vision, image_token_index=10, spatial_merge_size=2,
+                          multimodal_projector_bias=False, projector_hidden_act="gelu", vision_feature_layer=-1,
+                          tie_word_embeddings=False)
+
+
+class SchemaModel(nn.Module):
+    """Backbone (optionally LoRA-wrapped) + head. `forward_micro` runs one micro-batch of encoded requests."""
+
+    def __init__(self, backbone, head, taps, peft_model=None):
+        super().__init__()
+        self.backbone = backbone          # Mistral3ForConditionalGeneration (LoRA layers injected in place)
+        self.peft_model = peft_model      # the PeftModel wrapper (for save_pretrained) or None
+        self.head = head
+        self.taps = list(taps)
+        self._captured = {}
+        self._capturing = False   # hooks record only during our own forward, not during checkpoint recompute
+        self._hooks = []
+        layers = decoder_layers(backbone)
+        for t in self.taps:
+            if t == "final":
+                continue
+            idx = int(t)
+            if not 0 <= idx < len(layers):
+                raise ValueError(f"tap layer {idx} out of range (0..{len(layers) - 1})")
+            self._hooks.append(layers[idx].register_forward_hook(self._make_hook(idx)))
+
+    def _make_hook(self, idx):
+        def hook(module, inputs, output):
+            if self._capturing:
+                self._captured[idx] = output[0] if isinstance(output, (tuple, list)) else output
+        return hook
+
+    def device(self):
+        return next(self.head.parameters()).device
+
+    def backbone_states(self, batch):
+        """Run the inner model on a padded batch; returns list of [B, T, H] tensors in tap order."""
+        self._captured = {}
+        dev = batch["input_ids"].device
+        kwargs = {k: v for k, v in batch.items() if k in ("input_ids", "attention_mask", "position_ids",
+                                                           "pixel_values", "image_sizes")}
+        use_bf16 = dev.type == "cuda"
+        self._capturing = True
+        try:
+            with torch.autocast(device_type=dev.type, dtype=torch.bfloat16, enabled=use_bf16):
+                out = self.backbone.model(**kwargs, use_cache=False, return_dict=True)
+        finally:
+            self._capturing = False
+        states = []
+        for t in self.taps:
+            states.append(out.last_hidden_state if t == "final" else self._captured[int(t)])
+        self._captured = {}
+        return states
+
+    def forward_micro(self, encs, batch):
+        """encs: list of encoded requests (render.Encoder.encode outputs, right-padded in `batch`).
+        Returns list (per request) of lists (per question) of fp32 logit vectors (canonical option order)."""
+        states = self.backbone_states(batch)
+        outputs = []
+        dev = states[0].device
+        with torch.autocast(device_type=dev.type, enabled=False):
+            for i, enc in enumerate(encs):
+                n = enc["n_tokens"]
+                taps = [s[i, :n] for s in states]   # bf16 on GPU; the head casts inside its checkpointed projection
+                roles = torch.tensor(enc.get("_roles") or token_roles(enc), device=dev)
+                outputs.append(self.head.forward_one(taps, enc, roles))
+        return outputs
+
+    def trainable_named(self):
+        return [(n, p) for n, p in self.named_parameters() if p.requires_grad]
+
+
+def pad_pixel_values(chunks):
+    """Join per-row pixel tensors [n_i, C, H_i, W_i] into one [sum n_i, C, maxH, maxW] tensor, zero-padded at the
+    bottom/right. Every row went through the processor on its own, so H/W differ between rows; this is exactly what
+    PixtralImageProcessor._pad_for_batching does for a multi-image call, and the vision tower crops each image back
+    to its `image_sizes` entry (patch_embeds[..., :h // patch, :w // patch]) so the padding never reaches the model."""
+    h = max(int(c.shape[-2]) for c in chunks)
+    w = max(int(c.shape[-1]) for c in chunks)
+    return torch.cat([torch.nn.functional.pad(c, (0, w - int(c.shape[-1]), 0, h - int(c.shape[-2])))
+                      for c in chunks], 0)
+
+
+def make_batch(encs, pad_id, device):
+    """Right-padded text batch (position ids restart per row); image micro-batches carry pixel values."""
+    longest = max(e["n_tokens"] for e in encs)
+    ids = torch.full((len(encs), longest), pad_id, dtype=torch.long)
+    mask = torch.zeros((len(encs), longest), dtype=torch.long)
+    for i, e in enumerate(encs):
+        ids[i, :e["n_tokens"]] = torch.tensor(e["input_ids"])
+        mask[i, :e["n_tokens"]] = 1
+    pos = torch.arange(longest).unsqueeze(0).expand(len(encs), -1).contiguous()
+    batch = {"input_ids": ids.to(device), "attention_mask": mask.to(device), "position_ids": pos.to(device)}
+    pix = [e.get("_pixel") for e in encs]
+    if any(p is not None for p in pix):
+        if not all(p is not None for p in pix):
+            raise ValueError("image micro-batches must be homogeneous")
+        batch["pixel_values"] = pad_pixel_values([torch.as_tensor(p["pixel_values"]) for p in pix]).to(device)
+        batch["image_sizes"] = torch.cat([torch.as_tensor(p["image_sizes"]) for p in pix], 0).to(device)
+        if device.type == "cuda":
+            batch["pixel_values"] = batch["pixel_values"].to(torch.bfloat16)
+    return batch
+
+
+def build_model(snapshot=None, tiny=None, lora_rank=64, lora_alpha=None, lora_dropout=0.0, taps=("final", 25),
+                head_overrides=None, device="cpu", grad_checkpointing=True, init_adapter=None, attn="sdpa",
+                seed=0):
+    """Load (or create tiny) backbone, delete lm_head, freeze, add LoRA (rank 0 = none), attach the head.
+
+    Returns (SchemaModel, info dict)."""
+    import peft
+    import transformers
+    torch.manual_seed(seed)
+    if tiny is not None:
+        torch.manual_seed(1234)   # fixed: the tiny random "base" must be identical in training and serving
+        cfg = tiny if not isinstance(tiny, dict) else tiny_config(**tiny)
+        backbone = transformers.Mistral3ForConditionalGeneration._from_config(cfg, attn_implementation=attn)
+        backbone = backbone.to(device)
+    else:
+        backbone = transformers.Mistral3ForConditionalGeneration.from_pretrained(
+            str(snapshot), local_files_only=True, trust_remote_code=False, dtype=torch.bfloat16,
+            attn_implementation=attn, device_map={"": device} if str(device) != "cpu" else None)
+        if str(device) == "cpu":
+            backbone = backbone.to(device)
+    torch.manual_seed(seed)
+    backbone.config.use_cache = False
+    lm_head_bytes = backbone.lm_head.weight.numel() * backbone.lm_head.weight.element_size()
+    del backbone.lm_head
+    backbone.lm_head = nn.Identity()  # never called; keeps attribute access working
+    for p in backbone.parameters():
+        p.requires_grad_(False)
+    if grad_checkpointing:
+        backbone.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+        vt = getattr(backbone.model, "vision_tower", None)
+        if vt is not None and hasattr(vt, "gradient_checkpointing_disable"):
+            vt.gradient_checkpointing_disable()
+    leaves = lora_leaves(backbone)
+    peft_model = None
+    if init_adapter:
+        peft_model = peft.PeftModel.from_pretrained(backbone, str(init_adapter), is_trainable=True)
+    elif lora_rank:
+        alpha = lora_alpha if lora_alpha is not None else 2 * lora_rank
+        peft_model = peft.get_peft_model(backbone, peft.LoraConfig(
+            r=lora_rank, lora_alpha=alpha, lora_dropout=lora_dropout, target_modules=leaves,
+            exclude_modules=VISION_EXCLUDE_REGEX, bias="none"))
+    leaked = 0
+    for name, p in backbone.named_parameters():
+        if "vision_tower" in name or "multi_modal_projector" in name:
+            if p.requires_grad:
+                leaked += 1
+            p.requires_grad_(False)
+        elif p.requires_grad:
+            p.data = p.data.float()  # LoRA weights fp32 (as the trainer)
+    text_cfg = backbone.config.text_config
+    hcfg = HeadConfig(hidden_size=text_cfg.hidden_size, n_taps=len(taps), taps=list(taps),
+                      **(head_overrides or {}))
+    head = JointSchemaHead(hcfg).to(device=device, dtype=torch.float32)
+    model = SchemaModel(backbone, head, taps, peft_model)
+    info = {"lora_leaves": leaves, "lora_rank": lora_rank,
+            "lora_alpha": (lora_alpha if lora_alpha is not None else 2 * lora_rank) if lora_rank else None,
+            "lora_params": sum(p.numel() for n, p in backbone.named_parameters() if p.requires_grad),
+            "head_params": sum(p.numel() for p in head.parameters()), "lm_head_bytes_freed": lm_head_bytes,
+            "vision_lora_leaked_frozen": leaked, "taps": list(taps), "head_config": hcfg.to_dict(),
+            "layers": len(decoder_layers(backbone)), "hidden_size": text_cfg.hidden_size}
+    return model, info
+
+
+def save_head(head, folder, extra=None):
+    from safetensors.torch import save_file
+    folder = Path(folder)
+    folder.mkdir(parents=True, exist_ok=True)
+    save_file({k: v.detach().contiguous().cpu() for k, v in head.state_dict().items()},
+              str(folder / "schema_head.safetensors"))
+    cfg = {"head": head.cfg.to_dict(), **(extra or {})}
+    tmp = folder / "schema_head_config.json.tmp"
+    tmp.write_text(json.dumps(cfg, indent=2) + "\n")
+    os.replace(tmp, folder / "schema_head_config.json")
+
+
+def load_head_weights(head, folder):
+    from safetensors.torch import load_file
+    state = load_file(str(Path(folder) / "schema_head.safetensors"))
+    missing, unexpected = head.load_state_dict(state, strict=True), None
+    return missing

code/render.py

@@ -0,0 +1,424 @@
+"""Template `schema-v1`: render a schema request into one user turn and locate every question/option span.
+
+Design: request rendering for the joint schema head.
+
+A *schema request* (the internal, normalised form used everywhere in this package):
+
+    {"state": str | dict | list,
+     "images": [data URL, ...],                       # optional
+     "questions": {key: {"type": "choice" | "noul" | "score",
+                         "instructions": str | dict | list,
+                         "options": [{"name": str, "description": str | dict | list | None}, ...]}}}
+
+Options are always stored in CANONICAL order: choice = caller order, noul = [true, false], score = levels 0..K-1.
+Rendering may show them in another order (augmentation); spans are always returned per canonical index.
+
+Encoding (``Encoder.encode``) returns token ids plus token spans. Spans are half-open [start, end) token indices.
+"""
+import hashlib
+import json
+import math
+import random
+
+TEMPLATE_ID = "schema-v1"
+MAX_LENGTH = 32768
+MAX_OPTIONS = 255
+MAX_QUESTIONS = 256
+MAX_SCORE_LEVELS = 10
+QTYPES = ("choice", "noul", "score")
+# Token roles used by the head's role embedding.
+ROLE_OTHER, ROLE_STATE, ROLE_PREVIEW, ROLE_QUESTION, ROLE_OPTION = 0, 1, 2, 3, 4
+N_ROLES = 5
+
+# Wording variants. W0 is canonical (eval always uses W0); W1..W3 are augmentation only.
+WORDINGS = {
+    "W0": {"preamble": "Read the state and answer every question. Each question lists its possible answers.",
+           "preview": "Questions to answer:", "state": "State:", "questions": "Questions:",
+           "choice": "(choose one)", "noul": "(true or false)", "score": "(score 0 to {top})", "bullet": "- "},
+    "W1": {"preamble": "Answer each question below using only the information provided. Every question lists "
+                       "the answers it allows.",
+           "preview": "You will be asked:", "state": "Context:", "questions": "Answer these:",
+           "choice": "[pick one]", "noul": "[true/false]", "score": "[rate 0-{top}]", "bullet": "* "},
+    "W2": {"preamble": "Use the input to decide every question. Pick exactly one of the listed answers for each.",
+           "preview": "Questions:", "state": "Input:", "questions": "Decide:",
+           "choice": "(one of)", "noul": "(true/false)", "score": "(scale 0..{top})", "bullet": "- "},
+    "W3": {"preamble": "Below is some material followed by questions. Choose one listed answer per question.",
+           "preview": "Asked below:", "state": "Material:", "questions": "Questions and answers:",
+           "choice": "[choose one]", "noul": "[true or false]", "score": "[score from 0 to {top}]", "bullet": "* "},
+}
+
+
+class EncodeError(ValueError):
+    """A request that cannot be encoded; `reason` is a short machine-readable code (drop statistics)."""
+
+    def __init__(self, reason, message=""):
+        super().__init__(f"{reason}: {message}" if message else reason)
+        self.reason = reason
+
+
+def require(condition, reason, message=""):
+    if not condition:
+        raise EncodeError(reason, message)
+
+
+def render_value(value):
+    """Plain strings verbatim; structured values as sorted, indented JSON (= jev-adapter render_native_state_text)."""
+    if value is None:
+        return ""
+    if isinstance(value, str):
+        return value
+    return json.dumps(value, sort_keys=True, ensure_ascii=False, indent=2, allow_nan=False)
+
+
+# ----------------------------------------------------------------------------------------------- validation
+def validate_request(req):
+    """Structural checks of a schema request (raises EncodeError). Returns the list of question keys."""
+    require(isinstance(req, dict) and isinstance(req.get("questions"), dict), "bad_request", "questions missing")
+    keys = list(req["questions"])
+    require(1 <= len(keys) <= MAX_QUESTIONS, "bad_question_count", str(len(keys)))
+    for key in keys:
+        require(isinstance(key, str) and key.strip() and "\n" not in key and "[" not in key and "]" not in key,
+                "bad_question_key", repr(key))
+        q = req["questions"][key]
+        t = q.get("type")
+        require(t in QTYPES, "bad_type", repr(t))
+        opts = q.get("options")
+        require(isinstance(opts, list), "bad_options", key)
+        names = [o.get("name") for o in opts]
+        require(all(isinstance(n, str) and n.strip() and "\n" not in n for n in names), "bad_option_name", key)
+        require(len(set(names)) == len(names), "duplicate_option_name", key)
+        if t == "choice":
+            require(2 <= len(opts) <= MAX_OPTIONS, "bad_option_count", f"{key}: {len(opts)}")
+        elif t == "noul":
+            require(names == ["true", "false"], "bad_noul_options", f"{key}: {names}")
+        else:
+            require(2 <= len(opts) <= MAX_SCORE_LEVELS, "bad_score_levels", f"{key}: {len(opts)}")
+            require(names == [str(i) for i in range(len(opts))], "bad_score_names", key)
+    return keys
+
+
+# ----------------------------------------------------------------------------------------------- rendering
+def render(req, wording="W0", question_order=None, option_orders=None, preview=True):
+    """Render the user text. Returns (content, segments).
+
+    question_order: list of keys in display order (default: request order).
+    option_orders: {key: [canonical index shown at display position 0, 1, ...]} (default identity; score must be
+                   identity because levels are ordinal).
+    segments: {"state": (c0, c1), "preview": {key: (c0, c1)}, "question": {key: (c0, c1)},
+               "option": {key: [(c0, c1) per CANONICAL option index]}} as character offsets into `content`.
+    The state segment includes its header line so it is never empty.
+    """
+    keys = validate_request(req)
+    order = list(question_order) if question_order is not None else keys
+    require(sorted(order) == sorted(keys), "bad_question_order")
+    w = WORDINGS[wording]
+    parts, pos = [], 0
+    segs = {"state": None, "preview": {}, "question": {}, "option": {}}
+
+    def emit(text):
+        nonlocal pos
+        start = pos
+        parts.append(text)
+        pos += len(text)
+        return start, pos
+
+    emit(w["preamble"] + "\n\n")
+    if preview:
+        emit(w["preview"] + "\n")
+        for key in order:
+            text = f"[{key}] {render_value(req['questions'][key]['instructions'])}"
+            segs["preview"][key] = emit(text)
+            emit("\n")
+        emit("\n")
+    s0, _ = emit(w["state"] + "\n")
+    _, s1 = emit(render_value(req.get("state", "")))
+    segs["state"] = (s0, s1)
+    emit("\n\n" + w["questions"])
+    for key in order:
+        q = req["questions"][key]
+        n = len(q["options"])
+        tag = w[q["type"]].format(top=n - 1)
+        emit("\n")
+        segs["question"][key] = emit(f"[{key}] {tag} {render_value(q['instructions'])}")
+        perm = list(range(n)) if not option_orders or key not in option_orders else list(option_orders[key])
+        require(sorted(perm) == list(range(n)), "bad_option_order", key)
+        require(q["type"] != "score" or perm == list(range(n)), "score_order_must_be_identity", key)
+        spans = [None] * n
+        for canonical in perm:
+            o = q["options"][canonical]
+            desc = render_value(o.get("description"))
+            text = o["name"] if not desc else f"{o['name']}: {desc}"
+            emit("\n" + w["bullet"])
+            spans[canonical] = emit(text)
+        segs["option"][key] = spans
+    return "".join(parts), segs
+
+
+# ----------------------------------------------------------------------------------------------- augmentation
+def _draw(row_id, epoch, what):
+    return int(hashlib.sha256(f"{row_id}|{epoch}|{what}|schema-aug-v1".encode()).hexdigest()[:16], 16)
+
+
+def augmentation(req, row_id, epoch, p_choice=1.0, p_noul=0.5, p_field=1.0, p_wording=0.5):
+    """Deterministic per (row id, epoch) augmentation plan: wording, question order, option orders."""
+    keys = list(req["questions"])
+    frac = lambda what: _draw(row_id, epoch, what) / float(1 << 64)  # noqa: E731
+    wording = "W0"
+    if frac("wording") < p_wording:
+        wording = ("W1", "W2", "W3")[_draw(row_id, epoch, "wording-pick") % 3]
+    order = list(keys)
+    if len(keys) > 1 and frac("fields") < p_field:
+        random.Random(_draw(row_id, epoch, "field-perm")).shuffle(order)
+    option_orders = {}
+    for key in keys:
+        q = req["questions"][key]
+        n = len(q["options"])
+        p = p_choice if q["type"] == "choice" else p_noul if q["type"] == "noul" else 0.0
+        if p > 0 and frac("opt-gate|" + key) < p:
+            perm = list(range(n))
+            random.Random(_draw(row_id, epoch, "opt-perm|" + key)).shuffle(perm)
+            option_orders[key] = perm
+    return {"wording": wording, "question_order": order, "option_orders": option_orders}
+
+
+# ----------------------------------------------------------------------------------------------- encoding
+class Encoder:
+    """Request -> token ids + token spans, with the trainer's chat-template call shape.
+
+    tokenizer: the canonical tokenizer (AutoTokenizer as the trainer loads it); used for apply_chat_template and
+               the round-trip check.
+    fast: a `tokenizers.Tokenizer` loaded from the same tokenizer.json (character offsets).
+    processor: AutoProcessor, only needed for requests with images.
+    """
+
+    def __init__(self, tokenizer, fast, processor=None, max_length=MAX_LENGTH, image_token_id=None,
+                 image_end_id=None, check_roundtrip=True):
+        self.tokenizer, self.fast, self.processor = tokenizer, fast, processor
+        self.max_length = max_length
+        self.check_roundtrip = check_roundtrip
+        self.image_token_id = image_token_id
+        self.image_end_id = image_end_id
+        self.special_ids = set(getattr(tokenizer, "all_special_ids", []) or [])
+
+    @classmethod
+    def from_snapshot(cls, snapshot, with_processor=True, max_length=MAX_LENGTH, mistral_regex_fix=True):
+        import transformers
+        from tokenizers import Tokenizer
+        kwargs = {"fix_mistral_regex": True} if mistral_regex_fix else {}
+        tok = transformers.AutoTokenizer.from_pretrained(str(snapshot), local_files_only=True,
+                                                         trust_remote_code=False, token=False, **kwargs)
+        fast = Tokenizer.from_file(str(snapshot) + "/tokenizer.json")
+        proc = img = img_end = None
+        if with_processor:
+            proc = transformers.AutoProcessor.from_pretrained(str(snapshot), local_files_only=True,
+                                                              trust_remote_code=False, token=False, **kwargs)
+            img = tok.convert_tokens_to_ids(proc.image_token)
+            img_end = tok.convert_tokens_to_ids(proc.image_end_token)
+        enc = cls(tok, fast, proc, max_length, img, img_end)
+        enc.snapshot = str(snapshot)
+        return enc
+
+    # -- helpers
+    def _chat_text(self, content, n_images):
+        if n_images:
+            messages = [{"role": "user", "content": [{"type": "image"} for _ in range(n_images)]
+                         + [{"type": "text", "text": content}]}]
+            # no enable_thinking here: the Ministral template has no such variable, and the processor's
+            # apply_chat_template treats unknown kwargs as processor kwargs (warning on every image encode). The text
+            # is identical either way (tests/test_image_batch.py::test_processor_kwargs_accepted).
+            return self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True), messages
+        messages = [{"role": "user", "content": content}]
+        return self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True,
+                                                  enable_thinking=False), messages
+
+    def encode(self, req, wording="W0", question_order=None, option_orders=None, preview=True, pil_images=None):
+        """Returns a dict (all plain Python ints/lists):
+        input_ids, n_tokens, keys (rendered question order), qtype, n_opt,
+        q_span[i] = (s, e) question line, p_span[i] = (s, e) preview line or None,
+        o_span[i] = [(s, e) per CANONICAL option], o_display[i] = canonical index per display position,
+        state_spans = [(s, e), ...] (state text incl. header, plus the image block), roles = run-length
+        [(role, start, end)], image info."""
+        content, segs = render(req, wording, question_order, option_orders, preview)
+        images = req.get("images") or []
+        chat, _messages = self._chat_text(content, len(images))
+        require(chat.count(content) == 1, "content_not_unique_in_chat")
+        c0 = chat.index(content)
+        enc = self.fast.encode(chat, add_special_tokens=False)
+        ids, offsets = list(enc.ids), list(enc.offsets)
+        if self.check_roundtrip and not images:
+            ref = self.tokenizer.apply_chat_template([{"role": "user", "content": content}], tokenize=True,
+                                                     add_generation_prompt=True, enable_thinking=False,
+                                                     return_dict=False)
+            require(list(ref) == ids, "roundtrip_mismatch")
+        image_span = None
+        pixel = None
+        if images:
+            require(self.processor is not None, "no_processor")
+            img, img_end = self.image_token_id, self.image_end_id
+            require(ids.count(img) == len(images), "image_placeholder_count")
+            if pil_images is None:
+                pil_images = decode_images(images)
+            # flat add_special_tokens: _merge_kwargs routes it to the tokenizer (no extra BOS; checked below and in
+            # tests/test_image_batch.py)
+            out = self.processor(text=chat, images=pil_images, return_tensors="pt", add_special_tokens=False)
+            expanded = out["input_ids"][0].tolist()
+            plain_ids, plain_off = ids, offsets
+            first_plain = plain_ids.index(img)
+            last_plain = len(plain_ids) - 1 - plain_ids[::-1].index(img)
+            require(plain_off[last_plain][1] <= c0, "images_not_before_text")
+            require(img in expanded and img_end in expanded, "no_image_tokens_after_processor")
+            first_exp = expanded.index(img)
+            last_exp = len(expanded) - 1 - expanded[::-1].index(img_end)
+            require(expanded[:first_exp] == plain_ids[:first_plain], "image_prefix_mismatch")
+            require(expanded[last_exp + 1:] == plain_ids[last_plain + 1:], "image_suffix_mismatch")
+            shift = last_exp - last_plain
+            image_span = (first_exp, last_exp + 1)
+            ids = expanded
+            offsets = [None] * len(expanded)
+            for i in range(first_plain):
+                offsets[i] = plain_off[i]
+            for i in range(last_plain + 1, len(plain_ids)):
+                offsets[i + shift] = plain_off[i]
+            pixel = {"pixel_values": out["pixel_values"], "image_sizes": out.get("image_sizes")}
+        n = len(ids)
+        require(1 <= n <= self.max_length, "too_long", str(n))
+
+        # char -> segment id map over the content region
+        seg_names = []
+        seg_ranges = []
+
+        def add(name, rng):
+            seg_names.append(name)
+            seg_ranges.append((rng[0] + c0, rng[1] + c0))
+            return len(seg_names) - 1
+
+        keys = list(question_order) if question_order is not None else list(req["questions"])
+        sid_state = add(("state",), segs["state"])
+        sid_prev, sid_q, sid_o = {}, {}, {}
+        for key in keys:
+            if key in segs["preview"]:
+                sid_prev[key] = add(("preview", key), segs["preview"][key])
+            sid_q[key] = add(("question", key), segs["question"][key])
+            sid_o[key] = [add(("option", key, j), r) for j, r in enumerate(segs["option"][key])]
+        # sort segment ranges for binary search
+        order = sorted(range(len(seg_ranges)), key=lambda i: seg_ranges[i][0])
+        starts = [seg_ranges[i][0] for i in order]
+        import bisect
+
+        def seg_of(ch):
+            k = bisect.bisect_right(starts, ch) - 1
+            if k < 0:
+                return -1
+            sid = order[k]
+            a, b = seg_ranges[sid]
+            return sid if a <= ch < b else -1
+
+        token_seg = [-1] * n
+        content_end = c0 + len(content)
+        for t, off in enumerate(offsets):
+            if off is None:
+                continue
+            a, b = off
+            if b <= c0 or a >= content_end:
+                continue
+            ch = a
+            while ch < b and chat[ch].isspace():
+                ch += 1
+            if ch >= b:
+                ch = a
+            token_seg[t] = seg_of(ch)
+            # special tokens must not appear inside the caller's content (e.g. a key spelled like a control token)
+            require(ids[t] not in self.special_ids, "special_token_in_content", str(ids[t]))
+
+        # spans per segment + coverage check (every non-space char of a segment lies in one of its tokens)
+        spans = {}
+        for t, sid in enumerate(token_seg):
+            if sid < 0:
+                continue
+            s = spans.get(sid)
+            spans[sid] = (t, t + 1) if s is None else (s[0], t + 1)
+        for sid, (a, b) in spans.items():
+            require(all(token_seg[t] == sid for t in range(a, b)), "span_not_contiguous", str(seg_names[sid]))
+        for sid, (ca, cb) in enumerate(seg_ranges):
+            require(sid in spans, "empty_span", str(seg_names[sid]))
+            a, b = spans[sid]
+            covered_lo = offsets[a][0]
+            covered_hi = offsets[b - 1][1]
+            first = ca
+            while first < cb and chat[first].isspace():
+                first += 1
+            last = cb
+            while last > first and chat[last - 1].isspace():
+                last -= 1
+            require(covered_lo <= first and covered_hi >= last, "span_straddle", str(seg_names[sid]))
+
+        state_spans = [spans[sid_state]]
+        if image_span is not None:
+            state_spans.insert(0, image_span)
+        out = {"template": TEMPLATE_ID, "wording": wording, "input_ids": ids, "n_tokens": n, "keys": keys,
+               "qtype": [req["questions"][k]["type"] for k in keys],
+               "n_opt": [len(req["questions"][k]["options"]) for k in keys],
+               "q_span": [spans[sid_q[k]] for k in keys],
+               "p_span": [spans[sid_prev[k]] if k in sid_prev else None for k in keys],
+               "o_span": [[spans[s] for s in sid_o[k]] for k in keys],
+               "o_display": [list(option_orders[k]) if option_orders and k in option_orders
+                             else list(range(len(req["questions"][k]["options"]))) for k in keys],
+               "state_spans": state_spans, "n_images": len(images), "image_tokens": (image_span[1] - image_span[0]
+                                                                                     if image_span else 0)}
+        if pixel is not None:
+            out["_pixel"] = pixel
+        return out
+
+
+def decode_images(data_urls):
+    import base64
+    import io
+    from PIL import Image
+    out = []
+    for url in data_urls:
+        require(isinstance(url, str) and url.startswith("data:image/") and "," in url, "bad_image")
+        out.append(Image.open(io.BytesIO(base64.b64decode(url.split(",", 1)[1]))).convert("RGB"))
+    return out
+
+
+def token_roles(encoded):
+    """Per-token role ids (list of length n_tokens)."""
+    roles = [ROLE_OTHER] * encoded["n_tokens"]
+    for a, b in encoded["state_spans"]:
+        for t in range(a, b):
+            roles[t] = ROLE_STATE
+    for i in range(len(encoded["keys"])):
+        if encoded["p_span"][i]:
+            a, b = encoded["p_span"][i]
+            for t in range(a, b):
+                roles[t] = ROLE_PREVIEW
+        a, b = encoded["q_span"][i]
+        for t in range(a, b):
+            roles[t] = ROLE_QUESTION
+        for a, b in encoded["o_span"][i]:
+            for t in range(a, b):
+                roles[t] = ROLE_OPTION
+    return roles
+
+
+def request_from_systemone(body):
+    """/v1/systemone request JSON (jev-adapter protocol) -> schema request (canonical option order)."""
+    questions = {}
+    for key, q in body["questions"].items():
+        t = q["type"]
+        if t == "choice":
+            opts = [{"name": n, "description": d} for n, d in q["criteria"].items()]
+        elif t == "noul":
+            crit = q.get("criteria") or {}
+            opts = [{"name": "true", "description": crit.get("true")}, {"name": "false", "description": crit.get("false")}]
+        elif t == "score":
+            opts = [{"name": str(i), "description": d} for i, d in enumerate(q["criteria"])]
+        else:
+            raise EncodeError("bad_type", repr(t))
+        questions[key] = {"type": t, "instructions": q["instructions"], "options": opts}
+    return {"state": body.get("state", ""), "images": list(body.get("images") or []), "questions": questions}
+
+
+def entropy_confidence(p):
+    h = -sum(x * math.log(x) for x in p if x > 0)
+    return min(1.0, max(0.0, 1.0 - h / math.log(len(p))))

code/requirements.txt

@@ -0,0 +1,4 @@
+torch>=2.4
+transformers==5.12.1
+peft==0.21.0
+safetensors

code/scoring.py

@@ -0,0 +1,183 @@
+"""In-process /v1/systemone scoring with the joint schema head (shared by serve_head.py and eval_head.py).
+
+Response JSON follows the jev-adapter contract (protocol.reduce_probabilities / service.score), so the unchanged
+runners (for example `jev_adapter.benchmarks.run`) work as is:
+  answers[q] = {type, probabilities | noul, choice, confidence, score, legend, logprobs?}
+  usage = {input_tokens, output_tokens: 0}
+  metadata.evaluations = questions x permutations (the runners' per-question accounting; the head needs only
+  `prefills` = permutations forward passes for the whole request), plus prefill_ms / head_ms.
+Option permutations (options.permutations = K) rotate choice and noul options (display position p shows original
+(p + r) mod n, as the adapter's rotation_order); score levels always stay in ascending order.
+"""
+import json
+import math
+import time
+from pathlib import Path
+
+import torch
+
+from render import EncodeError, entropy_confidence, request_from_systemone, token_roles
+
+CONFIDENCE_METHOD = "1 - normalized_entropy"
+
+
+class RequestError(ValueError):
+    def __init__(self, code, message, field=None, status=422):
+        super().__init__(message)
+        self.code, self.field, self.status = code, field, status
+
+
+def validate_body(body):
+    """Light structural validation (the adapter's pydantic schema is used when importable)."""
+    try:
+        from jev_adapter.protocol import SystemOneRequest  # noqa: F401
+        SystemOneRequest.model_validate(body)
+    except ImportError:
+        pass
+    except Exception as e:  # pydantic ValidationError
+        errors = getattr(e, "errors", lambda: [{"msg": str(e), "loc": ()}])()
+        first = errors[0] if errors else {"msg": str(e), "loc": ()}
+        raise RequestError("invalid_request", first.get("msg", str(e)),
+                           ".".join(str(p) for p in first.get("loc", ()))) from None
+    if not isinstance(body, dict) or not isinstance(body.get("questions"), dict) or not body["questions"]:
+        raise RequestError("invalid_request", "questions are required", "questions")
+
+
+class HeadScorer:
+    def __init__(self, model, encoder, served_model, pad_id, temperature_by_type=None, max_concurrency=1,
+                 max_batch_tokens=32768, preview=True):
+        self.model, self.encoder, self.served_model, self.pad_id = model, encoder, served_model, pad_id
+        self.preview = bool(preview)   # must match training (head config "preview"; --no-preview ablation)
+        self.temperature_by_type = dict(temperature_by_type or {})
+        self.max_concurrency = max_concurrency
+        self.max_batch_tokens = max_batch_tokens
+        self.device = model.device()
+
+    def plan(self, body):
+        validate_body(body)
+        req = request_from_systemone(body)
+        opts = body.get("options") or {}
+        k = int(opts.get("permutations", 1))
+        encs = []
+        for r in range(k):
+            orders = {}
+            for key, q in req["questions"].items():
+                n = len(q["options"])
+                if q["type"] != "score" and r:
+                    orders[key] = [(p + r) % n for p in range(n)]
+            try:
+                encs.append(self.encoder.encode(req, "W0", list(req["questions"]), orders, preview=self.preview))
+            except EncodeError as e:
+                raise RequestError(e.reason, str(e), "questions" if e.reason != "too_long" else "state") from None
+        return req, opts, encs
+
+    @torch.no_grad()
+    def run(self, encs):
+        """Forward the encodings (one padded micro-batch); returns (per-encoding logits lists, prefill_s, head_s)."""
+        from model import make_batch
+        dev = self.device
+        sync = (lambda: torch.cuda.synchronize()) if dev.type == "cuda" else (lambda: None)
+        batch = make_batch(encs, self.pad_id, dev)
+        sync()
+        t0 = time.perf_counter()
+        states = self.model.backbone_states(batch)
+        sync()
+        t1 = time.perf_counter()
+        outs = []
+        with torch.autocast(device_type=dev.type, enabled=False):
+            for i, enc in enumerate(encs):
+                n = enc["n_tokens"]
+                taps = [s[i, :n].float() for s in states]
+                roles = torch.tensor(token_roles(enc), device=dev)
+                outs.append([z.float().cpu() for z in self.model.head.forward_one(taps, enc, roles)])
+        sync()
+        t2 = time.perf_counter()
+        return outs, t1 - t0, t2 - t1
+
+    def respond(self, req, opts, encs, outs, prefill_s, head_s, started):
+        temperature = float(opts.get("temperature", 1.0)) if opts.get("temperature_scaling", True) else 1.0
+        return_logprobs = bool(opts.get("return_logprobs", False))
+        answers = {}
+        keys = list(req["questions"])
+        for key in keys:
+            q = req["questions"][key]
+            names = [o["name"] for o in q["options"]]
+            t_type = self.temperature_by_type.get(q["type"], 1.0)
+            probs_sum = [0.0] * len(names)
+            raw = []
+            for enc, out in zip(encs, outs):
+                z = out[enc["keys"].index(key)].double()
+                logp = torch.log_softmax(z, -1)
+                raw.append({n: float(v) for n, v in zip(names, logp.tolist())})
+                p = torch.softmax(z / (t_type * temperature), -1).tolist()
+                for j, v in enumerate(p):
+                    probs_sum[j] += v
+            total = math.fsum(probs_sum)
+            probs = [v / total for v in probs_sum]
+            ans = {"type": q["type"]}
+            if q["type"] == "noul":
+                ans["noul"] = probs[0]
+            else:
+                ans["probabilities"] = dict(zip(names, probs))
+                ans["confidence"] = entropy_confidence(probs)
+                if q["type"] == "choice":
+                    ans["choice"] = names[max(range(len(names)), key=probs.__getitem__)]
+                else:
+                    ans["score"] = math.fsum(i * p for i, p in enumerate(probs))
+                    ans["legend"] = {o["name"]: o.get("description") for o in q["options"]}
+            if return_logprobs:
+                ans["logprobs"] = raw
+            answers[key] = ans
+        input_tokens = sum(e["n_tokens"] for e in encs)
+        usage = {"input_tokens": input_tokens, "output_tokens": 0}
+        if req["images"]:
+            usage["input_tokens_details"] = {"image_tokens": sum(e["image_tokens"] for e in encs)}
+        return {"model": self.served_model, "answers": answers, "usage": usage,
+                "metadata": {"confidence_method": CONFIDENCE_METHOD, "temperature": temperature,
+                             **({"temperature_by_type": self.temperature_by_type} if self.temperature_by_type else {}),
+                             "evaluations": len(keys) * len(encs), "prefills": len(encs),
+                             "service_max_concurrency": self.max_concurrency, "label_scheme": "schema-v1",
+                             "cached_tokens": 0, "prefill_ms": prefill_s * 1000, "head_ms": head_s * 1000,
+                             "adapter_elapsed_ms": (time.perf_counter() - started) * 1000}}
+
+    def score(self, body):
+        started = time.perf_counter()
+        if body.get("model") not in (None, self.served_model, "jev-latest"):
+            raise RequestError("model_not_found", "Requested model is not configured.", "model", 404)
+        req, opts, encs = self.plan(body)
+        outs, ps, hs = self.run(encs)
+        return self.respond(req, opts, encs, outs, ps, hs, started)
+
+
+def load_scorer(snapshot=None, head_dir=None, adapter=None, tiny=None, tokenizer=None, served_model=None,
+                temperature_file=None, device="auto", merge=True, max_concurrency=1, memory_cap_bytes=0):
+    """Backbone (+ LoRA adapter merged) + trained head -> HeadScorer.
+    memory_cap_bytes > 0 caps the torch allocator (set_per_process_memory_fraction) before the model is loaded."""
+    from model import build_model, load_head_weights
+    from render import Encoder
+    dev = torch.device(("cuda" if torch.cuda.is_available() else "cpu") if device == "auto" else device)
+    if dev.type == "cuda" and memory_cap_bytes:
+        total = torch.cuda.get_device_properties(0).total_memory
+        torch.cuda.set_per_process_memory_fraction(min(1.0, float(memory_cap_bytes) / total), 0)
+    cfg = json.loads((Path(head_dir) / "schema_head_config.json").read_text()) if head_dir else {}
+    hc = cfg.get("head", {})
+    taps = cfg.get("taps") or hc.get("taps") or ["final", 25]
+    overrides = {k: hc[k] for k in ("d", "heads", "ffn", "route_blocks", "interact_blocks", "ord_hidden", "max_nodes")
+                 if k in hc}
+    overrides["dropout"] = 0.0
+    overrides["grad_checkpoint"] = False
+    model, info = build_model(snapshot, tiny=tiny, lora_rank=0, taps=taps, device=dev, grad_checkpointing=False,
+                              init_adapter=adapter, head_overrides=overrides)
+    if adapter and merge and model.peft_model is not None:
+        model.backbone = model.peft_model.merge_and_unload()
+        model.peft_model = None
+    if head_dir:
+        load_head_weights(model.head, head_dir)
+    model.eval()
+    if dev.type == "cuda":
+        model.backbone.to(torch.bfloat16)
+    encoder = Encoder.from_snapshot(tokenizer or snapshot, with_processor=True)
+    pad = encoder.tokenizer.pad_token_id if encoder.tokenizer.pad_token_id is not None else encoder.tokenizer.eos_token_id
+    temps = json.loads(Path(temperature_file).read_text()).get("temperature_by_type") if temperature_file else None
+    return HeadScorer(model, encoder, served_model or "standardthinking/standard-schema-8b", pad, temps,
+                      max_concurrency=max_concurrency, preview=cfg.get("preview", True)), info

code/serve_head.py

@@ -0,0 +1,156 @@
+#!/usr/bin/env python3
+"""HTTP server implementing the jev-adapter /v1/systemone contract with the joint schema head (DESIGN.md 5.1).
+
+  python3 serve_head.py --snapshot <base or merged dir> --adapter <run>/final/adapter --head <run>/final/head \
+      --port 30171 [--temperature-file calib.json]
+
+Endpoints: POST /v1/systemone, GET /v1/models, /model_info, /server_info (also /get_model_info, /get_server_info),
+/health. The engine snapshot fields the benchmark runner checks are reported (disable_radix_cache=true,
+mm_preprocess_cache_size_mb=0, speculative_algorithm=null): every request is a full prefill, nothing is cached.
+Dynamic batching: requests waiting in the queue are forwarded together while the padded token count stays within
+--max-batch-tokens (image requests run alone). Stdlib HTTP server; binds 127.0.0.1 by default.
+"""
+import argparse
+import json
+import queue
+import sys
+import threading
+import time
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from pathlib import Path
+
+HERE = Path(__file__).resolve().parent
+sys.path.insert(0, str(HERE))
+
+from scoring import RequestError, load_scorer  # noqa: E402
+
+
+class Batcher:
+    def __init__(self, scorer, max_batch_tokens, max_batch_requests, wait_ms):
+        self.scorer, self.max_tokens, self.max_requests, self.wait = scorer, max_batch_tokens, max_batch_requests, wait_ms / 1000
+        self.q = queue.Queue()
+        threading.Thread(target=self.loop, daemon=True).start()
+
+    def submit(self, body):
+        started = time.perf_counter()
+        req, opts, encs = self.scorer.plan(body)        # validation + encoding in the caller's thread
+        job = {"req": req, "opts": opts, "encs": encs, "started": started, "done": threading.Event()}
+        self.q.put(job)
+        job["done"].wait()
+        if "error" in job:
+            raise job["error"]
+        return job["response"]
+
+    def loop(self):
+        while True:
+            jobs = [self.q.get()]
+            image = bool(jobs[0]["req"]["images"])
+            deadline = time.perf_counter() + self.wait
+            while not image and len(jobs) < self.max_requests:
+                try:
+                    nxt = self.q.get(timeout=max(0.0, deadline - time.perf_counter()))
+                except queue.Empty:
+                    break
+                encs = [e for j in jobs + [nxt] for e in j["encs"]]
+                cost = len(encs) * max(e["n_tokens"] for e in encs)
+                if nxt["req"]["images"] or cost > self.max_tokens:
+                    self.q.put(nxt)   # goes to the next batch
+                    break
+                jobs.append(nxt)
+            try:
+                encs = [e for j in jobs for e in j["encs"]]
+                outs, ps, hs = self.scorer.run(encs)
+                k = 0
+                for j in jobs:
+                    n = len(j["encs"])
+                    j["response"] = self.scorer.respond(j["req"], j["opts"], j["encs"], outs[k:k + n], ps, hs, j["started"])
+                    j["response"]["metadata"]["batch_requests"] = len(jobs)
+                    k += n
+            except Exception as e:  # noqa: BLE001
+                for j in jobs:
+                    j["error"] = e
+            for j in jobs:
+                j["done"].set()
+
+
+def make_handler(scorer, batcher, info):
+    model_info = {"model_path": info["model_path"], "served_model_name": scorer.served_model, "is_generation": False,
+                  "has_image_understanding": True, "model_type": "schema-head", "architectures": ["SchemaHeadModel"]}
+    server_info = {"model_path": info["model_path"], "served_model_name": scorer.served_model,
+                   "disable_radix_cache": True, "mm_preprocess_cache_size_mb": 0, "enable_prefix_mm_cache": False,
+                   "enable_mm_global_cache": False, "speculative_algorithm": None, "context_length": 32768,
+                   "dtype": info["dtype"], "attention_backend": "torch-sdpa", "version": "schema-head-serve-v1"}
+
+    class Handler(BaseHTTPRequestHandler):
+        protocol_version = "HTTP/1.1"
+
+        def log_message(self, fmt, *args):
+            pass
+
+        def _send(self, code, obj):
+            data = json.dumps(obj, ensure_ascii=False, allow_nan=False).encode()
+            self.send_response(code)
+            self.send_header("Content-Type", "application/json")
+            self.send_header("Content-Length", str(len(data)))
+            self.end_headers()
+            self.wfile.write(data)
+
+        def do_GET(self):
+            path = self.path.split("?")[0].rstrip("/")
+            if path == "/health":
+                return self._send(200, {"status": "ok"})
+            if path == "/v1/models":
+                return self._send(200, {"object": "list", "data": [{"id": scorer.served_model, "object": "model",
+                                                                     "owned_by": "schema-head", "label_scheme": "schema-v1"}]})
+            if path in ("/model_info", "/get_model_info"):
+                return self._send(200, model_info)
+            if path in ("/server_info", "/get_server_info"):
+                return self._send(200, server_info)
+            return self._send(404, {"error": {"code": "not_found", "message": path}})
+
+        def do_POST(self):
+            if self.path.split("?")[0].rstrip("/") != "/v1/systemone":
+                return self._send(404, {"error": {"code": "not_found", "message": self.path}})
+            try:
+                body = json.loads(self.rfile.read(int(self.headers.get("Content-Length", "0"))))
+                return self._send(200, batcher.submit(body))
+            except RequestError as e:
+                return self._send(e.status, {"error": {"code": e.code, "message": str(e), "field": e.field}})
+            except json.JSONDecodeError as e:
+                return self._send(400, {"error": {"code": "invalid_json", "message": str(e)}})
+            except Exception as e:  # noqa: BLE001
+                return self._send(500, {"error": {"code": "scoring_failed", "message": f"{type(e).__name__}: {e}"}})
+    return Handler
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--snapshot", type=Path)
+    ap.add_argument("--tokenizer", type=Path)
+    ap.add_argument("--tiny", help="JSON tiny config (tests only)")
+    ap.add_argument("--adapter", type=Path)
+    ap.add_argument("--head", type=Path, required=True)
+    ap.add_argument("--served-model-name", default="standardthinking/standard-schema-8b")
+    ap.add_argument("--temperature-file", type=Path)
+    ap.add_argument("--host", default="127.0.0.1")
+    ap.add_argument("--port", type=int, default=30171)
+    ap.add_argument("--device", default="auto")
+    ap.add_argument("--max-batch-tokens", type=int, default=32768)
+    ap.add_argument("--max-batch-requests", type=int, default=8)
+    ap.add_argument("--batch-wait-ms", type=float, default=5.0)
+    ap.add_argument("--memory-cap-bytes", type=float, default=0, help="torch allocator cap on GPU (0 = none)")
+    args = ap.parse_args()
+    tiny = json.loads(args.tiny) if args.tiny else None
+    scorer, info = load_scorer(args.snapshot, args.head, args.adapter, tiny, args.tokenizer, args.served_model_name,
+                               args.temperature_file, args.device, max_concurrency=args.max_batch_requests,
+                               memory_cap_bytes=args.memory_cap_bytes)
+    batcher = Batcher(scorer, args.max_batch_tokens, args.max_batch_requests, args.batch_wait_ms)
+    meta = {"model_path": str(args.snapshot or "tiny"), "dtype": "bfloat16" if scorer.device.type == "cuda" else "float32"}
+    server = ThreadingHTTPServer((args.host, args.port), make_handler(scorer, batcher, meta))
+    print(json.dumps({"serving": f"http://{args.host}:{args.port}", "model": args.served_model_name,
+                      "head_params": info["head_params"], "device": str(scorer.device)}), flush=True)
+    server.serve_forever()
+
+
+if __name__ == "__main__":
+    main()

head/schema_head.safetensors

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

head/schema_head_config.json

@@ -0,0 +1,31 @@
+{
+  "head": {
+    "hidden_size": 4096,
+    "n_taps": 2,
+    "d": 1024,
+    "heads": 16,
+    "ffn": 2816,
+    "route_blocks": 2,
+    "interact_blocks": 2,
+    "dropout": 0.1,
+    "n_roles": 5,
+    "ord_hidden": 256,
+    "max_nodes": 4096,
+    "grad_checkpoint": true,
+    "taps": [
+      "final",
+      25
+    ],
+    "template": "schema-v1"
+  },
+  "taps": [
+    "final",
+    25
+  ],
+  "template": "schema-v1",
+  "step": 982,
+  "lora_rank": 64,
+  "snapshot": "StandardThinking/StandardOne-8B",
+  "preview": true,
+  "snapshot_revision": "ac495d1c221297f0ae951941eff27bcf58b5ce52"
+}