Upload folder using huggingface_hub

.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.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,79 @@
+---
+library_name: transformers
+license: gemma
+base_model: google/gemma-4-12B
+pipeline_tag: text-classification
+tags:
+  - custom_code
+  - gemma4
+  - typed-decisions
+---
+
+# Truetype System One (Gemma 4 12B, 26-letter output layer)
+
+`google/gemma-4-12B` with its 262k-token language-model head replaced by a
+26-row layer that outputs one logit per letter `A`-`Z`. The decoder is
+unchanged. One forward pass gives the 26 logits; questions (`noul`, `choice`,
+`score`) are answered by softmax over the letters each question declares legal.
+
+Requires `transformers>=5.17`. The repo ships custom code, so pass
+`trust_remote_code=True`.
+
+## Model API
+
+```python
+from transformers import AutoModelForMultimodalLM, AutoTokenizer
+
+repo = "stephenlb/system-one-model"
+tokenizer = AutoTokenizer.from_pretrained(repo)
+model = AutoModelForMultimodalLM.from_pretrained(
+    repo, trust_remote_code=True, dtype="bfloat16", device_map="auto"
+)
+
+result = model.system_one(
+    tokenizer,
+    state="I was charged twice for order A-104.",
+    questions={
+        "refund": {"type": "noul", "instructions": "Does the text request a refund?"},
+        "team": {
+            "type": "choice",
+            "instructions": "Which team should handle this?",
+            "criteria": {"billing": "Charges", "returns": "Refunds"},
+        },
+    },
+)
+print(result["answers"])
+```
+
+## Pipeline API
+
+```python
+from transformers import pipeline
+
+pipe = pipeline("system-one", model=repo, trust_remote_code=True, dtype="bfloat16", device_map="auto")
+print(pipe({"state": "I was charged twice.", "questions": {...}}))
+```
+
+## Output
+
+```json
+{"answers": {
+  "refund": {"type": "noul", "noul": 0.93},
+  "team": {"type": "choice", "choice": "billing",
+           "probabilities": {"billing": 0.97, "returns": 0.03}, "confidence": 0.8},
+  "urgency": {"type": "score", "score": 3.1, "legend": {"0": "..."},
+              "probabilities": {"0": 0.0, "1": 0.1}, "confidence": 0.5}}}
+```
+
+`noul` is the probability of yes. `choice` returns the top label with
+probabilities per option. `score` returns the expected level over an ordered
+`criteria` list. Pass `temperature=` (default 0.7) to sharpen or soften.
+
+## Raw letter logits
+
+```python
+enc = tokenizer(prompts, return_tensors="pt", padding=True)  # padding_side="left"
+logits = model.letter_logits(enc.input_ids, enc.attention_mask)  # [batch, 26]
+```
+
+Source and prompt format: https://github.com/stephenlb/truetype.ai-open

config.json

