Release Drex v2.0 code and model card

README.md

@@ -9,78 +9,56 @@ tags:
 - kev
 - drex
 - moe
-private: true
 ---
 
 # Drex v2.0
 
-A decision model: it reads a state (text or JSON) and typed questions (`choice`, `noul`, ordinal `score`) and returns a probability for every option in one forward pass per question. Nothing is generated. It serves the TypeSafe-compatible `/v1/systemone` API through the bundled Kev code.
+**Nace's flagship decision model: 61.25 on the public Decision Index 0.3.1, 3.3 points ahead of Jev 1.13.0 and 3.2 ahead of Drex v1.5.**
 
-Mixture-of-experts backbone (Qwen3.6-35B-A3B, 3B active) plus a pointer head.
+We introduce Drex v2.0 from [Nace.AI](https://www.nace.ai/), a decision model that answers typed questions about a given context. Pass the context as `state` ([request format](https://github.com/nace-ai/drex-decision-models/blob/main/docs/request-format.md)) with one or more named questions, and get a probability for every option back in one forward pass. Nothing is generated. Drex v2.0 beats Jev head to head in the eight-game arena (60.9%) and on email classification (AUC 0.984 against 0.981), and it runs on your own hardware with Python, llama.cpp or Ollama. It is built on [Qwen3.6-35B-A3B](https://huggingface.co/Qwen/Qwen3.6-35B-A3B), a mixture of experts with 35B parameters and about 3B active per token, with a [Kev](https://github.com/jaredpalmer/kev) pointer head, and is released under the Nace.AI Open RAIL-M license.
 
-## Inference
+[Drex repository](https://github.com/nace-ai/drex-decision-models) · [Drex v1.5](https://huggingface.co/nace-ai/drex-v1.5) · [llama.cpp fork](https://github.com/nace-ai/llama.cpp) (branch `drex-v2.0`) · [Ollama fork](https://github.com/nace-ai/ollama) (branch `drex-v2.0`)
 
-```bash
-pip install -e code
-python inference.py --model NischayDnk/Drex-v2.0                     # built-in example
-python inference.py --model NischayDnk/Drex-v2.0 --request req.json  # your own System One request
-```
+It runs with Python, llama.cpp or Ollama. The default limit is 16,384 tokens per request and it can be raised to 131,072 ([Context length](https://github.com/nace-ai/drex-decision-models/blob/main/docs/context-length.md)). The bf16 weights take about 70 GB and the Q8_0 GGUF about 37 GB.
 
-HTTP server (`POST /v1/systemone`):
-
-```bash
-KEV_DTYPE=bf16 python -m kev.serve --run <folder with head.pt> --port 8000
-# long documents: KEV_SERVE_MAX_STATE=32768 KEV_SERVE_MAX_BRANCH=32768
-```
+<details>
+<summary>How it works</summary>
 
-Request shape:
+The backbone is a mixture-of-experts `Qwen3_5MoeForConditionalGeneration` with 40 layers, 256 experts and 8 active per token: 35B parameters in total, about 3B active per token. Only its language model is used; the vision tower is dropped on load. Each question is scored in one forward pass over the state plus that question, and the pointer head in `head.pt` turns the result into probabilities. Nothing is generated, so the sampling values in `generation_config.json` do not apply.
 
-```json
-{"state": {"ticket": "I was charged twice for order A-104."},
- "questions": {"team": {"type": "choice", "instructions": "Which team handles this?",
-                        "criteria": {"billing": "Charges and refunds", "shipping": "Delivery", "other": null}},
-               "refund": {"type": "noul", "instructions": "Does the customer ask for a refund?"}}}
-```
+</details>
 
-Fast scoring (vLLM pooling, pointer head in-process — do not use generate):
+## Performance
 
-```bash
-export VLLM_ENABLE_V1_MULTIPROCESSING=0
-# LLM(..., runner="pooling", convert="embed", dtype="bfloat16",
-#      pooler_config={"task": "token_embed", "use_activation": False})
-# PointerHead from head.pt: softmax(head(h[decide], h[opts]))
-```
+### Decision Index 0.3.1 (public only)
 
-## Benchmarks
+| Model | Public index | Raw | Knowledge & Reasoning | Language | Retrieval | Tools | Arts |
+|---|---|---|---|---|---|---|---|
+| **Drex v2.0** (35B) | **61.25** | 70.55 | 47.9 | 65.3 | 63.0 | 76.5 | 53.9 |
+| Drex v1.5 (9B) | 58.08 | 68.08 | 44.6 | 60.4 | 61.7 | 75.0 | 48.8 |
+| Jev 1.13.0 | 57.96 | 68.55 | 53.9 | 59.2 | 55.4 | 75.1 | 39.1 |
+| Bespoke Nimble 9B v3 | 57.19 | 68.38 | 41.0 | 60.7 | 55.5 | 84.2 | 44.0 |
+| Cloudflare clef-flash | 56.15 | 66.53 | 53.0 | 47.2 | 52.2 | 81.9 | 48.4 |
+| ezjev 4B s2 | 50.82 | 63.16 | 31.0 | 60.5 | 56.3 | 69.9 | 31.1 |
 
-### Decision Index 0.2.1 (official kit, full suite, 150,317 requests)
+Public index only: the 37 public benchmarks in five areas, without the private tests. Area columns are chance-corrected skill × 100. The Drex rows are from our own runs of the official kit on the full public suite; the other rows are from the [leaderboard](https://huggingface.co/spaces/multimodalart/jev-decision-index).
 
-| Rank | Model | Index | Raw | Knowledge & Reasoning | Language | Retrieval | Tools | Arts |
-|---|---|---|---|---|---|---|---|---|
-| **1** | **Drex v2.0** | **60.78** | **69.70** | 47.7 | **65.9** | 62.9 | **76.0** | **49.1** |
-| 2 | Drex v1.5 | 58.28 | 67.80 | 44.2 | 62.5 | 62.0 | 75.3 | 45.1 |
-| 3 | Jev 1.13.0 | 57.91 | 68.09 | **51.4** | 62.0 | 55.4 | 75.1 | 37.7 |
-| 4 | Surogate Rune 26B-A4B v3 | 57.44 | 67.30 | 43.4 | 63.1 | **63.5** | 71.2 | 41.9 |
-| 5 | Decider chat · Gemma-4-31B | 57.33 | 67.22 | 44.3 | 60.4 | 63.1 | 75.6 | 38.3 |
-| 6 | AutoJev-27B | 56.40 | 66.89 | 40.9 | 63.5 | 54.9 | 79.3 | 39.4 |
-
-Area columns are chance-corrected skill × 100. Drex v2.0 is ahead of Jev on 21 of the 38 index benchmarks and on 4 of 5 areas (Knowledge trails). Unanswered requests count as wrong.
-
-![Decision Index 0.2.1](decision_index.svg)
-
-### JevBench (231 public items)
+<details>
+<summary>JevBench (231 public items)</summary>
 
 | Model | Accuracy | Easy | Standard | Hard |
 |---|---|---|---|---|
-| Jev 1.13.0 | **86.6%** (200/231) | 100.0% | **98.6%** | 73.0% |
+| Jev 1.13.0 | **86.6%** (200/231) | 100.0% | **98.6%** | **73.0%** |
 | **Drex v2.0** | 85.7% (198/231) | 100.0% | 95.8% | **73.0%** |
-| Drex v1.5 | 85.7% (198/231) | 100.0% | 95.8% | 73.0% |
 
-Choice = top option, yes/no = noul ≥ 0.5, ordinal = top level. Public items only.
+Choice = top option, yes/no = noul ≥ 0.5, ordinal = top level.
+
+</details>
 
-### DecideBench v1.1 (400 contrastive decisions)
+<details>
+<summary>DecideBench v1.1 (400 contrastive decisions)</summary>
 
-Published decision models only (chat LLMs omitted). Few-shot uses the benchmark’s examples-in-instructions protocol.
+Published decision models only (chat LLMs omitted). Few-shot uses the benchmark's examples-in-instructions protocol. Not an official leaderboard submission.
 
 | Entry | Accuracy | Pair accuracy |
 |---|---|---|
@@ -95,9 +73,10 @@ Published decision models only (chat LLMs omitted). Few-shot uses the benchmark
 | Drex v1.5 | 90.0% | 80.5% |
 | JevK5 v0.3 | 88.8% | 78.5% |
 
-Not an official leaderboard submission.
+</details>
 
-### Email classification
+<details>
+<summary>Email classification</summary>
 
 Held-out addresses. Score is AUC; caught @ 5% is recall at a 5% false-positive rate.
 
@@ -106,14 +85,17 @@ Held-out addresses. Score is AUC; caught @ 5% is recall at a 5% false-positive r
 | AUC | **0.984** | 0.981 | 0.972 |
 | Caught @ 5% FP | **89.4%** | 86.4% | 80.3% |
 
-### Game arena against Jev 1.13.0
+</details>
 
-Eight games, 32 matches each (16 openings × both colours). Score = (wins + ½ draws) ÷ games.
+<details>
+<summary>Game arena against Jev</summary>
+
+Eight OpenSpiel games, 32 matches each (16 openings × both colours). Score = (wins + ½ draws) ÷ matches.
 
 |  | W – D – L | Score |
 |---|---|---|
 | **Drex v2.0 vs Jev** | **132 – 48 – 76** | **60.9%** |
-| Drex v1.5 vs Jev (published) | 122 – 47 – 87 | 56.8% |
+| Drex v1.5 vs Jev | 122 – 47 – 87 | 56.8% |
 
 | Game | Drex v2.0 vs Jev | Drex v1.5 vs Jev |
 |---|---|---|
@@ -126,12 +108,106 @@ Eight games, 32 matches each (16 openings × both colours). Score = (wins + ½ d
 | Dots and Boxes | 15 – 0 – 17 (46.9%) | 18 – 0 – 14 (56.3%) |
 | Nine Men's Morris | 15 – 13 – 4 (67.2%) | 5 – 19 – 8 (45.3%) |
 
-## Files
+</details>
+
+## Python
+
+Needs a CUDA GPU; memory use grows with document length. The server answers `POST /v1/systemone`.
+
+```bash
+hf download nace-ai/drex-v2.0 --local-dir drex-v2.0
+pip install -r drex-v2.0/requirements.txt
+python drex-v2.0/inference.py --request drex-v2.0/examples/request.json
+python drex-v2.0/serve.py --port 8000
+```
+
+To accept requests longer than 16,384 tokens, set `KEV_CONTEXT` (up to 131,072); see [Context length](https://github.com/nace-ai/drex-decision-models/blob/main/docs/context-length.md).
+
+## llama.cpp
+
+Runs on Apple silicon (Metal) and on the CPU. Run these from the repository root:
+
+```bash
+git clone --branch drex-v2.0 --single-branch https://github.com/nace-ai/llama.cpp.git
+cmake -S llama.cpp -B llama.cpp/build
+cmake --build llama.cpp/build -j --target llama-server
+
+python llama.cpp/convert_hf_to_gguf.py drex-v2.0 --no-mtp --outtype q8_0 \
+  --outfile drex-v2.0/drex-v2.0-Q8_0.gguf
+
+llama.cpp/build/bin/llama-server -m drex-v2.0/drex-v2.0-Q8_0.gguf \
+  --embedding --pooling none -np 2 -c 32768 -b 16384 -ub 2048 -ngl 28 --port 8000
+```
+
+The server answers `POST /v1/systemone` with the same request and response as the Python server. The Q8_0 GGUF is 36.9 GB, so on a 48 GB Mac `-ngl 28` offloads 28 layers to Metal and runs the rest on the CPU; `-ngl 0` runs entirely on the CPU. The conversion needs the bf16 weights, about 70 GB, on disk.
+
+`-np 2` encodes the state once and shares it across questions. `-c` is split across the two slots, so this command accepts requests up to 16,384 tokens. For up to 131,072 tokens, set `SYSTEMONE_CONTEXT=131072` and use `-c 262144` ([Context length](https://github.com/nace-ai/drex-decision-models/blob/main/docs/context-length.md)).
+
+<details>
+<summary>Tested hardware</summary>
+
+Apple M5 Pro (18-core CPU, 20-core GPU, 48 GB unified memory, macOS 26.5.2): the Q8_0 GGUF gives the same answers on Metal with `-ngl 28` and on the CPU, with probabilities within 0.03 of each other.
+
+</details>
+
+## Ollama
+
+Use the Nace fork of Ollama with a `llama-server` built from the `drex-v2.0` branch above. Convert the GGUF first, as in the llama.cpp steps.
+
+<details>
+<summary>Build the fork</summary>
+
+From the repository root:
+
+```bash
+git clone --branch drex-v2.0 --single-branch https://github.com/nace-ai/ollama.git
+cd ollama
+export OLLAMA_LLAMA_CPP_SOURCE="$PWD/../llama.cpp"
+cmake -S llama/server --preset darwin
+cmake --build build/llama-server-darwin --target llama-server --parallel 8
+GOTOOLCHAIN=auto go build -trimpath -o ollama .
+cd ..
+```
+
+For the CPU use the `cpu` preset and `build/llama-server-cpu` in place of `darwin` and `build/llama-server-darwin` here and below.
+
+</details>
+
+Write a `Modelfile` next to the GGUF:
+
+```bash
+cat > drex-v2.0/Modelfile <<'EOF'
+FROM ./drex-v2.0-Q8_0.gguf
+CAPABILITY decision
+PARAMETER num_ctx 16384
+PARAMETER num_gpu 28
+EOF
+```
+
+Start the daemon with the fork's binary and leave it running:
+
+```bash
+OLLAMA_HOST=127.0.0.1:11434 \
+OLLAMA_LLAMA_SERVER="$PWD/ollama/build/llama-server-darwin/bin/llama-server" ollama/ollama serve
+```
+
+In another terminal, import and call the model:
+
+```bash
+OLLAMA_HOST=127.0.0.1:11434 ollama/ollama create drex-v2.0 -f drex-v2.0/Modelfile
+jq '. + {model: "drex-v2.0"}' req.json | curl http://127.0.0.1:11434/v1/systemone \
+  -H 'Content-Type: application/json' -d @-
+```
+
+Ollama needs `"model"` in the request, which the `jq` filter adds. `num_gpu 28` offloads 28 layers to Metal on a 48 GB Mac; use `PARAMETER num_gpu 0` for the CPU. The default is 16,384 tokens; for up to 131,072, set `PARAMETER num_ctx 131072`, import again, and start the daemon with `SYSTEMONE_CONTEXT=131072`.
+
+<details>
+<summary>Tested hardware</summary>
+
+Through Ollama on the Apple M5 Pro above, the Q8_0 GGUF with `num_gpu 28` gives the same answers as `llama-server`.
+
+</details>
+
+## License
 
-| | |
-|---|---|
-| `model-*.safetensors`, `config.json` | backbone (bf16) |
-| `head.pt` | pointer head |
-| `code/` | Kev package for loading and serving |
-| `inference.py` | minimal scorer for System One requests |
-| `decision_index.svg` | Decision Index 0.2.1 card |
+Nace.AI Open RAIL-M license, see [LICENSE](LICENSE).

code/kev/__init__.py

@@ -1 +0,0 @@
-"""Kev: load and serve a decision model on the TypeSafe System One API."""

code/kev/api.py

@@ -1,3 +1,8 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2026 Jared Palmer
+# Derived from Kev (https://github.com/jaredpalmer/kev).
+# Modified by Nace.AI for Drex: model default and ordered score legends.
+# See NOTICE and docs/upstream.md for attribution and revision information.
 """TypeSafe-compatible request/response shapes (POST /v1/systemone) mapped onto the single pointer primitive.
 
 Noul   -> 2 options [false, true];            answer = p(true)
@@ -42,7 +47,7 @@ Question = Union[Noul, Choice, Score]
 
 class SystemOneRequest(BaseModel):
     state: JSONContent
-    model: str = "kev-latest"
+    model: str = "drex"
     questions: dict[str, Question] = Field(min_length=1)
 
 
@@ -112,7 +117,7 @@ def to_record(req: SystemOneRequest):
             opts = [option_text(k, v) for k, v in q.criteria.items()]
         else:
             opts = [render(x) for x in q.criteria]
-            m["legend"] = dict(zip(m["keys"], opts))
+            m["legend"] = opts
         qs.append({"instr": render(q.instructions), "options": opts, "label": 0}); meta.append(m)
     return {"state": render(req.state), "questions": qs}, meta
 

code/kev/checkpoint.py

@@ -1,209 +1,60 @@
-"""Load a decision-model checkpoint (a folder or Hub repo with `head.pt` and the backbone).
-
-    ck = Checkpoint("NischayDnk/Drex-v2.0")
-    tok, model = ck.load("cuda", LoadOptions.from_env())
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2026 Jared Palmer
+# Adapted from Kev (https://github.com/jaredpalmer/kev).
+# Modified by Nace.AI for local merged Drex checkpoints and pointer heads.
+# See NOTICE and docs/upstream.md for attribution and revision information.
+"""Load a merged Drex checkpoint.
+
+The backbone shards live next to ``head.pt``. There is no adapter to merge.
 """
-import datetime
-import json
-import os
-import re
-from dataclasses import dataclass, field
 from pathlib import Path
+import pickle
 
 import torch
 
-from .model import DecisionModel, is_hybrid, load_tokenizer, pad_id
-
-HUB_ID = re.compile(r"[\w.-]+/[\w.-]+(@[\w.-]+)?")
-
-
-def is_hub_id(run):
-    return not os.path.isdir(run) and HUB_ID.fullmatch(str(run)) is not None
-
-
-def resolve_run(run):
-    """Local run directory as given, or a Hub repo id like jaredpalmer/kev-4b, optionally pinned to a revision or tag
-    with `@` (jaredpalmer/kev-4b@qwen3), downloaded to the HF cache. Returns a str path."""
-    if os.path.isdir(run):
-        return str(run)
-    from huggingface_hub import snapshot_download
-    repo, _, revision = str(run).partition("@")
-    return snapshot_download(repo, revision=revision or None, allow_patterns=["*.json", "*.safetensors", "*.pt", "*.txt", "*.jinja"])
-
-
-@dataclass
-class Meta:
-    """Contents of `head.pt`. Every reader gets the same defaults for fields older checkpoints did not write.
-    `extra` keeps any other keys written into `head.pt` so a read-modify-write round trip loses nothing."""
-    base: str
-    head: dict | None = None
-    base_revision: str | None = None
-    lora: int = 0
-    head_dim: int = 256
-    option_isolation: bool = False
-    special_embeddings: bool = False
-    weights_dtype: str = "fp32"
-    temperature: float = 1.0
-    holdout: list = field(default_factory=list)
-    state_bidir: bool = False
-    extra: dict = field(default_factory=dict)
-
-    KNOWN = ("base", "head", "base_revision", "lora", "head_dim", "option_isolation", "special_embeddings", "weights_dtype", "temperature", "holdout", "state_bidir")
-
-    @classmethod
-    def from_dict(cls, d):
-        return cls(**{k: d[k] for k in cls.KNOWN if k in d}, extra={k: v for k, v in d.items() if k not in cls.KNOWN})
-
-    def to_dict(self):
-        return {**self.extra, **{k: getattr(self, k) for k in self.KNOWN}}   # known fields win over a stray key in extra
-
-
-def read_meta(run):
-    return Meta.from_dict(torch.load(f"{run}/head.pt", map_location="cpu"))
-
-
-def write_meta(run, meta):
-    torch.save(meta.to_dict(), f"{run}/head.pt")
-
+from .model import DecisionModel, load_tokenizer
 
-@dataclass(frozen=True)
-class LoadOptions:
-    """How a checkpoint is turned into a model. Defaults are the exact path every reported number uses; the fields
-    are the same knobs the KEV_* environment variables expose to the command-line tools (see from_env).
 
-    dtype        None = fp32, the exact path every reported number uses (bf16 when the checkpoint was trained with a bf16
-                 backbone). kev.serve defaults to bf16 on CUDA and MPS instead: half the memory, 2-4.5x lower latency on an
-                 L4 (Kev-4B: 209 -> 118 ms at 101 tokens, 850 -> 189 ms at 330 tokens), probabilities within ~0.01 and
-                 the same argmax on the checks run so far. KEV_DTYPE=fp32 restores the exact path when serving.
-    merge        fold the LoRA into the base weights in fp32 before any cast. Exact in fp32; in bf16 it is faster (~15%)
-                 and closer to the fp32 numbers than the unmerged adapter (kev-4b, 24 dev records: max |dp| 0.017 vs
-                 0.029, 0 vs 1 argmax flips). Ignored for adapters that carry trained token embeddings.
-    attn         attention backend; None = the model default (SDPA on CUDA, eager elsewhere). "sdpa" on MPS measured
-                 parity with eager and is a few percent faster.
-    lora_scale   WiSE-FT-style interpolation between base (0) and fine-tuned weights (1), at inference.
-    temperature  None = the temperature the checkpoint carries (fitted by scripts/calibrate_checkpoint.py); 1.0 = raw logits.
-    backend      None = torch, the path every reported number uses. "mlx" = kev.mlx_model (Metal kernels for the hybrid
-                 Qwen3.5 backbones through mlx-lm; the pointer head and encoder are shared; refused for attention-only
-                 bases, which MPS already runs well). "auto" = mlx when the device is mps, the base is hybrid, mlx-lm is
-                 installed and fp32 was not asked for (an explicit dtype=float32 means "the exact path"), else torch;
-                 kev.serve uses auto. The MLX path always merges the adapter and ignores `attn` and `dtype` (the backbone
-                 runs as stored, bf16).
-    """
-    dtype: torch.dtype | None = None
-    merge: bool = True
-    attn: str | None = None
-    lora_scale: float = 1.0
-    temperature: float | None = None
-    backend: str | None = None
-
-    BACKENDS = (None, "torch", "mlx", "auto")
-
-    @classmethod
-    def from_env(cls, env=os.environ):
-        """KEV_DTYPE=bf16|fp16|fp32, KEV_MERGE=0, KEV_ATTN=sdpa|eager, KEV_LORA_SCALE, KEV_TEMPERATURE, KEV_BACKEND=torch|mlx|auto.
-        For command-line entry points only; library code passes an explicit LoadOptions. Explicit values that equal a
-        library default are kept (fp32 as torch.float32, "torch" as a string) so a caller with its own default, like
-        kev.serve, can tell "asked for it" from "did not say"."""
-        backend = env.get("KEV_BACKEND") or None
-        if backend not in cls.BACKENDS: raise ValueError(f"KEV_BACKEND must be one of torch, mlx, auto; got {backend!r}")
-        return cls(dtype={"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}.get(env.get("KEV_DTYPE", "")),
-                   merge=env.get("KEV_MERGE", "1") != "0", attn=env.get("KEV_ATTN") or None,
-                   lora_scale=float(env.get("KEV_LORA_SCALE", "1")),
-                   temperature=float(env["KEV_TEMPERATURE"]) if env.get("KEV_TEMPERATURE") else None, backend=backend)
-
-
-def mlx_available():
+def read_meta(path):
+    path = Path(path)
     try:
-        import mlx_lm  # noqa: F401
-        return True
-    except ImportError:
-        return False
-
-
-class Checkpoint:
-    def __init__(self, run):
-        self.requested = str(run)                    # what the caller asked for (a Hub id stays a Hub id in labels)
-        self.path = resolve_run(run)
-        self.meta = read_meta(self.path)
-
-    def file(self, name):
-        return Path(self.path) / name
+        meta = torch.load(path / "head.pt", map_location="cpu", weights_only=True)
+    except pickle.UnpicklingError as exc:
+        raise ValueError("head.pt must contain only tensors and primitive metadata; unsafe pickle objects are rejected") from exc
+    if not isinstance(meta, dict) or not isinstance(meta.get("head"), dict):
+        raise ValueError("head.pt must be a metadata dictionary containing a head state dictionary")
+    return meta
 
-    def adapter_config(self):
-        return json.loads(self.file("adapter_config.json").read_text(encoding="utf-8"))
 
-    def release_date(self):
-        """ISO date for the TypeSafe model card: the Hub commit date for a Hub checkpoint (falls back to the cached file's
-        date offline), the time head.pt was written for a local run."""
-        if is_hub_id(self.requested):
-            from huggingface_hub import HfApi
-            repo, _, revision = self.requested.partition("@")
-            try:
-                return HfApi().model_info(repo, revision=revision or None).last_modified.date().isoformat()
-            except Exception:
-                pass
-        return datetime.date.fromtimestamp(self.file("head.pt").stat().st_mtime).isoformat()
+def load(path, device=None, dtype=None):
+    """Return ``(tokenizer, model)`` in eval mode.
 
-    def hybrid_base(self):
-        """Whether the base has Gated DeltaNet layers (Qwen3.5), read from its config without loading weights."""
-        from transformers import AutoConfig
-        return is_hybrid(AutoConfig.from_pretrained(self.meta.base, revision=self.meta.base_revision).get_text_config())
-
-    def backend(self, device, opts=LoadOptions()):
-        """The backend `load` will use: LoadOptions.backend resolved ("auto" -> mlx only where it pays and is installed)."""
-        if opts.backend not in LoadOptions.BACKENDS: raise ValueError(f"unknown backend {opts.backend!r}")
-        if opts.backend != "auto": return opts.backend or "torch"
-        exact = opts.dtype is torch.float32   # KEV_DTYPE=fp32: the caller wants the reported-numbers path, not a faster one
-        return "mlx" if str(device) == "mps" and not exact and mlx_available() and self.hybrid_base() else "torch"
-
-    def load(self, device, opts=LoadOptions()):
-        """-> (tokenizer, model) in eval mode with the LoRA applied and the pointer head loaded. The model is a
-        DecisionModel (torch) or an MLXDecisionModel (backend mlx); both expose the same scoring interface."""
-        meta = self.meta
-        tok_src = self.path if self._standalone_weights() else meta.base
-        tok = load_tokenizer(tok_src, revision=None if self._standalone_weights() else meta.base_revision)
-        m = self._load_mlx(tok, opts) if self.backend(device, opts) == "mlx" else self._load_torch(tok, device, opts)
-        m.head.load_state_dict(meta.head); m.eval()
-        m.head.temperature = meta.temperature if opts.temperature is None else opts.temperature
-        return tok, m
-
-    def _load_mlx(self, tok, opts):
-        from .mlx_model import MLXDecisionModel, merge_lora
-        if not opts.merge: raise ValueError("the MLX backend always merges the adapter (KEV_MERGE=0 needs backend=torch)")
-        if self.meta.option_isolation: raise ValueError("option_isolation needs the packed mask; not available on the MLX backend")
-        if not self.hybrid_base(): raise ValueError(f"the MLX backend is for the hybrid (Qwen3.5) bases; {self.meta.base} is attention-only and runs on MPS with backend=torch")
-        base_dir = resolve_run(f"{self.meta.base}@{self.meta.base_revision or ''}")   # the base snapshot the torch path already cached
-        m = MLXDecisionModel(base_dir, pad_id(tok), head_dim=self.meta.head_dim)
-        merge_lora(m.lm, self.path, opts.lora_scale)
-        return m
-
-    def _standalone_weights(self):
-        return self.file("model.safetensors.index.json").exists() and not self.file("adapter_config.json").exists()
-
-    def _load_torch(self, tok, device, opts):
-        from peft import PeftModel
-        meta = self.meta
-        if self.file("MERGED").exists() or self._standalone_weights():
-            base = str(self.path) if self._standalone_weights() else meta.base
-            return DecisionModel(base, tok, device, lora=None, revision=None, head_dim=meta.head_dim,
-                                 option_isolation=meta.option_isolation, dtype=torch.bfloat16, attn=opts.attn)
-        dtype, merge = opts.dtype or torch.float32, opts.merge
-        if meta.weights_dtype == "bf16":
-            dtype, merge = torch.bfloat16, False
-        merge = merge and not self.adapter_config().get("trainable_token_indices")
-        m = DecisionModel(meta.base, tok, device, lora=None, revision=meta.base_revision, head_dim=meta.head_dim,
-                          option_isolation=meta.option_isolation, dtype=torch.float32 if merge else dtype, attn=opts.attn)
-        m.lm = PeftModel.from_pretrained(m.lm, self.path, torch_device=str(device)).to(device)
-        if opts.lora_scale != 1:
-            for module in m.lm.modules():
-                if isinstance(getattr(module, "scaling", None), dict):
-                    for k in module.scaling: module.scaling[k] *= opts.lora_scale
-            m.lora_scale = opts.lora_scale
-        if merge: m.lm = m.lm.merge_and_unload()
-        if dtype != torch.float32: m.lm = m.lm.to(dtype)
-        return m
-
-
-def load(run, device, opts=LoadOptions()):
-    """Convenience: Checkpoint(run).load(device, opts)."""
-    return Checkpoint(run).load(device, opts)
+    ``device`` defaults to CUDA, then Apple MPS, then CPU. ``dtype`` defaults
+    to bfloat16 on CUDA and MPS, and float32 on CPU.
+    """
+    path = Path(path)
+    meta = read_meta(path)
+    if device is None:
+        if torch.cuda.is_available():
+            device = "cuda"
+        elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
+            device = "mps"
+        else:
+            device = "cpu"
+    if dtype is None:
+        dtype = torch.bfloat16 if device in ("cuda", "mps") else torch.float32
+    tok = load_tokenizer(str(path))
+    model = DecisionModel(
+        str(path),
+        tok,
+        device,
+        lora=None,
+        head_dim=int(meta.get("head_dim", 256)),
+        option_isolation=bool(meta.get("option_isolation", False)),
+        dtype=dtype,
+        state_bidir=bool(meta.get("state_bidir", True)),
+    )
+    model.head.load_state_dict(meta["head"])
+    model.head.temperature = float(meta.get("temperature", 1.0))
+    model.eval()
+    return tok, model

code/kev/device.py

@@ -1,24 +0,0 @@
-"""The accelerator this process uses: cuda, then mps, then cpu."""
-import torch
-
-
-def default_device():
-    return "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
-
-
-def sync(device):
-    """Wait for queued kernels, so wall-clock timings around a forward pass are real."""
-    if device == "mps": torch.mps.synchronize()
-    elif device == "cuda": torch.cuda.synchronize()
-
-
-def empty_cache(device):
-    if device == "mps": torch.mps.empty_cache()
-    elif device == "cuda": torch.cuda.empty_cache()
-
-
-def allocated_bytes(device):
-    """Bytes currently allocated on the device (MPS) or the peak since the process started (CUDA); 0 on CPU."""
-    if device == "mps": return torch.mps.current_allocated_memory()
-    if device == "cuda": return torch.cuda.max_memory_allocated()
-    return 0

code/kev/mlx_model.py

@@ -1,144 +0,0 @@
-"""Apple Silicon backend for the Qwen3.5 checkpoints: mlx-lm's Metal implementation of the hybrid backbone under Kev's
-own encoder and pointer head.
-
-MPS has no Gated DeltaNet kernels, so the PyTorch path runs reference code there (Kev-4B ~0.8 s per request). This module
-runs the same computation on Metal through mlx-lm and keeps everything Kev-specific unchanged: `kev.model.encode` builds the
-tokens, `rows_of` splits them into one causal row per question (the hybrid form the torch path uses too), and the readout is
-the very same `PointerHead` (fp32, with the checkpoint's temperature) applied to the branch hidden states.
-
-Same contract as `DecisionModel` for serving and scoring: encode / forward / probs / probs_and_prefix / probs_with_prefix /
-head / dtype. Selected by `LoadOptions(backend="mlx")` (or "auto" on Apple Silicon) in `kev.checkpoint`; never by the
-benchmark, whose reported numbers stay on the fp32 torch path. Parity against that path is measured in
-tests/test_mlx.py (max |dp| and argmax flips on development records, prefix vs full pass, one question vs several).
-"""
-import json
-from pathlib import Path
-
-import mlx.core as mx
-import numpy as np
-import torch
-import torch.nn.functional as F
-from mlx.utils import tree_flatten
-from mlx_lm.models.cache import make_prompt_cache
-from mlx_lm.utils import load_model
-
-from .model import PointerHead, encode, rows_of, rows_per_pass
-
-
-def merge_lora(lm, adapter_dir, scale=1.0):
-    """Fold a PEFT adapter into the mlx-lm model's weights the way the torch path does: W + (B @ A) * alpha / r in fp32,
-    rounded once to the backbone dtype. `scale` is LoadOptions.lora_scale (WiSE-FT interpolation). Returns the tensor count."""
-    adapter_dir = Path(adapter_dir)
-    cfg = json.loads((adapter_dir / "adapter_config.json").read_text(encoding="utf-8"))
-    if cfg.get("trainable_token_indices"):
-        raise ValueError("the MLX backend does not carry trained token embeddings (special_embeddings checkpoints); use backend=torch")
-    alpha = cfg["lora_alpha"] / (cfg["r"] ** 0.5 if cfg.get("use_rslora") else cfg["r"])
-    weights = mx.load(str(adapter_dir / "adapter_model.safetensors"))
-    params = dict(tree_flatten(lm.parameters()))
-    merged = {}
-    with mx.stream(mx.cpu):   # the GPU's fp32 matmul is a reduced-precision fast path (~1e-3 relative on an M5); the merge is one-time and must be exact
-        for name, a in weights.items():
-            if not name.endswith(".lora_A.weight"):
-                continue
-            stem = name[: -len(".lora_A.weight")]
-            # peft names the wrapped text model `base_model.model.<layers...>`; mlx-lm nests it as `language_model.model.<layers...>`
-            target = stem.replace("base_model.model.", "language_model.model.", 1) + ".weight"
-            if target not in params:
-                raise ValueError(f"adapter tensor {stem} has no weight in the mlx-lm model (looked for {target})")
-            base = params[target]
-            delta = (weights[stem + ".lora_B.weight"].astype(mx.float32) @ a.astype(mx.float32)) * (alpha * scale)
-            merged[target] = (base.astype(mx.float32) + delta).astype(base.dtype)
-        mx.eval(list(merged.values()))
-    lm.load_weights(list(merged.items()), strict=False)
-    mx.eval(lm.parameters())
-    return len(merged)
-
-
-class MLXDecisionModel:
-    """Prefill-only scorer: hidden states from mlx-lm, logits from the shared torch PointerHead."""
-    backend, device, hybrid, option_isolation = "mlx", "mlx", True, False
-    prefix_min_tokens = 0   # kev.serve caches the state prefix for every request: on Metal the branch-only pass is always the cheaper one
-
-    def __init__(self, base_dir, pad_id, head_dim=256):
-        self.lm, _ = load_model(Path(base_dir))                       # weights as stored (bf16 for the Qwen3.5 bases)
-        self.text = self.lm.language_model.model                      # Qwen3_5TextModel: embeddings -> layers -> final norm = `.model.last_hidden_state`
-        self.pad_id = pad_id
-        self.head = PointerHead(self.text.embed_tokens.weight.shape[1], dp=head_dim).eval()
-
-    @property
-    def dtype(self):
-        return str(self.text.embed_tokens.weight.dtype).removeprefix("mlx.core.")
-
-    def eval(self):
-        self.head.eval(); return self
-
-    def encode(self, tok, rec, **kw):
-        return encode(tok, rec, option_isolation=False, **kw)
-
-    def _hidden(self, rows, cache=None):
-        """[N, L, d] hidden states of right-padded token rows. Pads sit after every real token and both layer kinds are
-        causal (attention: causal mask; DeltaNet: a left-to-right recurrence), so no real token sees a pad."""
-        L = max(len(r) for r in rows)
-        ids = mx.array([r + [self.pad_id] * (L - len(r)) for r in rows], dtype=mx.int32)
-        h = self.text(ids, cache=cache)
-        mx.eval(h)
-        return h
-
-    def _logits(self, h, decide, opts):
-        """One question's logits through the fp32 pointer head (temperature included, eval mode)."""
-        idx = mx.array([decide, *opts], dtype=mx.int32)
-        picked = torch.from_numpy(np.asarray(h[idx].astype(mx.float32)))
-        with torch.no_grad():
-            return self.head(picked[0], picked[1:])
-
-    def forward_rows(self, enc):
-        """Row form, as the torch path computes it: every question is one causal row of state + branch tokens, the state
-        recomputed per row. The reference the prefix form is checked against (tests/test_mlx.py); serving uses `forward`."""
-        S, _, rows = rows_of(enc)
-        chunk, out = rows_per_pass([S + r["ids"] for r in rows]), []
-        for start in range(0, len(rows), chunk):
-            part = rows[start:start + chunk]
-            h = self._hidden([S + r["ids"] for r in part])
-            out += [self._logits(h[i], len(S) + r["decide"], [len(S) + o for o in r["opts"]]) for i, r in enumerate(part)]
-        return out
-
-    # --- state prefix: the state runs once into an mlx-lm prompt cache (KV for the attention layers, conv + recurrent state
-    # for the DeltaNet layers); the branches run as one batch on a replicated copy, so the prefix stays pristine and can be
-    # reused by the next request with the same state. On Metal this is also the cheapest way to answer a single request
-    # (state once instead of once per question), so it is the only path `forward` / `probs` take.
-
-    def prefix(self, enc):
-        Ls = enc["seg"].count(0)
-        cache = make_prompt_cache(self.lm)
-        self._hidden([enc["ids"][:Ls]], cache)
-        return Ls, cache
-
-    def _branch_logits(self, enc, cache):
-        """Branches as rows on a replicated copy of the state cache, rows_per_pass rows (and cache copies) at a time."""
-        _, _, rows = rows_of(enc)
-        chunk, out = rows_per_pass([r["ids"] for r in rows], enc["seg"].count(0)), []
-        for start in range(0, len(rows), chunk):
-            part = rows[start:start + chunk]
-            batch = [type(c).merge([c] * len(part)) for c in cache]      # merge copies the arrays: `cache` is not mutated
-            h = self._hidden([r["ids"] for r in part], batch)
-            out += [self._logits(h[i], r["decide"], r["opts"]) for i, r in enumerate(part)]
-        return out
-
-    def _branch_probs(self, enc, cache):
-        return [F.softmax(z, -1) for z in self._branch_logits(enc, cache)]
-
-    def forward(self, enc):
-        """List of logits tensors, one per question."""
-        return self._branch_logits(enc, self.prefix(enc)[1])
-
-    def probs(self, enc):
-        return self._branch_probs(enc, self.prefix(enc)[1])
-
-    def probs_and_prefix(self, enc):
-        prefix = self.prefix(enc)
-        return self._branch_probs(enc, prefix[1]), prefix
-
-    def probs_with_prefix(self, enc, prefix):
-        Ls, cache = prefix
-        if enc["seg"].count(0) != Ls: raise ValueError("prefix does not match this record's state")
-        return self._branch_probs(enc, cache)

code/kev/model.py

@@ -1,5 +1,12 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2026 Jared Palmer
+# Derived from Kev (https://github.com/jaredpalmer/kev).
+# Modified by Nace.AI for Drex: bidirectional state
+# attention, and configurable serving limits.
+# See NOTICE and docs/upstream.md for attribution and revision information.
 """Decision model: causal LM backbone + block-causal branch mask + pointer readout."""
-import copy, math, os, re
+import copy, json, math, os, re
+from pathlib import Path
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -11,13 +18,33 @@ SPECIAL = ["<|fim_prefix|>", "<|fim_middle|>", "<|box_start|>", "<|box_end|>", "
 # training context: state tokens, tokens per question branch, and the whole packed record. Frozen suites are admitted with
 # this rule (kev.suite) and training applies it to records built on the fly, so train and eval see the same population.
 MAX_STATE, MAX_BRANCH, MAX_PACKED = 384, 1024, 2048
-# serving context (kev.serve): per-branch cap mirrors Jev's ~32k, bounded by the base model window; longer than training, so untested there
-SERVE_MAX_STATE = int(os.environ.get("KEV_SERVE_MAX_STATE", "8192"))
-SERVE_MAX_BRANCH = int(os.environ.get("KEV_SERVE_MAX_BRANCH", "8192"))
-SERVE_MAX_PACKED = SERVE_MAX_STATE + SERVE_MAX_BRANCH   # one row at most: a packed request longer than this runs in the row form (the block-causal mask is L x L)
-# the longest state a checkpoint may be trained on (kev.train --max_state) and still leave every question its training
-# branch budget when served: serving's row limit is SERVE_MAX_BRANCH = state + branch
-MAX_TRAIN_STATE = SERVE_MAX_BRANCH - (MAX_BRANCH - MAX_STATE)
+# Position window of the weights. Serving recommends 16,384 and can be raised to this (32,768 for non-hybrid backbones such as DLM).
+CONTEXT_LENGTH = 131072
+NON_HYBRID_CONTEXT_LENGTH = 32768
+RECOMMENDED_CONTEXT = 16384
+
+
+def _serve_limit(name, default):
+    """A serving cap from the environment. Unset keeps `default`. An invalid or too large value is an error."""
+    raw = os.environ.get(name)
+    if raw is None or raw == "":
+        return default
+    try:
+        value = int(raw)
+    except ValueError:
+        raise ValueError(f"{name}={raw!r} is not an integer") from None
+    if value < 1 or value > CONTEXT_LENGTH:
+        raise ValueError(f"{name}={value} must be between 1 and {CONTEXT_LENGTH}")
+    return value
+
+
+# KEV_CONTEXT selects the window. KEV_SERVE_MAX_STATE, KEV_SERVE_MAX_BRANCH, and KEV_SERVE_MAX_PACKED override one cap.
+_context = _serve_limit("KEV_CONTEXT", RECOMMENDED_CONTEXT)
+SERVE_MAX_STATE = _serve_limit("KEV_SERVE_MAX_STATE", _context)
+SERVE_MAX_BRANCH = _serve_limit("KEV_SERVE_MAX_BRANCH", _context)
+SERVE_MAX_PACKED = _serve_limit("KEV_SERVE_MAX_PACKED", _context)
+# Training admission stays on the historical 8,192-token row, so a wider serving window does not widen training.
+MAX_TRAIN_STATE = 8192 - (MAX_BRANCH - MAX_STATE)
 
 
 def training_context(max_state=MAX_STATE):
@@ -30,14 +57,12 @@ def training_context(max_state=MAX_STATE):
     return {"max_state": max_state, "max_branch": MAX_BRANCH + extra, "max_packed": MAX_PACKED + extra}
 
 
-def rows_per_pass(rows, prefix_len=0, budget=None):
+def rows_per_pass(rows, prefix_len=0, budget=SERVE_MAX_PACKED):
     """How many causal rows one inference forward pass takes: as many as fit `budget` tokens counting the cached state
     each row carries (prefix_len) plus its own tokens, at least one. Memory per pass is bounded by one maximal row however
     many questions a request has, and the rows are independent, so the answers do not depend on the split. Kev-4B on MLX,
     a 4.8k-token state with 64 questions: 24.7 GB peak in one pass, 8.9 GB one row at a time (and faster: the cost was
     64 copies of the state cache)."""
-    if budget is None:
-        budget = int(os.environ.get("KEV_ROW_BUDGET", str(SERVE_MAX_PACKED)))
     return max(1, budget // (prefix_len + max(len(r) for r in rows)))
 
 
@@ -47,7 +72,29 @@ class ContextOverflow(ValueError):
 
 
 def load_tokenizer(name, revision=None):
-    return AutoTokenizer.from_pretrained(name, revision=revision)
+    return AutoTokenizer.from_pretrained(name, revision=revision, trust_remote_code=False)
+
+
+def load_backbone(name, revision, dtype, attn):
+    """Load a local safetensors encoder without executing checkpoint Python code."""
+    cfg = AutoConfig.from_pretrained(name, revision=revision, trust_remote_code=False)
+    if getattr(cfg, "model_type", None) == "qwen3_5_moe":
+        # Qwen3.6-35B-A3B is a vision-language checkpoint: the language model sits at model.language_model and the
+        # vision tower is not used for decisions, so it is dropped before the weights move to the device.
+        from transformers import Qwen3_5MoeForConditionalGeneration
+
+        full = Qwen3_5MoeForConditionalGeneration.from_pretrained(
+            name, config=cfg, revision=revision, dtype=dtype, attn_implementation=attn,
+            trust_remote_code=False, use_safetensors=True, weights_only=True,
+        )
+        language_model = full.model.language_model
+        full.model.language_model = None
+        del full
+        return language_model
+    return AutoModelForCausalLM.from_pretrained(
+        name, config=cfg, revision=revision, dtype=dtype, attn_implementation=attn,
+        trust_remote_code=False, use_safetensors=True, weights_only=True,
+    ).model
 
 
 def pad_id(tok):
@@ -200,20 +247,6 @@ SCORING_INTERFACE = ("encode", "forward", "probs", "probs_and_prefix", "probs_wi
                      "head", "backend", "dtype", "device", "hybrid", "option_isolation", "prefix_min_tokens")
 
 
-def _load_backbone(name, revision, dtype, attn):
-    """The text backbone. Qwen3.6-35B-A3B is a vision-language MoE checkpoint: its language model lives at
-    model.language_model, and the vision tower is not used for decisions."""
-    cfg = AutoConfig.from_pretrained(name, revision=revision)
-    if getattr(cfg, "model_type", None) == "qwen3_5_moe":
-        from transformers import Qwen3_5MoeForConditionalGeneration
-        full = Qwen3_5MoeForConditionalGeneration.from_pretrained(name, revision=revision, dtype=dtype, attn_implementation=attn)
-        lm = full.model.language_model
-        full.model.language_model = None
-        del full
-        return lm
-    return AutoModelForCausalLM.from_pretrained(name, revision=revision, dtype=dtype, attn_implementation=attn).model
-
-
 class DecisionModel(nn.Module):
     def __init__(self, name, tok, device, lora=None, revision=None, attn=None, head_dim=256, option_isolation=False, special_embeddings=False, lora_targets="all", dtype=torch.float32, state_bidir=False):
         super().__init__()
@@ -221,7 +254,7 @@ class DecisionModel(nn.Module):
         # eager on MPS/CPU (known-good with our float 4D mask); SDPA on CUDA (accepts arbitrary additive masks).
         attn = attn or ("sdpa" if str(device).startswith("cuda") else "eager")
         # dtype: fp32 for training and exact evaluation; bf16 is a serving option for large backbones (8B on a 32 GB Mac)
-        self.lm = _load_backbone(name, revision, dtype, attn)
+        self.lm = load_backbone(name, revision, dtype, attn)
         self.pad_id = pad_id(tok)
         # hybrid backbones (Qwen3.5: Gated DeltaNet layers, recurrent) cannot honour the block-causal mask, so every
         # question runs as its own causal row continuing from the state (rows_of). Attention-only backbones keep the
@@ -317,9 +350,7 @@ class DecisionModel(nn.Module):
         rows_per_pass at a time; training keeps one batch (its batches are small and autograd needs the whole graph anyway).
         With `cache`, the rows are branches continuing the cached state: the cache is replicated once per chunk (a copy,
         so the caller's prefix stays pristine) and the cached tokens are marked real in the attention mask."""
-        # One pass is bounded by the serving row budget in training too. A 35B MoE forward of every question
-        # at once does not fit beside the frozen weights; rows are independent, so the split does not change the loss.
-        chunk = rows_per_pass([ids for ids, _ in rows], prefix_len)
+        chunk = len(rows) if self.training else rows_per_pass([ids for ids, _ in rows], prefix_len)
         out = []
         for start in range(0, len(rows), chunk):
             part = rows[start:start + chunk]

code/kev/serve.py

@@ -1,183 +0,0 @@
-"""FastAPI sidecar for the playground: loads one checkpoint, exposes prefill-only decisions.
-
-Run: uv run --extra serve python -m kev.serve --run runs/kev --port 8008
-
-TypeSafe-compatible: POST /v1/systemone, GET /v1/models, the `x-typesafe-request-id` response header, and bearer auth
-when KEV_API_KEY is set (unset = open server, the local default). Demo extras: POST /v1/systemone/permute (one Choice
-under several option orders) and POST /v1/systemone/separate (each question in its own pass, for the packed-vs-separate
-comparison). KEV_PREFIX_CACHE / KEV_PREFIX_MIN_TOKENS size the state-prefix cache; KEV_DATE_FACTS=1 opts into the
-date preprocessing (api.with_date_facts). Backend and precision follow LoadOptions (KEV_BACKEND, KEV_DTYPE, ...): on Apple
-Silicon the hybrid Qwen3.5 checkpoints run on MLX by default, elsewhere on torch in bf16.
-"""
-import argparse, hmac, os, random, threading, time, uuid
-import torch
-from dataclasses import dataclass, field, replace
-from fastapi import FastAPI, HTTPException
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import JSONResponse
-from pydantic import BaseModel, Field
-from .api import SystemOneRequest, to_record, to_answers, output_tokens, with_date_facts
-from .checkpoint import Checkpoint, LoadOptions, is_hub_id
-from .device import default_device, sync
-from .model import SERVE_MAX_BRANCH, SERVE_MAX_STATE
-
-PREFIX_CACHE_SIZE = int(os.environ.get("KEV_PREFIX_CACHE", "4"))          # states kept (KV + hidden); 0 disables
-PREFIX_MIN_TOKENS = os.environ.get("KEV_PREFIX_MIN_TOKENS")               # states shorter than this are not cached; default = the model's prefix_min_tokens (0 for hybrid backbones and MLX, 384 for attention-only torch models)
-DATE_FACTS = os.environ.get("KEV_DATE_FACTS", "0") == "1"
-API_KEY = os.environ.get("KEV_API_KEY")                                  # unset = open server; set = require Authorization: Bearer <key>, as the TypeSafe clients always send
-MODEL_NAMES = ("kev-latest", "jev-latest")                               # both names serve this checkpoint; jev-latest is the TypeSafe SDK default model, so an unconfigured client works
-
-
-@dataclass
-class Server:
-    """The loaded checkpoint and the state-prefix cache shared by every request (one model, one lock)."""
-    checkpoint: Checkpoint
-    tok: object
-    model: object
-    device: str
-    lock: threading.Lock = field(default_factory=threading.Lock)
-    prefix_cache: dict = field(default_factory=dict)   # (state token ids, option_isolation) -> prefix, in LRU order
-    prefix_hits: int = 0
-    prefix_misses: int = 0
-    release_date: str = field(default="")   # for the TypeSafe model card; resolved once (may ask the Hub)
-
-    def __post_init__(self):
-        self.release_date = self.release_date or self.checkpoint.release_date()
-
-    @property
-    def prefix_min_tokens(self):
-        return int(PREFIX_MIN_TOKENS) if PREFIX_MIN_TOKENS else self.model.prefix_min_tokens
-
-    def probs(self, rec):
-        """One forward pass. The state prefix (tokens up to the first question) is cached across requests, so a repeated
-        state only pays for its question branches. Exact: the state's activations do not depend on the branches."""
-        try: enc = self.model.encode(self.tok, rec, max_state=SERVE_MAX_STATE, max_branch=SERVE_MAX_BRANCH)
-        except ValueError as e: raise HTTPException(422, str(e))
-        Ls = enc["seg"].count(0); key = (tuple(enc["ids"][:Ls]), bool(enc.get("option_isolation")))
-        cache, hit = self.prefix_cache, False
-        with self.lock:
-            sync(self.device); t = time.time()
-            eligible = PREFIX_CACHE_SIZE and Ls >= self.prefix_min_tokens
-            if eligible and key in cache:
-                prefix = cache.pop(key)                            # pop + reinsert = LRU order
-                ps = self.model.probs_with_prefix(enc, prefix); cache[key] = prefix
-                self.prefix_hits += 1; hit = True
-            elif eligible:
-                ps, prefix = self.model.probs_and_prefix(enc)      # one pass, and the state prefix is kept for next time
-                cache[key] = prefix
-                while len(cache) > PREFIX_CACHE_SIZE: cache.pop(next(iter(cache)))
-                self.prefix_misses += 1
-            else:
-                ps = self.model.probs(enc)
-            sync(self.device); dt = time.time() - t
-        return [p.tolist() for p in ps], {"tokens": len(enc["ids"]), "state_tokens": Ls, "latency_ms": round(dt * 1000, 1), "prefix_cache_hit": hit}
-
-    def answer(self, req):
-        """The /v1/systemone response body for one request."""
-        rec, meta = to_record(prepare(req))
-        ps, m = self.probs(rec)
-        answers = to_answers(ps, meta)
-        return {"model": req.model, "answers": answers, "usage": {"input_tokens": m["tokens"], "output_tokens": output_tokens(self.tok, answers)}, "latency_ms": m["latency_ms"]}
-
-
-def prepare(req):
-    """Opt-in preprocessing applied to every request before the model sees it."""
-    return req.model_copy(update={"state": with_date_facts(req.state)}) if DATE_FACTS else req
-
-
-app = FastAPI(title="kev")
-app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], expose_headers=["x-typesafe-request-id"])
-
-
-@app.middleware("http")
-async def typesafe(request, call_next):
-    """Bearer auth (when API_KEY is set) and the request id every TypeSafe client reads off the response."""
-    if API_KEY and request.url.path.startswith("/v1") and not hmac.compare_digest(request.headers.get("authorization", ""), f"Bearer {API_KEY}"):
-        resp = JSONResponse({"detail": "missing or invalid API key; send Authorization: Bearer <KEV_API_KEY>"}, 401, {"www-authenticate": "Bearer"})
-    else:
-        resp = await call_next(request)
-    resp.headers["x-typesafe-request-id"] = request.headers.get("x-typesafe-request-id") or uuid.uuid4().hex
-    return resp
-
-
-def server() -> Server:
-    return app.state.server
-
-
-@app.post("/v1/systemone")
-def systemone(req: SystemOneRequest):
-    """TypeSafe-compatible endpoint: typed questions in, typed answers out, one prefill pass."""
-    return server().answer(req)
-
-
-class PermuteSystemOne(BaseModel):
-    request: SystemOneRequest
-    question: str
-    n_perm: int = Field(default=6, ge=1, le=64)   # each order is a forward pass; 0 divided by nothing, unbounded counts ran forever (#30)
-    seed: int = 0
-
-
-@app.post("/v1/systemone/permute")
-def systemone_permute(r: PermuteSystemOne):
-    """Re-run one Choice question under n_perm option orders. Returns per-order probabilities keyed by option name."""
-    q = r.request.questions.get(r.question)
-    if q is None or q.type != "choice": raise HTTPException(422, "question must be an existing choice question")
-    rng = random.Random(r.seed); keys = list(q.criteria); runs = []
-    for i in range(r.n_perm):
-        order = list(keys)
-        if i > 0: rng.shuffle(order)
-        one = r.request.model_copy(update={"questions": {r.question: q.model_copy(update={"criteria": {k: q.criteria[k] for k in order}})}})
-        resp = server().answer(one); a = resp["answers"][r.question]
-        runs.append({"order": order, "probabilities": a["probabilities"], "choice": a["choice"], "latency_ms": resp["latency_ms"]})
-    spread = {k: max(x["probabilities"][k] for x in runs) - min(x["probabilities"][k] for x in runs) for k in keys}
-    return {"runs": runs, "argmax_stable": len({x["choice"] for x in runs}) == 1, "spread": spread}
-
-
-@app.post("/v1/systemone/separate")
-def systemone_separate(req: SystemOneRequest):
-    """Answer each question in its own request against the same state (N passes). For packed-vs-separate comparison."""
-    parts = [server().answer(req.model_copy(update={"questions": {qid: q}})) for qid, q in req.questions.items()]
-    answers = {qid: a for p in parts for qid, a in p["answers"].items()}
-    return {"model": req.model, "answers": answers,
-            "usage": {"input_tokens": sum(p["usage"]["input_tokens"] for p in parts), "output_tokens": output_tokens(server().tok, answers)},
-            "latency_ms": round(sum(p["latency_ms"] for p in parts), 1)}
-
-
-@app.get("/v1/models")
-def models():
-    """One TypeSafe model card (name, description, release_date) per accepted model name, plus the Kev serving details
-    a client may ignore: the run, the base, the device, the backend and precision, the temperature, prefix-cache stats."""
-    s = server()
-    ck, meta = s.checkpoint, s.checkpoint.meta
-    card = {"description": f"Kev pointer head on {meta.base}, serving {ck.requested} at temperature {s.model.head.temperature:.2f}",
-            "release_date": s.release_date,
-            "run": ck.requested, "base": meta.base, "lora": meta.lora, "device": s.device, "backend": s.model.backend, "dtype": s.model.dtype,
-            "temperature": s.model.head.temperature,
-            "prefix_cache": {"size": PREFIX_CACHE_SIZE, "min_state_tokens": s.prefix_min_tokens, "hits": s.prefix_hits,
-                             "misses": s.prefix_misses, "cached_states": len(s.prefix_cache)}}
-    return {"models": [{"name": name, **card} for name in MODEL_NAMES]}
-
-
-def main():
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--run", default="runs/kev")
-    ap.add_argument("--fallback", default="runs/smoke")
-    ap.add_argument("--port", type=int, default=8008)
-    a = ap.parse_args()
-    run = a.run if is_hub_id(a.run) or os.path.exists(f"{a.run}/head.pt") else a.fallback
-    if run != a.run: print(f"{a.run} not found, falling back to {run}")
-    dev = default_device()
-    opts = LoadOptions.from_env()
-    if dev == "mps" and opts.attn is None: opts = replace(opts, attn="sdpa")   # serving default on Apple GPUs (parity measured)
-    if dev != "cpu" and opts.dtype is None: opts = replace(opts, dtype=torch.bfloat16)   # serving default: 2-4.5x faster than fp32 on an L4, same answers (LoadOptions.dtype); KEV_DTYPE=fp32 for the exact path
-    if opts.backend is None: opts = replace(opts, backend="auto")   # serving default: MLX for the hybrid Qwen3.5 checkpoints on Apple Silicon (LoadOptions.backend); KEV_BACKEND=torch to decline
-    ck = Checkpoint(run)
-    tok, model = ck.load(dev, opts)
-    app.state.server = Server(ck, tok, model, dev)
-    print(f"serving {ck.requested} ({ck.path}) on {dev} via {model.backend} ({model.dtype}) :{a.port}")   # /v1/models reports the run as given, not the resolved cache path
-    import uvicorn
-    uvicorn.run(app, host="127.0.0.1", port=a.port)
-
-
-if __name__ == "__main__":
-    main()

code/kev/serve_batch.py

@@ -1,85 +0,0 @@
-"""Same /v1/systemone server as kev.serve, but concurrent requests are batched into shared forward passes.
-
-The reference server serializes requests behind one lock (one forward per request, ~100-250 ms of mostly Python/kernel-launch
-overhead on the 35B-A3B MoE). Here a batcher thread gathers whatever arrived within a few ms, groups records of similar
-length under a token budget and runs DecisionModel.forward_batch once per group (right-padded causal rows, so the answers match
-the one-at-a-time path up to bf16 noise). Prefix cache is skipped (single-question records recompute nothing extra).
-usage: python -m kev.serve_batch --run <merged ckpt> --port 8100   env: KEV_BATCH_WAIT_MS (4), KEV_BATCH_MAX (48), KEV_ROW_BUDGET (tokens/pass)
-"""
-import os, queue, threading, time
-import torch
-import torch.nn.functional as F
-from fastapi import HTTPException
-from . import serve as S
-from .model import SERVE_MAX_BRANCH, SERVE_MAX_STATE
-
-WAIT = float(os.environ.get("KEV_BATCH_WAIT_MS", "4")) / 1000
-MAXB = int(os.environ.get("KEV_BATCH_MAX", "48"))
-BUDGET = int(os.environ.get("KEV_ROW_BUDGET", "32768"))
-Q = queue.Queue()
-STATS = {"batches": 0, "records": 0, "max_batch": 0}
-
-
-def _groups(items):
-    """items sorted by token length; yield groups with n * max_len <= BUDGET (at least one item)."""
-    items = sorted(items, key=lambda it: len(it["enc"]["ids"]))
-    cur = []
-    for it in items:
-        L = len(it["enc"]["ids"])
-        if cur and (len(cur) + 1) * L > BUDGET: yield cur; cur = []
-        cur.append(it)
-    if cur: yield cur
-
-
-def _run(model, group):
-    encs = [it["enc"] for it in group]
-    try:
-        with torch.no_grad():
-            out = model.forward_batch(encs)
-            res = [[F.softmax(z, -1).cpu().tolist() for z in zs] for zs in out]
-        for it, r in zip(group, res): it["res"] = r
-    except Exception as e:   # e.g. OOM on a big group: retry one by one so a single bad request cannot sink the batch
-        torch.cuda.empty_cache()
-        for it in group:
-            try:
-                with torch.no_grad(): it["res"] = [F.softmax(z, -1).cpu().tolist() for z in model.forward(it["enc"])]
-            except Exception as e2: it["err"] = e2
-    for it in group: it["ev"].set()
-
-
-def batcher(model):
-    while True:
-        first = Q.get(); items = [first]; t_end = time.time() + WAIT
-        while len(items) < MAXB:
-            left = t_end - time.time()
-            if left <= 0: break
-            try: items.append(Q.get(timeout=left))
-            except queue.Empty: break
-        STATS["batches"] += 1; STATS["records"] += len(items); STATS["max_batch"] = max(STATS["max_batch"], len(items))
-        for g in _groups(items): _run(model, g)
-
-
-def probs(self, rec):
-    try: enc = self.model.encode(self.tok, rec, max_state=SERVE_MAX_STATE, max_branch=SERVE_MAX_BRANCH)
-    except ValueError as e: raise HTTPException(422, str(e))
-    it = {"enc": enc, "ev": threading.Event()}; t = time.time(); Q.put(it); it["ev"].wait()
-    if "err" in it: raise HTTPException(500, f"{type(it['err']).__name__}: {str(it['err'])[:200]}")
-    return it["res"], {"tokens": len(enc["ids"]), "state_tokens": enc["seg"].count(0), "latency_ms": round((time.time() - t) * 1000, 1), "prefix_cache_hit": False}
-
-
-S.Server.probs = probs
-
-
-@S.app.get("/batch_stats")
-def batch_stats(): return STATS
-
-
-@S.app.on_event("startup")
-async def _startup():
-    import anyio
-    anyio.to_thread.current_default_thread_limiter().total_tokens = 512   # many concurrent clients -> many waiting handler threads
-    threading.Thread(target=batcher, args=(S.server().model,), daemon=True).start()
-
-
-if __name__ == "__main__":
-    S.main()

code/pyproject.toml

@@ -1,24 +0,0 @@
-[build-system]
-requires = ["setuptools>=77"]
-build-backend = "setuptools.build_meta"
-
-[project]
-name = "kev"
-version = "0.1.0"
-description = "Decision-model loader and TypeSafe-compatible /v1/systemone server"
-requires-python = ">=3.12"
-dependencies = [
-    "accelerate>=1.15.0",
-    "numpy>=2.0",
-    "peft>=0.21",
-    "pydantic>=2.9",
-    "torch>=2.6",
-    "transformers>=5.17,<6",
-    "huggingface_hub",
-]
-
-[project.optional-dependencies]
-serve = ["fastapi>=0.115", "uvicorn>=0.30"]
-
-[tool.setuptools]
-packages = ["kev"]

examples/request.json

@@ -0,0 +1,29 @@
+{
+  "state": {
+    "ticket": "I was charged twice for the same order. Please refund the extra payment."
+  },
+  "questions": {
+    "team": {
+      "type": "choice",
+      "instructions": "Which team should handle this ticket?",
+      "criteria": {
+        "billing": "Payments, charges, and refunds",
+        "technical": "Bugs and outages",
+        "other": "Anything else"
+      }
+    },
+    "refund": {
+      "type": "noul",
+      "instructions": "Does the customer explicitly ask for a refund?"
+    },
+    "urgency": {
+      "type": "score",
+      "instructions": "How urgent is this ticket?",
+      "criteria": [
+        "Routine",
+        "Soon",
+        "Urgent"
+      ]
+    }
+  }
+}

inference.py

@@ -1,61 +1,150 @@
-"""Score System One decisions with Drex v2.0.
+"""Score a System One request with Drex v2.0.
 
-    pip install -e code
-    python inference.py --model .                                  # built-in example
-    python inference.py --model NischayDnk/Drex-v2.0 --request req.json
+    pip install -r requirements.txt
+    python inference.py
+    python inference.py --request examples/request.json
 
-A request is TypeSafe System One JSON: {"state": ..., "questions": {qid: {"type": "noul"|"choice"|"score", ...}}}.
-The model returns a probability for every option in one forward pass per question; nothing is generated.
+The model returns a probability for every option in one forward pass.
+A state or question past the context limit returns an error, so the score covers the whole document.
 """
 import argparse
 import json
+import math
 import sys
+import time
 from pathlib import Path
 
 import torch
 
 sys.path.insert(0, str(Path(__file__).resolve().parent / "code"))
-from kev.api import SystemOneRequest, to_answers, to_record  # noqa: E402
-from kev.checkpoint import Checkpoint, LoadOptions  # noqa: E402
-from kev.model import SERVE_MAX_BRANCH, SERVE_MAX_STATE  # noqa: E402
+from kev.api import SystemOneRequest, output_tokens, to_answers, to_record  # noqa: E402
+from kev.checkpoint import load as _load  # noqa: E402
+from kev.model import NON_HYBRID_CONTEXT_LENGTH, SERVE_MAX_BRANCH, SERVE_MAX_PACKED, SERVE_MAX_STATE  # noqa: E402
 
 EXAMPLE = {
-    "state": {"premise": "The lawyer and the judge supported the students .", "hypothesis": "The judge supported the students ."},
+    "state": {
+        "ticket": "I was charged twice for the same order. Please refund the extra payment.",
+    },
     "questions": {
-        "entails": {"type": "noul", "instructions": "If the premise is true, must the hypothesis also be true?"},
-        "topic": {"type": "choice", "instructions": "What is the setting?", "criteria": {"court": "A legal setting", "sports": "A sports event", "kitchen": "Cooking"}},
+        "team": {
+            "type": "choice",
+            "instructions": "Which team should handle this ticket?",
+            "criteria": {
+                "billing": "Payments, charges, and refunds",
+                "technical": "Bugs and outages",
+                "other": "Anything else",
+            },
+        },
+        "refund": {
+            "type": "noul",
+            "instructions": "Does the customer explicitly ask for a refund?",
+        },
+        "urgency": {
+            "type": "score",
+            "instructions": "How urgent is this ticket?",
+            "criteria": ["Routine", "Soon", "Urgent"],
+        },
     },
 }
 
 
-def load(model):
-    path = Path(model)
-    if not path.exists():
-        from huggingface_hub import snapshot_download
-        path = Path(snapshot_download(model))
-    ck = Checkpoint(str(path))
-    ck.meta.base = str(path)
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    tok, m = ck.load(device, LoadOptions(dtype=torch.bfloat16 if device == "cuda" else None))
-    return tok, m.eval()
+def load(path):
+    tok, model = _load(path)
+    if not model.hybrid and max(SERVE_MAX_STATE, SERVE_MAX_BRANCH, SERVE_MAX_PACKED) > NON_HYBRID_CONTEXT_LENGTH:
+        raise ValueError(f"this model supports at most {NON_HYBRID_CONTEXT_LENGTH} tokens; lower KEV_CONTEXT")
+    return tok, model
+
+
+def _encode(tok, model, request):
+    if not isinstance(request, dict) or "requests" in request:
+        raise ValueError("each decision must be an object with state and questions, without a requests wrapper")
+    rec, meta = to_record(SystemOneRequest.model_validate(request))
+    enc = model.encode(tok, rec, max_state=SERVE_MAX_STATE, max_branch=SERVE_MAX_BRANCH, strict=True)
+    return enc, meta
+
+
+def _result(tok, enc, meta, probs, latency_ms):
+    if len(probs) != len(meta) or any(len(row) != len(question["keys"]) or
+                                      any(not math.isfinite(value) for value in row)
+                                      for row, question in zip(probs, meta)):
+        raise RuntimeError("model readout is incomplete or non-finite")
+    answers = to_answers(probs, meta)
+    return {
+        "answers": answers,
+        "usage": {"input_tokens": len(enc["ids"]), "output_tokens": _run_scorer(output_tokens, tok, answers)},
+        "latency_ms": latency_ms,
+    }
+
+
+def _run_scorer(fn, *args):
+    try:
+        return fn(*args)
+    except ValueError as exc:
+        raise RuntimeError("model scoring failed") from exc
 
 
 def decide(tok, model, request):
-    rec, meta = to_record(SystemOneRequest(state=request["state"], questions=request["questions"]))
-    enc = model.encode(tok, rec, max_state=SERVE_MAX_STATE, max_branch=SERVE_MAX_BRANCH)
+    enc, meta = _encode(tok, model, request)
+    started = time.perf_counter()
     with torch.no_grad():
-        ps = [p.float().tolist() for p in model.probs(enc)]
-    return to_answers(ps, meta)
+        probs = [p.float().tolist() for p in _run_scorer(model.probs, enc)]
+    return _result(tok, enc, meta, probs, round((time.perf_counter() - started) * 1000, 1))
+
+
+def decide_batch(tok, model, requests):
+    if not isinstance(requests, list) or not requests:
+        raise ValueError("requests must be a non-empty list")
+    if len(requests) > 1024:
+        raise ValueError("requests exceeds 1024 items")
+    prepared = [_encode(tok, model, request) for request in requests]
+    results = [None] * len(prepared)
+    index = 0
+    while index < len(prepared):
+        enc, _meta = prepared[index]
+        if len(enc["ids"]) > SERVE_MAX_PACKED:
+            started = time.perf_counter()
+            with torch.no_grad():
+                probs = [p.float().tolist() for p in _run_scorer(model.probs, enc)]
+            results[index] = _result(tok, enc, _meta, probs, round((time.perf_counter() - started) * 1000, 1))
+            index += 1
+            continue
+        end = index
+        longest = 0
+        # Bound padded tokens and the quadratic attention mask before batching.
+        while end < len(prepared) and end - index < 8:
+            size = len(prepared[end][0]["ids"])
+            if size > SERVE_MAX_PACKED:
+                break
+            padded = (max(longest, size) + 127) // 128 * 128
+            count = end - index + 1
+            if count * padded > SERVE_MAX_PACKED or count * padded * padded > SERVE_MAX_PACKED ** 2:
+                break
+            longest = max(longest, size)
+            end += 1
+        if end == index:
+            end = index + 1
+        group = prepared[index:end]
+        started = time.perf_counter()
+        with torch.no_grad():
+            logits = _run_scorer(model.forward_batch, [item[0] for item in group])
+        latency_ms = round((time.perf_counter() - started) * 1000, 1)
+        if len(logits) != len(group):
+            raise RuntimeError("batch readout did not return one result per request")
+        for offset, ((enc, meta), row) in enumerate(zip(group, logits)):
+            probs = [torch.softmax(item.float(), dim=-1).tolist() for item in row]
+            results[index + offset] = _result(tok, enc, meta, probs, latency_ms)
+        index = end
+    return results
 
 
 def main():
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--model", default=str(Path(__file__).resolve().parent))
-    ap.add_argument("--request", help="JSON file with one System One request (default: a built-in example)")
-    a = ap.parse_args()
-    tok, model = load(a.model)
-    req = json.load(open(a.request)) if a.request else EXAMPLE
-    print(json.dumps(decide(tok, model, req), indent=2))
+    parser = argparse.ArgumentParser(description="Score one System One request with a Drex model.")
+    parser.add_argument("--model", default=str(Path(__file__).resolve().parent), help="Checkpoint directory (default: this directory)")
+    parser.add_argument("--request", help="JSON file with one System One request")
+    args = parser.parse_args()
+    tok, model = load(args.model)
+    request = json.loads(Path(args.request).read_text()) if args.request else EXAMPLE
+    print(json.dumps(decide(tok, model, request), indent=2))
 
 
 if __name__ == "__main__":

requirements.txt

@@ -0,0 +1,4 @@
+torch>=2.6,<2.9
+transformers>=5.17,<6
+pydantic>=2
+huggingface_hub>=0.34

serve.py

@@ -0,0 +1,107 @@
+"""Local System One server for Drex v2.0.
+
+    python serve.py --port 8000
+
+POST /v1/systemone accepts one decision, or ``requests``: a list of decisions.
+Several decisions that fit share one forward. This process scores the pointer head
+of the checkpoint.
+"""
+import argparse
+import json
+import threading
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from pathlib import Path
+
+from inference import decide, decide_batch, load
+
+
+def reject_nonfinite_json(value):
+    raise ValueError(f"non-finite JSON number: {value}")
+
+
+def make_handler(tok, model, model_name):
+    lock = threading.Lock()
+
+    class Handler(BaseHTTPRequestHandler):
+        protocol_version = "HTTP/1.1"
+
+        def log_message(self, fmt, *args):
+            print("%s %s" % (self.address_string(), fmt % args))
+
+        def _send(self, status, payload):
+            body = json.dumps(payload).encode()
+            self.send_response(status)
+            self.send_header("Content-Type", "application/json")
+            self.send_header("Content-Length", str(len(body)))
+            if self.close_connection:
+                self.send_header("Connection", "close")
+            self.end_headers()
+            self.wfile.write(body)
+
+        def do_GET(self):
+            if self.path.rstrip("/") == "/health":
+                self._send(200, {"status": "ok", "model": model_name})
+                return
+            self._send(404, {"error": "not found"})
+
+        def do_POST(self):
+            if self.path.rstrip("/") != "/v1/systemone":
+                self._send(404, {"error": "not found"})
+                return
+            try:
+                length = int(self.headers.get("Content-Length", "0"))
+            except ValueError:
+                self.close_connection = True
+                self._send(400, {"error": "Content-Length must be an integer"})
+                return
+            if length <= 0 or length > 1_000_000:
+                self.close_connection = True
+                self._send(400, {"error": "request body must be between 1 byte and 1 MB"})
+                return
+            try:
+                request = json.loads(self.rfile.read(length), parse_constant=reject_nonfinite_json)
+                if not isinstance(request, dict):
+                    raise ValueError("request must be an object")
+                if "model" in request and not isinstance(request["model"], str):
+                    raise ValueError("model must be a string")
+                with lock:
+                    if "requests" in request:
+                        if "state" in request or "questions" in request:
+                            raise ValueError("a batch cannot also contain state or questions")
+                        items = request["requests"]
+                        if not isinstance(items, list) or not items:
+                            raise ValueError("requests must be a non-empty list")
+                        if len(items) > 1024:
+                            raise ValueError("requests exceeds 1024 items")
+                        results = decide_batch(tok, model, items)
+                        payload = {"model": request.get("model") or model_name, "results": results}
+                    else:
+                        payload = {"model": request.get("model") or model_name, **decide(tok, model, request)}
+            except ValueError as exc:
+                self._send(400, {"error": str(exc)})
+                return
+            except Exception as exc:
+                self.log_error("System One scoring failed: %r", exc)
+                self._send(500, {"error": "internal scoring error"})
+                return
+            self._send(200, payload)
+
+    return Handler
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Serve Drex v2.0 at POST /v1/systemone.")
+    parser.add_argument("--model", default=str(Path(__file__).resolve().parent), help="Checkpoint directory (default: this directory)")
+    parser.add_argument("--host", default="127.0.0.1")
+    parser.add_argument("--port", type=int, default=8000)
+    parser.add_argument("--name", default="drex-v2.0")
+    args = parser.parse_args()
+    name = args.name
+    tok, model = load(args.model)
+    server = ThreadingHTTPServer((args.host, args.port), make_handler(tok, model, name))
+    print(f"{name} listening on http://{args.host}:{args.port}/v1/systemone")
+    server.serve_forever()
+
+
+if __name__ == "__main__":
+    main()