phocinae-server 0.1.0__py3-none-any.whl

phocinae/__init__.py

@@ -0,0 +1,8 @@
+"""phocinae-server: self-contained typed-decisions (system-one) inference server.
+
+Zero runtime dependencies beyond fastapi + uvicorn + torch:
+tokenizer, safetensors loading, mmBERT encoder and decision head are all
+implemented here in pure python/torch (no laya, no transformers, no tokenizers).
+"""
+
+__version__ = "0.1.0"

phocinae/engine.py

@@ -0,0 +1,423 @@
+"""Inference engine: sequence construction, batched forward, decoding.
+
+The sequence construction mirrors laya's build_sequence protocol exactly
+(head "[qtype] question: <ins>" + [MASK]-prefixed option markers + state +
+EOS separators), so the server serves the released checkpoint with the same
+input distribution it was trained with. P0 questions are mapped onto that
+protocol:
+
+  noul   -> options "false: no, the statement does not hold" /
+            "true: yes, the statement holds"; answer = P(true) >= threshold
+  choice -> options rendered as "<i>: <option text>" with the P0 list index
+            as the label; answer = index of the argmax option
+  score  -> 9 ordinal levels "level 2" .. "level 10";
+            answer = 2 + argmax (an integer in 2..10)
+
+Permutation averaging (PHOC_PERM_AVG / /v1/systemone/permute) follows the
+canonical protocol: K=4 option orders (original, reversed, two seeded
+shuffles, seed 0) with choice probabilities averaged before argmax.
+"""
+
+import json
+import logging
+import math
+import os
+import random
+import threading
+import time
+
+import torch
+
+from .model import DecisionModel, MMBertEncoder
+from .tokenizer import META, GemmaTokenizer
+from .safetensors_lite import load_safetensors
+
+log = logging.getLogger("phocinae.engine")
+
+QTYPES = {"choice": 0, "score": 1, "noul": 2}
+
+DEFAULT_INS = {
+    "choice": "Pick the option that best matches the statement.",
+    "score": "Rate the statement.",
+    "noul": "Does the statement hold?",
+}
+NOUL_OPTIONS = [
+    "false: no, the statement does not hold",
+    "true: yes, the statement holds",
+]
+SCORE_LEVELS = 9          # answers 2..10
+MAX_OPTIONS = 255
+MAX_QUESTIONS = 64
+MAX_POS = 8192            # encoder max_position_embeddings
+PERM_K = 4                # canonical permutation averaging order count
+BUDGET_TIERS = [(192, 512), (512, 1024), (1024, 2048), (2048, 4096),
+                (4096, MAX_POS), (MAX_POS, MAX_POS)]
+
+
+class Engine:
+    def __init__(self, model_dir, device="auto", model_name="Phocinae-Largha-150M-v1",
+                 perm_avg=False, warm=True, threads=None):
+        self.model_name = model_name
+        self.perm_avg = perm_avg
+        if device in ("auto", None, ""):
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = torch.device(device)
+        self.is_cuda = self.device.type == "cuda"
+        self.dtype = torch.float16 if self.is_cuda else torch.float32
+        if not self.is_cuda and threads:
+            torch.set_num_threads(int(threads))
+        self._lock = threading.Lock()
+        # Constant question-head texts recur on every request: cache once.
+        self._head_cache = {}
+        # Per-run encode memo for option strings (perm orders repeat them).
+        self._opt_cache = {}
+        # Lazily torch.compile'd model on CUDA (reduce-overhead, dynamic shapes);
+        # falls back to eager if compilation fails.
+        self._compiled = None
+        self._compile_failed = False
+
+        with open(os.path.join(model_dir, "encoder", "config.json"),
+                  encoding="utf-8") as fh:
+            ecfg = json.load(fh)
+        with open(os.path.join(model_dir, "rl_agent_config.json"),
+                  encoding="utf-8") as fh:
+            self.cfg = json.load(fh)
+        self.max_len = int(self.cfg.get("max_len", 512))
+        self.head_max_len = int(self.cfg.get("head_max_len", 192))
+
+        log.info("loading tokenizer from %s", model_dir)
+        t0 = time.time()
+        self.tok = GemmaTokenizer(os.path.join(model_dir, "tokenizer"))
+        self._verify_vocab()
+        log.info("tokenizer ready in %.2fs", time.time() - t0)
+
+        log.info("loading weights from %s", os.path.join(model_dir, "model.safetensors"))
+        t0 = time.time()
+        full_layers = {i for i, t in enumerate(ecfg["layer_types"])
+                       if t == "full_attention"}
+        encoder = MMBertEncoder(
+            vocab=ecfg["vocab_size"],
+            d=ecfg["hidden_size"],
+            n_layers=ecfg["num_hidden_layers"],
+            heads=ecfg["num_attention_heads"],
+            inter=ecfg["intermediate_size"],
+            eps=float(ecfg.get("layer_norm_eps", 1e-5)),
+            full_layers=full_layers,
+            theta=float(ecfg["rope_parameters"]["full_attention"]["rope_theta"]),
+            local_window=int(ecfg["local_attention"]) // 2,
+        )
+        head = DecisionModel(
+            d=ecfg["hidden_size"],
+            head_layers=int(self.cfg.get("head_layers", 2)),
+            heads=ecfg["hidden_size"] // 64,
+            dropout=float(self.cfg.get("head_dropout", 0.1)),
+        )
+        head.encoder = encoder
+        self.model = head
+        weights = load_safetensors(os.path.join(model_dir, "model.safetensors"))
+        # The checkpoint's "temperature" tensor is a dummy (all ones); the
+        # canonical calibrated temperatures ship in rl_agent_config.json.
+        weights.pop("temperature", None)
+        self.temperature = [float(t) for t in self.cfg.get("temperature",
+                                                           [1.0, 1.0, 1.0])]
+        if len(self.temperature) != 3:
+            raise RuntimeError("expected temperature vector of size 3")
+        weights = {k: v.to(self.dtype) for k, v in weights.items()}
+        missing, unexpected = self.model.load_state_dict(weights, strict=False)
+        if missing:
+            raise RuntimeError("weights missing for: %s" % ", ".join(sorted(missing)))
+        if unexpected:
+            raise RuntimeError("unexpected weight keys: %s" % ", ".join(sorted(unexpected)))
+        self.model = self.model.to(self.device).eval()
+        log.info("model loaded in %.2fs (%s, %d params)",
+                 time.time() - t0, self.device, sum(p.numel() for p in self.model.parameters()))
+        if warm:
+            self.warmup()
+
+    # ------------------------------------------------------------ vocabulary
+    def _verify_vocab(self):
+        anchors = {0: "<pad>", 1: "<eos>", 2: "<bos>", 3: "<unk>", 4: "<mask>",
+                   476: META + "a", 108: "\n", 235248: META}
+        for i, want in anchors.items():
+            got = self.tok.id2tok.get(i)
+            if got != want:
+                raise RuntimeError("vocab anchor mismatch: id %d -> %r, expected %r"
+                                   % (i, got, want))
+        probes = {
+            "hello": [25612],
+            "hello world": [25612, 2134],
+            "the the the": [573, 573, 573],
+            "a\nb": [476, 108, 518],
+            "a": [476],
+            " \t": [235248, 226],
+        }
+        for text, want in probes.items():
+            got = self.tok.encode(text)
+            if got != want:
+                raise RuntimeError("vocab encode mismatch for %r: got %r want %r"
+                                   % (text, got, want))
+        rt = "hello world, 中文测试 🚀"
+        ids = self.tok.encode(rt)
+        back = self.tok.decode(ids)
+        if back.replace(" ", "") != rt.replace(" ", ""):
+            raise RuntimeError("vocab round-trip failed: %r -> %r" % (rt, back))
+        log.info("vocabulary decode-verified (%d tokens, %d merges)",
+                 len(self.tok.vocab), len(self.tok.merges))
+
+    # ---------------------------------------------------------- warmup buckets
+    def warmup(self):
+        buckets = [(1, 64, 2), (1, 128, 4), (2, 128, 4), (4, 128, 4), (4, 256, 9),
+                   (8, 256, 9), (16, 256, 9), (16, 512, 16), (32, 512, 16), (64, 512, 32)]
+        log.info("warming shape buckets: %s", buckets)
+        t0 = time.time()
+        with torch.no_grad():
+            for B, L, K in buckets:
+                b = {
+                    "input_ids": torch.randint(5, 2000, (B, L), dtype=torch.long),
+                    "attention_mask": torch.ones((B, L), dtype=torch.long),
+                    "marker_pos": torch.randint(0, L, (B, K), dtype=torch.long),
+                    "marker_mask": torch.ones((B, K), dtype=torch.bool),
+                    "qtype": torch.zeros((B,), dtype=torch.long),
+                }
+                self._forward(b)
+                if self.is_cuda:
+                    torch.cuda.synchronize()
+        log.info("warmup done in %.2fs", time.time() - t0)
+
+    # ------------------------------------------------------------ inference
+    def _amp_ctx(self):
+        if self.is_cuda:
+            return torch.autocast(self.device.type, dtype=torch.float16)
+        return torch.autocast("cpu", enabled=False)
+
+    def _forward(self, b):
+        with torch.no_grad(), self._amp_ctx():
+            m = self.model
+            if self.is_cuda and not self._compile_failed:
+                if self._compiled is None:
+                    try:
+                        self._compiled = torch.compile(
+                            self.model, mode="reduce-overhead", dynamic=True)
+                        log.info("torch.compile active (reduce-overhead, dynamic)")
+                    except Exception as exc:  # pragma: no cover
+                        self._compile_failed = True
+                        log.warning("torch.compile unavailable: %s", exc)
+                m = self._compiled if self._compiled is not None else self.model
+            logits, act = m(
+                b["input_ids"].to(self.device),
+                b["attention_mask"].to(self.device),
+                b["marker_pos"].to(self.device),
+                b["marker_mask"].to(self.device),
+                b["qtype"].to(self.device),
+            )
+        return logits.float().cpu(), torch.softmax(act.float(), -1).cpu()
+
+    # ------------------------------------------------------- sequence building
+    def _question_spec(self, q):
+        t = q["type"]
+        ins = DEFAULT_INS[t]
+        if t == "noul":
+            return t, ins, list(NOUL_OPTIONS)
+        if t == "choice":
+            return t, ins, [str(o) for o in q["options"]]
+        return t, ins, ["level %d" % (2 + i) for i in range(SCORE_LEVELS)]
+
+    def _build_sequence(self, spec, state_ids, order, head_max_len, max_len):
+        t, ins, opts = spec
+        head_ids = self._head_cache.get((t, ins))
+        if head_ids is None:
+            head_ids = self.tok.encode("%s question: %s" % (t, ins))
+            self._head_cache[(t, ins)] = head_ids
+        opt_ids = []
+        for i in order:
+            o = opts[i].replace(self.tok.mask_token, " ")
+            ot = self._opt_cache.get(o)
+            if ot is None:
+                ot = self.tok.encode(" " + o)[:48]
+                self._opt_cache[o] = ot
+                if len(self._opt_cache) > 65536:
+                    self._opt_cache.clear()
+            opt_ids.append([self.tok.mask_id] + ot)
+        opt_budget = head_max_len - sum(len(o) for o in opt_ids)
+        if opt_budget < 16:
+            per = max(4, (head_max_len - 16) // max(1, len(opt_ids)))
+            opt_ids = [o[:per] for o in opt_ids]
+            opt_budget = head_max_len - sum(len(o) for o in opt_ids)
+        head_ids = head_ids[:max(8, opt_budget)]
+        ids = [self.tok.cls_token_id] + head_ids + [self.tok.sep_token_id]
+        markers = []
+        for o in opt_ids:
+            markers.append(len(ids))
+            ids.extend(o)
+        ids.append(self.tok.sep_token_id)
+        room = max(0, max_len - len(ids) - 1)
+        ids = ids + state_ids[:room] + [self.tok.sep_token_id]
+        return ids[:max_len], [m for m in markers if m < max_len]
+
+    def _build_sequence_fit(self, spec, state_ids, order=None):
+        k = len(spec[2])
+        if order is None:
+            order = list(range(k))
+        for hml, ml in BUDGET_TIERS:
+            seq, markers = self._build_sequence(spec, state_ids, order, hml, ml)
+            if len(markers) == k:
+                return seq, markers
+        raise ValueError("options exceed the %d-token model budget" % MAX_POS)
+
+    def _encode_request(self, state, questions):
+        state_ids = self.tok.encode(state.replace(self.tok.mask_token, " "))
+        items = []
+        for q in questions:
+            spec = self._question_spec(q)
+            seq, markers = self._build_sequence_fit(spec, state_ids)
+            items.append({
+                "ids": seq,
+                "markers": markers,
+                "qtype": QTYPES[q["type"]],
+                "n_opts": len(spec[2]),
+                "spec": spec,
+                "q": q,
+            })
+        return state_ids, items
+
+    # ------------------------------------------------------------- collation
+    def _collate(self, items):
+        n = len(items)
+        L = max(len(it["ids"]) for it in items)
+        K = max(len(it["markers"]) for it in items)
+        ids = [[self.tok.pad_id] * L for _ in range(n)]
+        att = [[0] * L for _ in range(n)]
+        mpos = [[0] * K for _ in range(n)]
+        mmask = [[False] * K for _ in range(n)]
+        for i, it in enumerate(items):
+            ids[i][:len(it["ids"])] = it["ids"]
+            att[i][:len(it["ids"])] = [1] * len(it["ids"])
+            for j, m in enumerate(it["markers"]):
+                mpos[i][j] = m
+                mmask[i][j] = True
+        return {
+            "input_ids": torch.tensor(ids, dtype=torch.long),
+            "attention_mask": torch.tensor(att, dtype=torch.long),
+            "marker_pos": torch.tensor(mpos, dtype=torch.long),
+            "marker_mask": torch.tensor(mmask, dtype=torch.bool),
+            "qtype": torch.tensor([it["qtype"] for it in items], dtype=torch.long),
+        }
+
+    # --------------------------------------------------------------- decoding
+    @staticmethod
+    def _probs(logits_row, k, t_scale):
+        z = [float(v) / t_scale for v in logits_row[:k]]
+        m = max(z)
+        p = [math.exp(v - m) for v in z]
+        s = sum(p)
+        return [v / s for v in p]
+
+    def _decode_rows(self, logits, act, items):
+        """Returns per-item (answer_value, confidence, act_probability).
+
+        confidence follows laya's canonical answer_confidence:
+        float(clip(max(p[:k]), 0, 1)) — the calibrated top probability.
+        """
+        out = []
+        for j, it in enumerate(items):
+            k = it["n_opts"]
+            t = it["q"]["type"]
+            t_scale = self.temperature[QTYPES[t]]
+            p = self._probs(logits[j], k, t_scale)
+            conf = round(float(min(1.0, max(p[:k]))), 4)
+            act_prob = round(float(act[j, 0]), 4)
+            if t == "choice":
+                ans = max(range(k), key=lambda i: p[i])
+            elif t == "score":
+                ans = 2 + max(range(k), key=lambda i: p[i])
+            else:
+                thr = it["q"].get("threshold")
+                thr = 0.5 if thr is None else float(thr)
+                ans = bool(p[1] >= thr)
+            out.append((ans, conf, act_prob, p))
+        return out
+
+    # ------------------------------------------------------------- permutation
+    def _perm_orders(self, item, rng):
+        k = item["n_opts"]
+        if item["q"]["type"] != "choice" or k < 2:
+            return [list(range(k))]
+        o1 = list(reversed(range(k)))
+        o2 = list(range(k))
+        rng.shuffle(o2)
+        o3 = list(range(k))
+        rng.shuffle(o3)
+        return [list(range(k)), o1, o2, o3]
+
+    # ----------------------------------------------------------------- run
+    def run(self, state, questions, perm=None, usage=True, with_scores=False):
+        """Run one system-one request. Returns (answers, confidence, action, usage).
+
+        perm: None -> engine default; False -> single pass; True -> permute-avg.
+        with_scores: True -> also return the per-option calibrated probability
+        vectors as a 5th element: {qid: [p0, p1, ...]} in canonical option
+        order (choice: request option order; noul: [false, true];
+        score: levels 2..10). Choice vectors are averaged over permuted
+        orders when perm=True.
+        """
+        if perm is None:
+            perm = self.perm_avg
+        with self._lock:
+            state_ids, items = self._encode_request(state, questions)
+            rng = random.Random(0)
+            rows = []
+            row_meta = []  # (question_idx, order)
+            for qi, it in enumerate(items):
+                orders = self._perm_orders(it, rng) if perm \
+                    else [list(range(it["n_opts"]))]
+                for order in orders:
+                    seq, markers = self._build_sequence_fit(it["spec"], state_ids, order)
+                    rows.append({"ids": seq, "markers": markers,
+                                 "qtype": it["qtype"], "n_opts": it["n_opts"],
+                                 "spec": it["spec"], "q": it["q"]})
+                    row_meta.append((qi, order))
+            b = self._collate(rows)
+            logits, act = self._forward(b)
+            decoded = self._decode_rows(logits.numpy(), act.numpy(), rows)
+
+            # regroup per question; average choice probabilities over orders
+            n_q = len(questions)
+            per_q = [[] for _ in range(n_q)]
+            for (qi, order), d in zip(row_meta, decoded):
+                per_q[qi].append((order,) + d)
+            answers, confs, actions = {}, {}, {}
+            scores = {}
+            for qi, rowsq in enumerate(per_q):
+                q = questions[qi]
+                order0 = rowsq[0][0]
+                n_opts = len(order0)
+                if len(rowsq) == 1:
+                    _, ans, conf, actp, p = rowsq[0]
+                else:
+                    pal = [0.0] * n_opts
+                    for order, _ans, _conf, _actp, p in rowsq:
+                        for pos, opt_idx in enumerate(order):
+                            pal[opt_idx] += p[pos]
+                    m = len(rowsq)
+                    pal = [v / m for v in pal]
+                    p = pal
+                    conf = round(float(min(1.0, max(pal[:n_opts]))), 4)
+                    actp = rowsq[0][3]
+                    ans = max(range(n_opts), key=lambda i: pal[i])
+                if q["type"] == "choice":
+                    pass  # ans is already the option index
+                elif q["type"] == "score":
+                    ans = 2 + max(range(len(p)), key=lambda i: p[i])
+                else:
+                    thr = q.get("threshold")
+                    thr = 0.5 if thr is None else float(thr)
+                    ans = bool(p[1] >= thr)
+                answers[q["id"]] = ans
+                confs[q["id"]] = conf
+                actions[q["id"]] = {"act_probability": actp}
+                scores[q["id"]] = [round(float(v), 4) for v in p]
+            total_tokens = int(b["attention_mask"].sum().item())
+            usage = {"input_tokens": total_tokens, "output_tokens": 0}
+            if with_scores:
+                return answers, confs, actions, usage, scores
+            return answers, confs, actions, usage

phocinae/main.py

@@ -0,0 +1,20 @@
+"""uvicorn entrypoint: python -m phocinae.main"""
+
+import logging
+import os
+
+import uvicorn
+
+
+def main():
+    logging.basicConfig(
+        level=os.environ.get("PHOC_LOG", "INFO").upper(),
+        format="%(asctime)s %(levelname)s %(name)s: %(message)s")
+    host = os.environ.get("PHOC_HOST", "127.0.0.1")
+    port = int(os.environ.get("PHOC_PORT", "8155"))
+    uvicorn.run("phocinae.server:app", host=host, port=port, workers=1,
+                log_level="warning", access_log=os.environ.get("PHOC_ACCESS_LOG", "0") == "1")
+
+
+if __name__ == "__main__":
+    main()

phocinae/model.py

@@ -0,0 +1,159 @@
+"""Pure-torch mmBERT-small encoder + typed-decisions head.
+
+Re-implements (inference only, no runtime dependency on laya/transformers):
+  - mmBERT-small encoder: 384 hidden / 22 layers / 6 heads / GEGLU MLP /
+    RoPE (Gemma2-style, theta=160000) / sliding-window attention on all but
+    every third layer (full attention) / pre-norm, LayerNorm(eps=1e-5, no bias),
+    layer 0 attention norm is Identity.
+  - decision head: type embedding + 2-layer pre-norm TransformerEncoder
+    (nn.TransformerEncoderLayer, d=384, 6 heads, ff=1536, dropout 0.1) +
+    marker-gathered scorer + act head (d+4 -> 256 -> 2).
+
+State-dict keys match the released checkpoint exactly
+(encoder.layers.i.attn.Wqkv/Wo, mlp.Wi/Wo, attn_norm/mlp_norm, head.layers.*,
+scorer.*, act_head.*, type_emb.weight).
+"""
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def _rotate_half(x):
+    d = x.shape[-1] // 2
+    return torch.cat([-x[..., d:], x[..., :d]], dim=-1)
+
+
+class _Attn(nn.Module):
+    """Fused Wqkv attention with RoPE + SDPA (checkpoint keys attn.Wqkv/attn.Wo)."""
+
+    def __init__(self, d, heads):
+        super().__init__()
+        self.heads = heads
+        self.hd = d // heads
+        self.Wqkv = nn.Linear(d, 3 * d, bias=False)
+        self.Wo = nn.Linear(d, d, bias=False)
+
+    def forward(self, x, mask, cos, sin):
+        B, L, _ = x.shape
+        qkv = self.Wqkv(x).view(B, L, 3, self.heads, self.hd)
+        q, k, v = qkv.unbind(2)
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+        q = (q.float() * cos + _rotate_half(q.float()) * sin).to(q.dtype)
+        k = (k.float() * cos + _rotate_half(k.float()) * sin).to(k.dtype)
+        o = F.scaled_dot_product_attention(
+            q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False,
+            scale=1.0 / math.sqrt(self.hd))
+        return self.Wo(o.transpose(1, 2).reshape(B, L, -1))
+
+
+class _MLP(nn.Module):
+    """GEGLU MLP (checkpoint keys mlp.Wi/mlp.Wo): Wo(gelu(input) * gate)."""
+
+    def __init__(self, d, inter):
+        super().__init__()
+        self.Wi = nn.Linear(d, 2 * inter, bias=False)
+        self.Wo = nn.Linear(inter, d, bias=False)
+
+    def forward(self, x):
+        inp, gate = self.Wi(x).chunk(2, dim=-1)
+        return self.Wo(F.gelu(inp) * gate)
+
+
+class _EncoderLayer(nn.Module):
+    def __init__(self, d, heads, inter, eps, full, first):
+        super().__init__()
+        self.full = full
+        self.attn_norm = nn.Identity() if first else nn.LayerNorm(d, eps=eps, bias=False)
+        self.attn = _Attn(d, heads)
+        self.mlp_norm = nn.LayerNorm(d, eps=eps, bias=False)
+        self.mlp = _MLP(d, inter)
+
+    def forward(self, x, mask, cos, sin):
+        x = x + self.attn(self.attn_norm(x), mask, cos, sin)
+        x = x + self.mlp(self.mlp_norm(x))
+        return x
+
+
+class MMBertEncoder(nn.Module):
+    def __init__(self, vocab, d, n_layers, heads, inter, eps, full_layers,
+                 theta, local_window):
+        super().__init__()
+        self.hd = d // heads
+        self.local_window = local_window
+        self.embeddings = nn.Module()
+        self.embeddings.tok_embeddings = nn.Embedding(vocab, d, padding_idx=0)
+        self.embeddings.norm = nn.LayerNorm(d, eps=eps, bias=False)
+        self.layers = nn.ModuleList([
+            _EncoderLayer(d, heads, inter, eps, full=(i in full_layers), first=(i == 0))
+            for i in range(n_layers)
+        ])
+        self.final_norm = nn.LayerNorm(d, eps=eps, bias=False)
+        inv = 1.0 / (theta ** (torch.arange(0, self.hd, 2, dtype=torch.float32) / self.hd))
+        self.register_buffer("inv_freq", inv, persistent=False)
+
+    def forward(self, input_ids, attention_mask):
+        B, L = input_ids.shape
+        dev = input_ids.device
+        x = self.embeddings.tok_embeddings(input_ids)
+        x = self.embeddings.norm(x)
+        pos = torch.arange(L, device=dev, dtype=torch.float32)
+        freqs = torch.outer(pos, self.inv_freq.to(dev))  # [L, hd/2] fp32
+        emb = torch.cat([freqs, freqs], dim=-1)          # [L, hd]
+        cos = emb.cos()[None, None]                      # [1, 1, L, hd]
+        sin = emb.sin()[None, None]
+        pad = attention_mask.bool()                      # [B, L]
+        full_mask = pad[:, None, None, :]                # [B, 1, 1, L]
+        band = (pos[:, None] - pos[None, :]).abs() <= self.local_window
+        slide_mask = band[None, None] & pad[:, None, None, :]  # [B, 1, L, L]
+        for layer in self.layers:
+            x = layer(x, full_mask if layer.full else slide_mask, cos, sin)
+        return self.final_norm(x)
+
+
+class DecisionModel(nn.Module):
+    """Full decision model; checkpoint keys live at the top level:
+
+    encoder.*, head.layers.*, scorer.*, act_head.*, type_emb.weight
+    """
+
+    def __init__(self, d, head_layers, heads, dropout):
+        super().__init__()
+        self.type_emb = nn.Embedding(3, d)
+        layer = nn.TransformerEncoderLayer(
+            d, heads, 4 * d, dropout, batch_first=True, norm_first=True)
+        self.head = nn.TransformerEncoder(
+            layer, head_layers, enable_nested_tensor=False)
+        self.scorer = nn.Sequential(
+            nn.LayerNorm(d), nn.Linear(d, d), nn.GELU(), nn.Linear(d, 1))
+        self.act_head = nn.Sequential(
+            nn.Linear(d + 4, 256), nn.GELU(), nn.Linear(256, 2))
+
+    def forward(self, input_ids, attention_mask, marker_pos, marker_mask, qtype):
+        h = self.encoder(input_ids, attention_mask)
+        h = h + self.type_emb(qtype)[:, None, :]
+        pad = ~attention_mask.bool()
+        h = self.head(h, src_key_padding_mask=pad)
+        return self._score(h, marker_pos, marker_mask)
+
+    def _score(self, h, marker_pos, marker_mask):
+        idx = marker_pos.clamp(min=0)[:, :, None].expand(-1, -1, h.size(-1))
+        m = torch.gather(h, 1, idx)
+        logits = self.scorer(m).squeeze(-1).float()
+        logits = logits.masked_fill(~marker_mask, -1e4)
+        p = torch.softmax(logits.detach(), -1)
+        k = marker_mask.sum(-1).clamp(min=2).float()
+        ent = -(p * torch.log(p.clamp_min(1e-9))).sum(-1) / torch.log(k)
+        if p.size(-1) >= 2:
+            top2 = p.topk(2, -1).values
+        else:
+            top1 = p.topk(1, -1).values
+            top2 = torch.cat([top1, torch.zeros_like(top1)], dim=-1)
+        feats = torch.stack(
+            [top2[:, 0], top2[:, 0] - top2[:, 1], ent, k / 255.0], -1)
+        act_logits = self.act_head(torch.cat([h[:, 0].float(), feats], -1))
+        return logits, act_logits