@@ -0,0 +1,241 @@
+{
+  "architectures": [
+    "TrueTypeSystemOneModel"
+  ],
+  "audio_config": {
+    "_name_or_path": "",
+    "architectures": null,
+    "audio_embed_dim": 640,
+    "chunk_size_feed_forward": 0,
+    "dtype": "bfloat16",
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "model_type": "gemma4_unified_audio",
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "problem_type": null,
+    "return_dict": true,
+    "rms_norm_eps": 1e-06
+  },
+  "audio_token_id": 258881,
+  "boa_token_id": 256000,
+  "boi_token_id": 255999,
+  "dtype": "bfloat16",
+  "eoa_token_index": 258883,
+  "eoi_token_id": 258882,
+  "image_token_id": 258880,
+  "initializer_range": 0.02,
+  "letter_token_ids": [
+    562,
+    603,
+    565,
+    622,
+    645,
+    633,
+    667,
+    640,
+    564,
+    730,
+    751,
+    639,
+    587,
+    646,
+    708,
+    593,
+    1366,
+    630,
+    555,
+    558,
+    749,
+    782,
+    649,
+    1684,
+    895,
+    1491
+  ],
+  "model_type": "truetype_system_one",
+  "num_letters": 26,
+  "text_config": {
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "attention_k_eq_v": true,
+    "bos_token_id": 2,
+    "dtype": "bfloat16",
+    "enable_moe_block": false,
+    "eos_token_id": 1,
+    "final_logit_softcapping": 30.0,
+    "head_dim": 256,
+    "hidden_activation": "gelu_pytorch_tanh",
+    "hidden_size": 3840,
+    "hidden_size_per_layer_input": 0,
+    "initializer_range": 0.02,
+    "intermediate_size": 15360,
+    "layer_types": [
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention"
+    ],
+    "max_position_embeddings": 262144,
+    "model_type": "gemma4_unified_text",
+    "moe_intermediate_size": null,
+    "num_attention_heads": 16,
+    "num_experts": null,
+    "num_hidden_layers": 48,
+    "num_key_value_heads": 8,
+    "num_kv_shared_layers": 0,
+    "pad_token_id": 0,
+    "per_layer_config": {
+      "05": {
+        "head_dim": 512,
+        "num_key_value_heads": 1
+      },
+      "11": {
+        "head_dim": 512,
+        "num_key_value_heads": 1
+      },
+      "17": {
+        "head_dim": 512,
+        "num_key_value_heads": 1
+      },
+      "23": {
+        "head_dim": 512,
+        "num_key_value_heads": 1
+      },
+      "29": {
+        "head_dim": 512,
+        "num_key_value_heads": 1
+      },
+      "35": {
+        "head_dim": 512,
+        "num_key_value_heads": 1
+      },
+      "41": {
+        "head_dim": 512,
+        "num_key_value_heads": 1
+      },
+      "47": {
+        "head_dim": 512,
+        "num_key_value_heads": 1
+      }
+    },
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+      "full_attention": {
+        "partial_rotary_factor": 0.25,
+        "rope_theta": 1000000.0,
+        "rope_type": "proportional"
+      },
+      "sliding_attention": {
+        "rope_theta": 10000.0,
+        "rope_type": "default"
+      }
+    },
+    "sliding_window": 1024,
+    "tie_word_embeddings": true,
+    "top_k_experts": null,
+    "use_bidirectional_attention": "vision",
+    "use_cache": true,
+    "use_double_wide_mlp": false,
+    "vocab_size": 262144,
+    "vocab_size_per_layer_input": 262144
+  },
+  "tie_word_embeddings": false,
+  "transformers_version": "5.17.0",
+  "video_token_id": 258884,
+  "vision_config": {
+    "_name_or_path": "",
+    "architectures": null,
+    "chunk_size_feed_forward": 0,
+    "dtype": "bfloat16",
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "mm_embed_dim": 3840,
+    "mm_posemb_size": 1120,
+    "model_type": "gemma4_unified_vision",
+    "num_soft_tokens": 280,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_proj_dims": 3840,
+    "patch_size": 16,
+    "pooling_kernel_size": 3,
+    "problem_type": null,
+    "return_dict": true,
+    "rms_norm_eps": 1e-06
+  },
+  "auto_map": {
+    "AutoConfig": "configuration_truetype.TrueTypeSystemOneConfig",
+    "AutoModel": "modeling_truetype.TrueTypeSystemOneModel",
+    "AutoModelForMultimodalLM": "modeling_truetype.TrueTypeSystemOneModel",
+    "AutoModelForImageTextToText": "modeling_truetype.TrueTypeSystemOneModel"
+  },
+  "custom_pipelines": {
+    "system-one": {
+      "impl": "pipeline_truetype.SystemOnePipeline",
+      "pt": [
+        "AutoModelForMultimodalLM"
+      ],
+      "tf": []
+    }
+  }
+}

configuration_truetype.py

@@ -0,0 +1,23 @@
+"""Config for the Gemma 4 System One model with a 26-logit A-Z output layer."""
+
+from transformers.models.gemma4_unified.configuration_gemma4_unified import Gemma4UnifiedConfig
+
+LETTERS = tuple("ABCDEFGHIJKLMNOPQRSTUVWXYZ")
+
+
+class TrueTypeSystemOneConfig(Gemma4UnifiedConfig):
+    """Gemma 4 config plus the tokenizer ids the 26 output rows were copied from.
+
+    ``letter_token_ids[i]`` is the vocabulary id whose embedding row became output
+    row ``i`` (``A`` to ``Z``). It is kept for provenance and for verification
+    against the original vocabulary head.
+    """
+
+    model_type = "truetype_system_one"
+
+    def __init__(self, letter_token_ids=None, num_letters=len(LETTERS), **kwargs):
+        # The 26-row head is its own parameter, not the input embedding.
+        kwargs["tie_word_embeddings"] = False
+        super().__init__(**kwargs)
+        self.letter_token_ids = list(letter_token_ids or [])
+        self.num_letters = num_letters

generation_config.json

@@ -0,0 +1,14 @@
+{
+  "bos_token_id": 2,
+  "do_sample": true,
+  "eos_token_id": 1,
+  "pad_token_id": 0,
+  "suppress_tokens": [
+    258883,
+    258882
+  ],
+  "temperature": 1.0,
+  "top_k": 64,
+  "top_p": 0.95,
+  "transformers_version": "5.17.0"
+}

letters.py

@@ -0,0 +1,138 @@
+"""Extract single-letter logits and compute top-k softmax scores.
+
+The engine reads one next-token position, keeps logits for the 26 uppercase
+ASCII letters, and converts the top-k letter logits into a probability
+distribution. It does not sample or decode text.
+"""
+
+from __future__ import annotations
+
+import math
+import string
+from dataclasses import dataclass
+
+LETTERS: tuple[str, ...] = tuple(string.ascii_uppercase)
+LETTER_INDEX: dict[str, int] = {c: i for i, c in enumerate(LETTERS)}
+
+# Rank at which we stop tracking for the reported top-k. k <= 26 by construction.
+MAX_TOP_K = len(LETTERS)
+
+
+class LetterLogitError(ValueError):
+    """Raised when the model's next-token distribution cannot be read as letters."""
+
+
+@dataclass(frozen=True)
+class LetterReadout:
+    """The full 26-letter logit readout for one prompt, plus a top-k view.
+
+    ``logits`` keeps every letter when a question has more options than ``top_k``.
+    ``top`` is the reported top-k softmax distribution.
+    ``letter_mass`` is retained as a nullable compatibility field. The optimized
+    26-row output head intentionally cannot report it because it never computes
+    a full-vocabulary distribution.
+    """
+
+    logits: dict[str, float]
+    top: "SoftmaxDistribution"
+    letter_mass: float | None = None
+
+
+@dataclass(frozen=True)
+class SoftmaxDistribution:
+    """A probability distribution over a chosen slice of letters."""
+
+    ranked: tuple[tuple[str, float], ...]
+    probabilities: dict[str, float]
+    top_k: int
+    temperature: float
+
+    @property
+    def top_letter(self) -> str:
+        return self.ranked[0][0]
+
+    @property
+    def top_probability(self) -> float:
+        return self.probabilities[self.top_letter]
+
+
+def _softmax(logits: list[float], temperature: float) -> list[float]:
+    if temperature <= 0:
+        raise ValueError("temperature must be > 0")
+    scaled = [x / temperature for x in logits]
+    peak = max(scaled)
+    exps = [math.exp(x - peak) for x in scaled]
+    total = sum(exps)
+    if total <= 0 or not math.isfinite(total):
+        raise LetterLogitError("softmax denominator is not finite")
+    return [e / total for e in exps]
+
+
+def distribution_from_letter_logits(
+    letter_logits: dict[str, float],
+    *,
+    top_k: int = 5,
+    temperature: float = 1.0,
+) -> SoftmaxDistribution:
+    """Softmax the top-k letter logits into a probability distribution.
+
+    `top_k` is clamped to 2..26. The softmax is taken *only* over the k letters we
+    keep, which is what TypeSafe-compatible consumers expect: the reported
+    probabilities sum to 1 across the surfaced options.
+    """
+    unknown = set(letter_logits) - set(LETTER_INDEX)
+    if unknown:
+        raise LetterLogitError(f"non-letter logit keys: {sorted(unknown)}")
+    if not letter_logits:
+        raise LetterLogitError("no letter logits supplied")
+
+    k = max(1, min(int(top_k), len(letter_logits), MAX_TOP_K))
+    ranked = sorted(letter_logits.items(), key=lambda kv: kv[1], reverse=True)[:k]
+    probs = _softmax([v for _, v in ranked], temperature)
+
+    distribution = {letter: p for (letter, _), p in zip(ranked, probs)}
+    # Keep the distribution in descending-probability order for stable output.
+    ordered = dict(sorted(distribution.items(), key=lambda kv: kv[1], reverse=True))
+    return SoftmaxDistribution(
+        ranked=tuple((letter, logit) for letter, logit in ranked),
+        probabilities=ordered,
+        top_k=k,
+        temperature=temperature,
+    )
+
+
+def gather_letter_logits(next_token_logits, letter_token_ids):
+    """Return the ``[batch, 26]`` A-Z slice of a vocabulary-logit tensor.
+
+    ``letter_token_ids`` may be the device-resident tensor owned by the engine.
+    Keeping it on-device avoids allocating a Python list and an indexing tensor on
+    every request.
+    """
+    if next_token_logits.dim() != 2:
+        raise LetterLogitError("expected next-token logits of shape [batch, vocab]")
+
+    import torch
+
+    ids = (
+        letter_token_ids
+        if isinstance(letter_token_ids, torch.Tensor)
+        else torch.tensor(list(letter_token_ids.values()), device=next_token_logits.device)
+    )
+    return next_token_logits.index_select(1, ids)
+
+
+def batch_letter_logits(
+    gathered_logits, letters: tuple[str, ...] | dict[str, int] = LETTERS
+) -> list[dict[str, float]]:
+    """Convert a previously gathered ``[batch, 26]`` A-Z tensor to dictionaries."""
+    # Preserve the original public helper signature for downstream callers while
+    # allowing the engine to pass the already-gathered slice.
+    if isinstance(letters, dict):
+        gathered_logits = gather_letter_logits(gathered_logits, letters)
+        letters = tuple(letters)
+    if gathered_logits.dim() != 2 or gathered_logits.shape[1] != len(letters):
+        raise LetterLogitError("expected gathered letter logits of shape [batch, 26]")
+    return [
+        {letter: float(value) for letter, value in zip(letters, row)}
+        for row in gathered_logits
+    ]

model.safetensors

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

modeling_truetype.py

@@ -0,0 +1,108 @@
+"""Gemma 4 with a 26-logit A-Z output layer instead of the 262k-token LM head.
+
+The decoder is unchanged. Only the final projection differs: ``lm_head`` has 26
+rows (the A-Z rows of the original tied embedding), so one forward pass returns
+one logit per letter and the vocabulary-sized matmul disappears. Gemma's final
+logit soft-capping still runs on those 26 values, exactly as it did on the
+original rows.
+
+    model = AutoModelForMultimodalLM.from_pretrained(repo, trust_remote_code=True)
+    tokenizer = AutoTokenizer.from_pretrained(repo)
+    result = model.system_one(tokenizer, state="...", questions={...})
+"""
+
+from __future__ import annotations
+
+import json
+
+import torch
+from torch import nn
+from transformers.conversion_mapping import (
+    get_checkpoint_conversion_mapping,
+    register_checkpoint_conversion_mapping,
+)
+from transformers.models.gemma4_unified.modeling_gemma4_unified import (
+    Gemma4UnifiedForConditionalGeneration,
+)
+
+from .configuration_truetype import LETTERS, TrueTypeSystemOneConfig
+from .questions import build_question, score_question_tensor
+from .render import render_batch
+
+DEFAULT_TEMPERATURE = 0.7
+
+# Gemma's checkpoints keep the vision embedder under legacy names that
+# transformers renames per model_type. This model shares that layout.
+register_checkpoint_conversion_mapping(
+    TrueTypeSystemOneConfig.model_type, get_checkpoint_conversion_mapping("gemma4_unified"), overwrite=True
+)
+
+
+class TrueTypeSystemOneModel(Gemma4UnifiedForConditionalGeneration):
+    config_class = TrueTypeSystemOneConfig
+    _tied_weights_keys = {}
+
+    def __init__(self, config: TrueTypeSystemOneConfig):
+        super().__init__(config)
+        self.lm_head = nn.Linear(config.text_config.hidden_size, len(LETTERS), bias=False)
+        self.post_init()
+
+    def letter_logits(self, input_ids, attention_mask=None, **kwargs) -> torch.Tensor:
+        """``[batch, 26]`` float32 logits for A-Z at the next-token position.
+
+        Inputs must be left-padded so the last position is the answer slot.
+        """
+        out = self(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            logits_to_keep=1,
+            use_cache=False,
+            **kwargs,
+        )
+        return out.logits[:, -1, :].float()
+
+    @torch.inference_mode()
+    def system_one(
+        self,
+        tokenizer,
+        state,
+        questions: dict[str, dict],
+        temperature: float = DEFAULT_TEMPERATURE,
+        batch_size: int = 16,
+        include_letter_logits: bool = False,
+    ) -> dict:
+        """Answer typed questions (``noul``, ``choice``, ``score``) about ``state``.
+
+        ``questions`` maps an id to ``{"type", "instructions", "criteria"}``.
+        Returns ``{"answers": {id: {...probabilities and answer...}}}``.
+        """
+        if not questions:
+            raise ValueError("at least one question is required")
+        state_text = state if isinstance(state, str) else json.dumps(state, indent=2)
+        parsed = [build_question(qid, spec) for qid, spec in questions.items()]
+        prompts = [item.text for item in render_batch(parsed, state_text)]
+
+        rows = []
+        padding_side, tokenizer.padding_side = tokenizer.padding_side, "left"
+        try:
+            for start in range(0, len(prompts), max(1, batch_size)):
+                encoded = tokenizer(
+                    prompts[start : start + batch_size],
+                    return_tensors="pt",
+                    padding=True,
+                    add_special_tokens=True,
+                ).to(self.device)
+                rows.extend(self.letter_logits(encoded["input_ids"], encoded["attention_mask"]))
+        finally:
+            tokenizer.padding_side = padding_side
+
+        result = {"answers": {}}
+        for question, row in zip(parsed, rows):
+            result["answers"][question.id] = score_question_tensor(
+                question, row, temperature=temperature
+            ).to_dict()
+        if include_letter_logits:
+            result["letter_logits"] = {
+                question.id: dict(zip(LETTERS, row.tolist())) for question, row in zip(parsed, rows)
+            }
+        return result

pipeline_truetype.py

@@ -0,0 +1,58 @@
+"""``pipeline("system-one", ...)``: typed questions in, probabilities and answers out.
+
+    pipe = pipeline("system-one", model=repo, trust_remote_code=True)
+    pipe({"state": "I was charged twice.",
+          "questions": {"refund": {"type": "noul", "instructions": "Does the text request a refund?"}}})
+"""
+
+from __future__ import annotations
+
+import json
+
+import torch
+from transformers import Pipeline
+
+from .configuration_truetype import LETTERS
+from .questions import build_question, score_question_tensor
+from .render import render_batch
+
+
+class SystemOnePipeline(Pipeline):
+    def _sanitize_parameters(self, temperature=None, include_letter_logits=None, **kwargs):
+        postprocess = {}
+        if temperature is not None:
+            postprocess["temperature"] = temperature
+        if include_letter_logits is not None:
+            postprocess["include_letter_logits"] = include_letter_logits
+        return {}, {}, postprocess
+
+    def preprocess(self, inputs):
+        state, questions = inputs["state"], inputs["questions"]
+        if not questions:
+            raise ValueError("at least one question is required")
+        state_text = state if isinstance(state, str) else json.dumps(state, indent=2)
+        parsed = [build_question(qid, spec) for qid, spec in questions.items()]
+        prompts = [item.text for item in render_batch(parsed, state_text)]
+        self.tokenizer.padding_side = "left"
+        encoded = self.tokenizer(prompts, return_tensors="pt", padding=True, add_special_tokens=True)
+        return {"questions": parsed, "input_ids": encoded["input_ids"], "attention_mask": encoded["attention_mask"]}
+
+    def _forward(self, model_inputs):
+        with torch.inference_mode():
+            logits = self.model.letter_logits(
+                model_inputs["input_ids"].to(self.model.device),
+                model_inputs["attention_mask"].to(self.model.device),
+            )
+        return {"questions": model_inputs["questions"], "logits": logits}
+
+    def postprocess(self, model_outputs, temperature=0.7, include_letter_logits=False):
+        questions, logits = model_outputs["questions"], model_outputs["logits"]
+        result = {
+            "answers": {
+                q.id: score_question_tensor(q, row, temperature=temperature).to_dict()
+                for q, row in zip(questions, logits)
+            }
+        }
+        if include_letter_logits:
+            result["letter_logits"] = {q.id: dict(zip(LETTERS, row.tolist())) for q, row in zip(questions, logits)}
+        return result

questions.py

@@ -0,0 +1,270 @@
+"""Map Gemma letter logits onto TypeSafe's three question primitives.
+
+The letter distribution is the model's raw output. Each primitive renormalizes
+the distribution over its legal answer letters, so the probabilities sum to 1
+across the caller's options.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+
+from .letters import LETTER_INDEX, LETTERS, _softmax
+
+# Fixed bank of question-agnostic demonstrations. Their predicates differ from the
+# target question on purpose: they teach the letter-slot format, not the task.
+GENERIC_EXAMPLES: tuple[tuple[str, str, str], ...] = (
+    ("I need this fixed immediately, it is blocking everything.", "Does the text express urgency?", "Y"),
+    ("Here is the documentation link for reference.", "Does the text express urgency?", "N"),
+    ("The app crashes when I open the settings page.", "Does the text contain a bug report?", "Y"),
+    ("Just wanted to share the meeting notes from today.", "Does the text contain a bug report?", "N"),
+    ("Please send this to a human agent right now.", "Does the text ask for a human agent?", "Y"),
+    ("No rush on this, whenever you get a chance.", "Does the text ask for a human agent?", "N"),
+    ("URGENT: the payment system is failing right now.", "Does the text express urgency?", "Y"),
+    ("The sky is blue and the grass is green.", "Does the text express urgency?", "N"),
+)
+
+DEFAULT_EXAMPLE_COUNT = 6
+
+
+@dataclass(frozen=True)
+class QuestionOption:
+    letter: str
+    label: str
+    description: str | None = None
+
+
+@dataclass
+class Question:
+    id: str
+    type: str  # noul | choice | score
+    instructions: str
+    options: list[QuestionOption] = field(default_factory=list)
+    criteria: object = None
+
+
+@dataclass(frozen=True)
+class QuestionAnswer:
+    type: str
+    choice: str | None = None
+    probabilities: dict[str, float] = field(default_factory=dict)
+    confidence: float | None = None
+    noul: float | None = None
+    score: float | None = None
+    legend: dict[str, str] | None = None
+
+    def to_dict(self) -> dict:
+        payload: dict = {"type": self.type}
+        if self.type == "noul":
+            payload["noul"] = round(self.noul or 0.0, 6)
+        elif self.type == "choice":
+            payload["choice"] = self.choice
+            payload["probabilities"] = {k: round(v, 6) for k, v in self.probabilities.items()}
+            payload["confidence"] = round(self.confidence or 0.0, 6)
+        elif self.type == "score":
+            payload["score"] = round(self.score or 0.0, 6)
+            payload["legend"] = self.legend or {}
+            payload["probabilities"] = {k: round(v, 6) for k, v in self.probabilities.items()}
+            payload["confidence"] = round(self.confidence or 0.0, 6)
+        return payload
+
+
+def confidence_from_probabilities(probabilities: list[float]) -> float:
+    """1.0 when all mass is on one option, 0.0 when uniform."""
+    n = len(probabilities)
+    if n <= 1:
+        return 1.0
+    entropy = -sum(p * math.log(p) for p in probabilities if p > 0)
+    return max(0.0, 1.0 - entropy / math.log(n))
+
+
+def _predicate_forms(instructions: str) -> tuple[str, str]:
+    """Turn an instruction into affirmative/negative option text.
+
+    "Does the text request a refund?" -> ("Yes, the text requests a refund.",
+                                          "No, the text does not request a refund.")
+    """
+    text = instructions.strip().rstrip(".").strip()
+    lowered = text.lower()
+
+    if lowered.startswith("does the text "):
+        body = text[len("Does the text ") :]
+        affirmative = f"the text {body}"
+    elif lowered.startswith("is the text "):
+        body = text[len("Is the text ") :]
+        affirmative = f"the text {body}"
+    elif lowered.startswith("is this "):
+        body = text[len("Is this ") :]
+        affirmative = f"this {body}"
+    elif lowered.startswith("does this "):
+        body = text[len("Does this ") :]
+        affirmative = f"this {body}"
+    elif lowered.startswith("does the "):
+        body = text[len("Does the ") :]
+        affirmative = f"the {body}"
+    elif lowered.startswith("is the "):
+        body = text[len("Is the ") :]
+        affirmative = f"the {body}"
+    elif lowered.startswith("did the "):
+        body = text[len("Did the ") :]
+        affirmative = f"the {body}"
+    elif lowered.startswith("can the "):
+        body = text[len("Can the ") :]
+        affirmative = f"the {body}"
+    else:
+        affirmative = text
+
+    affirmative = affirmative.rstrip("?").strip()
+    if affirmative.lower().startswith("the "):
+        negative = "the " + affirmative[4:].split(" ", 1)[0] + " does not" + (
+            " " + affirmative[4:].split(" ", 1)[1] if " " in affirmative[4:] else ""
+        )
+    else:
+        negative = f"it is not the case that {affirmative}"
+    return f"Yes, {affirmative}.", f"No, {negative}."
+
+
+def build_question(qid: str, spec: dict) -> Question:
+    """Parse one TypeSafe question spec into a letter-addressable Question."""
+    qtype = spec.get("type")
+    if qtype not in {"noul", "choice", "score"}:
+        raise ValueError(f"question {qid!r} has unsupported type {qtype!r}")
+
+    instructions = spec.get("instructions")
+    if not isinstance(instructions, str) or not instructions.strip():
+        raise ValueError(f"question {qid!r} requires non-empty string instructions")
+
+    criteria = spec.get("criteria")
+
+    if qtype == "noul":
+        if isinstance(criteria, dict):
+            yes_desc = criteria.get("true")
+            no_desc = criteria.get("false")
+        else:
+            yes_desc = no_desc = None
+        yes_text, no_text = _predicate_forms(instructions)
+        return Question(
+            id=qid,
+            type=qtype,
+            instructions=instructions,
+            options=[
+                QuestionOption("Y", "yes", yes_desc or yes_text),
+                QuestionOption("N", "no", no_desc or no_text),
+            ],
+            criteria=criteria,
+        )
+
+    if qtype == "choice":
+        if not isinstance(criteria, dict) or not criteria:
+            raise ValueError(f"choice question {qid!r} requires a non-empty criteria map")
+        if len(criteria) > len(LETTERS):
+            raise ValueError(f"choice question {qid!r} supports at most 26 options")
+        entries = list(criteria.items())
+    else:
+        if not isinstance(criteria, list) or len(criteria) < 2:
+            raise ValueError(f"score question {qid!r} requires an ordered list of at least two levels")
+        if len(criteria) > len(LETTERS):
+            raise ValueError(f"score question {qid!r} supports at most 26 levels")
+        entries = [(f"level {i}", level) for i, level in enumerate(criteria)]
+
+    options = [
+        QuestionOption(
+            LETTERS[i],
+            label,
+            description if isinstance(description, str) else str(description),
+        )
+        for i, (label, description) in enumerate(entries)
+    ]
+    return Question(id=qid, type=qtype, instructions=instructions, options=options, criteria=criteria)
+
+
+def score_question(question: Question, letter_logits: dict[str, float], temperature: float = 1.0) -> QuestionAnswer:
+    """Turn raw letter logits into a typed answer, renormalized over legal options."""
+    option_letters = [option.letter for option in question.options]
+    missing = [letter for letter in option_letters if letter not in letter_logits]
+    if missing:
+        raise ValueError(f"question {question.id!r} is missing logits for {missing}")
+
+    logits = [letter_logits[letter] for letter in option_letters]
+    probs = _softmax(logits, temperature)
+
+    if question.type == "noul":
+        noul = probs[0]
+        return QuestionAnswer(
+            type="noul",
+            noul=noul,
+        )
+
+    if question.type == "choice":
+        pairs = sorted(zip(question.options, probs), key=lambda op: op[1], reverse=True)
+        top = pairs[0]
+        return QuestionAnswer(
+            type="choice",
+            choice=top[0].label,
+            probabilities={option.label: p for option, p in pairs},
+            confidence=confidence_from_probabilities(probs),
+        )
+
+    score = sum(i * p for i, p in enumerate(probs))
+    return QuestionAnswer(
+        type="score",
+        score=score,
+        probabilities={str(i): p for i, p in enumerate(probs)},
+        legend={str(i): option.description or "" for i, option in enumerate(question.options)},
+        confidence=confidence_from_probabilities(probs),
+    )
+
+
+def score_question_tensor(question: Question, letter_logits, temperature: float = 1.0) -> QuestionAnswer:
+    """Score a ``[26]`` device tensor, transferring only final answer values.
+
+    The logits, legal-option gather, softmax, entropy, ranking, and score
+    reduction all remain on the tensor's device.  Constructing ``QuestionAnswer``
+    necessarily creates Python/JSON values at the API boundary.
+    """
+    if temperature <= 0:
+        raise ValueError("temperature must be > 0")
+    if letter_logits.dim() != 1 or letter_logits.shape[0] != len(LETTERS):
+        raise ValueError("expected one [26] letter-logit tensor")
+
+    import torch
+
+    indices = torch.tensor(
+        [LETTER_INDEX[option.letter] for option in question.options],
+        device=letter_logits.device,
+        dtype=torch.long,
+    )
+    probs = torch.softmax(letter_logits.index_select(0, indices).float() / temperature, dim=0)
+
+    if question.type == "noul":
+        return QuestionAnswer(type="noul", noul=float(probs[0].item()))
+
+    if len(question.options) <= 1:
+        confidence_value = 1.0
+    else:
+        entropy = -(probs * torch.where(probs > 0, probs.log(), torch.zeros_like(probs))).sum()
+        confidence = 1.0 - entropy / torch.log(torch.tensor(float(len(question.options)), device=probs.device))
+        confidence_value = float(torch.clamp(confidence, min=0.0).item())
+
+    if question.type == "choice":
+        # Stable GPU argsort preserves the request's option order for exact ties.
+        order = torch.argsort(probs, descending=True, stable=True).tolist()
+        probability_values = probs.tolist()
+        return QuestionAnswer(
+            type="choice",
+            choice=question.options[order[0]].label,
+            probabilities={question.options[index].label: probability_values[index] for index in order},
+            confidence=confidence_value,
+        )
+
+    probability_values = probs.tolist()
+    levels = torch.arange(len(question.options), device=probs.device, dtype=probs.dtype)
+    score = float((levels * probs).sum().item())
+    return QuestionAnswer(
+        type="score",
+        score=score,
+        probabilities={str(index): value for index, value in enumerate(probability_values)},
+        legend={str(index): option.description or "" for index, option in enumerate(question.options)},
+        confidence=confidence_value,
+    )

render.py

@@ -0,0 +1,169 @@
+"""Render every question in a request to a base-completion prompt.
+
+Prompt layout
+-------------
+Every block puts the answer-relevant scaffolding *first* and the text being
+judged *last*, so a prompt is::
+
+    Choices:
+    Y. Yes, ...
+    N. No, ...
+    Question: <question>
+    Text: <example text>
+    Answer: Y
+
+    ... more examples ...
+
+    Choices:
+    ...
+    Question: <target question>
+    Text: <the caller's state>
+    Answer:
+
+Putting the state last makes the engine's prefix KV cache effective:
+everything above ``Text:`` depends only on the question, so it is byte-identical
+across calls and can be cached once. That took a warm Doom decision from 357ms to
+178ms (mean prompt tail 68 -> 20 tokens) with no accuracy change: 63/63 on the
+50-case suite plus the Doom probe and 8 held-out noul cases.
+
+Two example strategies, chosen by question type:
+
+- Noul (Y/N) uses a fixed bank of demonstrations with their own questions.
+  The predicate in the options ("Yes, the text requests a refund.") carries the
+  task; the examples teach the letter slot. Their questions
+  differ from the target question, so their Y/N labels stay truthful rather than
+  being relabelled under a question they were not written for. This lifted
+  accuracy from 3/6 to 6/6 in tests/example_transfer.py.
+
+- Choice and Score use examples generated from the caller's criteria. Each
+  option gets one synthetic demonstration mapping it to its letter, visited in a
+  non-alphabetical order so the model cannot exploit A,B,C positional cues. This
+  reached 10/10 on Choice and 11/12 on Score in tests/choice_score_bakeoff2.py and
+  needs no hand-authoring.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from .questions import DEFAULT_EXAMPLE_COUNT, GENERIC_EXAMPLES, Question
+
+
+@dataclass
+class RenderedPrompt:
+    question_id: str
+    text: str
+
+
+def _render_options(question: Question) -> str:
+    lines = []
+    for option in question.options:
+        if question.type == "noul":
+            lines.append(f"{option.letter}. {option.description or option.label}")
+        elif option.description:
+            lines.append(f"{option.letter}. {option.label} - {option.description}")
+        else:
+            lines.append(f"{option.letter}. {option.label}")
+    return "\n".join(lines)
+
+
+def _header(question: Question, instructions: str | None = None) -> str:
+    """Choices + question. Identical for every call on the same question."""
+    return (
+        f"Choices:\n{_render_options(question)}\n"
+        f"Question: {instructions or question.instructions}\n"
+    )
+
+
+def _noul_examples(question: Question, count: int) -> list[str]:
+    blocks = []
+    for ex_text, ex_question, ex_letter in GENERIC_EXAMPLES[:count]:
+        blocks.append(f"{_header(question, ex_question)}Text: {ex_text}\nAnswer: {ex_letter}")
+    return blocks
+
+
+def _criteria_roundtrip_examples(question: Question) -> list[str]:
+    """One demonstration per option, built from the option's own description.
+
+    Options use a non-alphabetical order that
+    interleaves the scale (low, high, middle, ...) so no positional shortcut exists.
+    """
+    header = _header(question)
+    blocks = []
+    for index in _interleaved_order(len(question.options)):
+        option = question.options[index]
+        if question.type == "choice":
+            detail = f" This example is about the topic of {option.label}."
+        else:
+            detail = " " + (option.description or option.label)
+        blocks.append(
+            f"{header}Text: An example of {option.label}:{detail}\nAnswer: {option.letter}"
+        )
+    return blocks
+
+
+def _interleaved_order(n: int) -> list[int]:
+    """Low, high, then remaining indices ascending. Breaks A,B,C positional cues."""
+    if n <= 2:
+        return list(range(n))
+    order = [0, n - 1]
+    order.extend(i for i in range(1, n - 1))
+    return order
+
+
+def render_example_prefix(
+    question: Question,
+    *,
+    example_count: int = DEFAULT_EXAMPLE_COUNT,
+) -> str:
+    """The part of the prompt that depends only on the question, never on the state.
+
+    This is what the engine keeps in a KV cache across calls: the few-shot examples
+    plus the target question's own Choices/Question header. For a fixed question it
+    is byte-identical every time, and it is ~90% of the prompt.
+    """
+    if question.type == "noul":
+        blocks = _noul_examples(question, example_count)
+    else:
+        blocks = _criteria_roundtrip_examples(question)
+    body = "\n\n".join(blocks) + "\n\n" if blocks else ""
+    return body + _header(question)
+
+
+def render_target_block(question: Question, state: str) -> str:
+    """The state-dependent tail, ending at the ``Answer:`` slot.
+
+    There is no trailing space after ``Answer:``. The letter tokens
+    the engine reads are space-prefixed (``"▁Y"``), which is how the examples above
+    tokenize (``['Answer', ':', '▁Y']``). Adding a trailing space here would emit a
+    standalone ``'��'`` token and strand the space, so the model would want a bare
+    ``"Y"`` while the engine read ``"▁Y"``. This measured 0.0000 probability mass on
+    the letters being scored, versus 0.9871 without the space.
+    """
+    return f"Text: {state}\nAnswer:"
+
+
+def render_question(
+    question: Question,
+    state: str,
+    *,
+    example_count: int = DEFAULT_EXAMPLE_COUNT,
+) -> RenderedPrompt:
+    # Composed from the two halves so the cached-prefix path and the plain path
+    # always produce the exact same prompt string.
+    text = render_example_prefix(question, example_count=example_count) + render_target_block(
+        question, state
+    )
+    return RenderedPrompt(question_id=question.id, text=text)
+
+
+def render_batch(
+    questions: list[Question],
+    state: str,
+    *,
+    example_count: int = DEFAULT_EXAMPLE_COUNT,
+) -> list[RenderedPrompt]:
+    """One prompt per question. Every question sees the same state, as in TypeSafe."""
+    return [
+        render_question(question, state, example_count=example_count) for question in questions
+    ]

tokenizer.json

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

tokenizer_config.json

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