tinyjev 0.1.0__py3-none-any.whl

tinyjev/convert.py

@@ -0,0 +1,243 @@
+"""Convert upstream checkpoints into the tinyjev layout (v2, transformers-compatible).
+
+    <dest>/config.json           standard Qwen3Model config (AutoModel.from_pretrained loads the backbone)
+    <dest>/model.safetensors     backbone with standard Qwen3Model keys (fp16 by default)
+    <dest>/head.safetensors      the decision head (fp32), keys without prefix
+    <dest>/tinyjev.json          family, head config, tokenizer ids, upstream provenance
+    <dest>/tokenizer.json ...    tokenizer files at the root, repaired for current loaders
+
+Both backends load exactly this. Conversion streams tensor by tensor so a 2.4 GB fp32
+checkpoint converts on a 16 GB machine.
+"""
+from __future__ import annotations
+
+import json
+import shutil
+from pathlib import Path
+from typing import Dict
+
+import numpy as np
+
+NANOJEV_HEAD_PREFIXES = ("norm.", "scalar.", "set_project.", "set_attention.", "set_output.")
+
+
+def _save(dest: Path, weights: Dict[str, np.ndarray]):
+    """Split into model.safetensors (backbone, standard keys) and head.safetensors (head, unprefixed keys)."""
+    from safetensors.numpy import save_file
+    backbone = {k[len("backbone."):]: v for k, v in weights.items() if k.startswith("backbone.")}
+    head = {k[len("head."):]: v for k, v in weights.items() if k.startswith("head.")}
+    save_file(backbone, str(dest / "model.safetensors"), metadata={"format": "pt"})
+    save_file(head, str(dest / "head.safetensors"), metadata={"format": "pt"})
+
+
+def _write_backbone_config(dest: Path, backbone_cfg: dict, dtype: str):
+    """A config.json transformers 4.x and 5.x both read correctly: rope base at the top level AND under
+    rope_parameters (4.x ignores the latter and would silently default to 10000 — the NanoJev bug)."""
+    cfg = dict(backbone_cfg)
+    rope = cfg.get("rope_parameters") or {}
+    theta = rope.get("rope_theta", cfg.get("rope_theta"))
+    if theta is not None:
+        cfg["rope_theta"] = theta
+        cfg["rope_parameters"] = {**rope, "rope_theta": theta, "rope_type": rope.get("rope_type", "default")}
+    cfg["architectures"] = ["Qwen3Model"]
+    cfg["model_type"] = cfg.get("model_type", "qwen3")
+    cfg["dtype"] = dtype
+    cfg["torch_dtype"] = dtype
+    cfg.pop("transformers_version", None)
+    (dest / "config.json").write_text(json.dumps(cfg, indent=2))
+
+
+def _cast(arr: np.ndarray, dtype: str) -> np.ndarray:
+    return arr.astype(np.float16 if dtype == "float16" else np.float32)
+
+
+def normalize_tokenizer_config(cfg: dict):
+    """Repair fields newer loaders refuse: NanoJev stores extra_special_tokens as a list."""
+    changed = False
+    extra = cfg.get("extra_special_tokens")
+    if isinstance(extra, list):
+        cfg["extra_special_tokens"] = {f"extra_{i}": t for i, t in enumerate(extra)}
+        changed = True
+    if cfg.get("tokenizer_class") in (None, "TokenizersBackend"):
+        cfg["tokenizer_class"] = "PreTrainedTokenizerFast"
+        cfg.pop("backend", None)
+        cfg.pop("is_local", None)
+        changed = True
+    return cfg, changed
+
+
+def _copy_tokenizer(src_dir: Path, dest: Path):
+    tok = dest
+    tok.mkdir(parents=True, exist_ok=True)
+    for name in ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json",
+                 "added_tokens.json", "vocab.json", "merges.txt", "chat_template.jinja"):
+        if (src_dir / name).exists():
+            shutil.copy2(src_dir / name, tok / name)
+    cfg_path = tok / "tokenizer_config.json"
+    if cfg_path.exists():
+        cfg, changed = normalize_tokenizer_config(json.loads(cfg_path.read_text()))
+        if changed:
+            cfg_path.write_text(json.dumps(cfg, indent=2))
+    return tok
+
+
+def _eos_pad_from_tokenizer(tok_dir: Path, backbone_cfg: dict):
+    from tokenizers import Tokenizer
+    t = Tokenizer.from_file(str(tok_dir / "tokenizer.json"))
+    eos = backbone_cfg.get("eos_token_id")
+    cfg_path = tok_dir / "tokenizer_config.json"
+    if cfg_path.exists():
+        cfg = json.loads(cfg_path.read_text())
+        for key in ("eos_token", "pad_token"):
+            v = cfg.get(key)
+            if isinstance(v, dict):
+                cfg[key] = v.get("content")
+        if isinstance(cfg.get("eos_token"), str) and t.token_to_id(cfg["eos_token"]) is not None:
+            eos = t.token_to_id(cfg["eos_token"])
+        pad = t.token_to_id(cfg["pad_token"]) if isinstance(cfg.get("pad_token"), str) else None
+    else:
+        pad = None
+    return int(eos), int(pad if pad is not None else eos)
+
+
+# ---------------------------------------------------------------- NanoJev
+def convert_nanojev(source, dest, dtype: str = "float16") -> Path:
+    from safetensors import safe_open
+    src, dst = Path(source).expanduser().resolve(strict=True), Path(dest).expanduser()
+    for rel in ("config.json", "backbone_config/config.json", "tokenizer/tokenizer.json", "best.safetensors"):
+        if not (src / rel).exists():
+            raise FileNotFoundError(f"NanoJev checkpoint is missing {rel}")
+    dst.mkdir(parents=True, exist_ok=True)
+    run = json.loads((src / "config.json").read_text())
+    backbone_cfg = json.loads((src / "backbone_config/config.json").read_text())
+
+    weights, n_body, n_head = {}, 0, 0
+    with safe_open(str(src / "best.safetensors"), framework="numpy") as f:
+        for key in f.keys():
+            arr = f.get_tensor(key)
+            if key.startswith(NANOJEV_HEAD_PREFIXES):
+                weights["head." + key] = _cast(arr, "float32"); n_head += 1
+            elif key.startswith("backbone."):
+                weights[key] = _cast(arr, dtype); n_body += 1
+            else:
+                raise ValueError(f"unexpected tensor {key}")
+    _save(dst, weights); del weights
+    _write_backbone_config(dst, backbone_cfg, dtype)
+    tok = _copy_tokenizer(src / "tokenizer", dst)
+    eos, pad = _eos_pad_from_tokenizer(tok, backbone_cfg)
+    (dst / "tinyjev.json").write_text(json.dumps({
+        "format": "tinyjev-v2", "family": "marker", "name": "nanojev",
+        "head": {"set_head": run.get("set_head", "attention")},
+        "tokenizer": {"eos_token_id": eos, "pad_token_id": eos},
+        "max_length": run.get("max_length", 8192),
+        "dtypes": {"backbone": dtype, "head": "float32"},
+        "upstream": {"repo": "C-Tianyu/NanoJev", "base_model": run.get("model"),
+                     "base_revision": run.get("resolved_model_revision"),
+                     "schema": run.get("schema_version")},
+    }, indent=2))
+    print(f"nanojev: {n_body} backbone tensors -> {dtype}, {n_head} head tensors -> float32\n-> {dst}")
+    return dst
+
+
+# ---------------------------------------------------------------- Kev
+def convert_kev(adapter_dir, base_dir=None, dest=None, dtype: str = "float16", name: str = "kev") -> Path:
+    """Convert a Kev run: either a LoRA run (adapter + head.pt, merged into `base_dir`) or a
+    full fine-tune run from `kev.train --lora 0` (the run dir holds the whole backbone via
+    save_pretrained, plus head.pt; `base_dir` is then ignored)."""
+    from safetensors import safe_open
+    try:
+        import torch
+    except ImportError as exc:
+        raise ImportError("converting a Kev checkpoint needs torch: pip install 'tinyjev[convert]'") from exc
+
+    adapter = Path(adapter_dir).expanduser().resolve(strict=True)
+    dst = Path(dest).expanduser()
+    dst.mkdir(parents=True, exist_ok=True)
+    full_ft = not (adapter / "adapter_model.safetensors").exists()
+    if full_ft:
+        base = adapter                      # the run dir IS the backbone
+        scale, acfg = 0.0, {"r": 0, "lora_alpha": 0}
+    else:
+        if base_dir is None:
+            raise ValueError("a LoRA run needs base_dir (the Qwen3 base the adapter was trained on)")
+        base = Path(base_dir).expanduser().resolve(strict=True)
+        acfg = json.loads((adapter / "adapter_config.json").read_text())
+        if acfg.get("trainable_token_indices"):
+            raise ValueError("adapters with trained token embeddings are not supported by this converter")
+        scale = float(acfg["lora_alpha"]) / float(acfg["r"])
+        if acfg.get("use_rslora"):
+            scale = float(acfg["lora_alpha"]) / float(acfg["r"]) ** 0.5
+    backbone_cfg = json.loads((base / "config.json").read_text())
+
+    # LoRA deltas keyed by the base tensor they modify (empty for a full fine-tune)
+    deltas: Dict[str, np.ndarray] = {}
+    adapter_file = adapter / "adapter_model.safetensors"
+    with (safe_open(str(adapter_file), framework="numpy") if not full_ft else _NoTensors()) as f:
+        keys = list(f.keys())
+        a_keys = [k for k in keys if k.endswith("lora_A.weight")]
+        for ka in a_keys:
+            kb = ka.replace("lora_A.weight", "lora_B.weight")
+            A, B = f.get_tensor(ka).astype(np.float32), f.get_tensor(kb).astype(np.float32)
+            # peft key: base_model.model.model.layers.N.self_attn.q_proj.lora_A.weight
+            target = ka.split("base_model.model.")[-1].replace(".lora_A.weight", ".weight")
+            if target.startswith("model."):
+                target = target[len("model."):]
+            deltas[target] = (B @ A) * scale
+    merged_count = 0
+    weights: Dict[str, np.ndarray] = {}
+    base_files = sorted(p for p in base.glob("*.safetensors") if p.name != "adapter_model.safetensors")
+    if not base_files:
+        raise FileNotFoundError(f"no backbone safetensors found in {base}")
+    for bf in base_files:
+        with safe_open(str(bf), framework="pt") as f:      # base is bf16; numpy cannot read it
+            for key in f.keys():
+                if key.startswith("lm_head."):
+                    continue                                # never generates text
+                t = f.get_tensor(key).to(torch.float32).numpy()
+                short = key[len("model."):] if key.startswith("model.") else key
+                if short in deltas:
+                    t = t + deltas.pop(short); merged_count += 1
+                weights["backbone." + short] = _cast(t, dtype)
+    if deltas:
+        raise ValueError(f"{len(deltas)} adapter tensors matched no base weight, e.g. {list(deltas)[:3]}")
+
+    meta = torch.load(str(adapter / "head.pt"), map_location="cpu")
+    head = meta["head"]
+    for k in ("q.weight", "q.bias", "k.weight", "k.bias"):
+        weights["head." + k] = head[k].to(torch.float32).numpy()
+    _save(dst, weights); del weights
+    _write_backbone_config(dst, backbone_cfg, dtype)
+    tok = _copy_tokenizer(adapter, dst)
+    eos, pad = _eos_pad_from_tokenizer(tok, backbone_cfg)
+    (dst / "tinyjev.json").write_text(json.dumps({
+        "format": "tinyjev-v2", "family": "pointer", "name": name,
+        "head": {"head_dim": int(meta.get("head_dim", 256)), "temperature": float(meta.get("temperature", 1.0)),
+                 "option_isolation": bool(meta.get("option_isolation", False))},
+        "tokenizer": {"eos_token_id": eos, "pad_token_id": pad},
+        "max_state": 8192, "max_branch": 8192,
+        "dtypes": {"backbone": dtype, "head": "float32"},
+        "upstream": {"repo": str(adapter_dir), "base_model": meta.get("base"),
+                     "base_revision": meta.get("base_revision"), "lora_rank": int(acfg["r"]),
+                     "lora_alpha": acfg["lora_alpha"], "merged_tensors": merged_count,
+                     "full_finetune": full_ft},
+    }, indent=2))
+    print(f"kev: merged {merged_count} LoRA deltas into the base, backbone -> {dtype}, head -> float32\n-> {dst}")
+    return dst
+
+
+class _NoTensors:
+    """Stand-in for safe_open when a run has no adapter file."""
+    def __enter__(self):
+        return self
+    def __exit__(self, *exc):
+        return False
+    def keys(self):
+        return []
+
+
+def convert(family: str, dest, dtype: str = "float16", **kw) -> Path:
+    if family == "nanojev":
+        return convert_nanojev(kw["source"], dest, dtype)
+    if family == "kev":
+        return convert_kev(kw["adapter"], kw.get("base"), dest, dtype, name=kw.get("name", "kev"))
+    raise ValueError(f"unknown family {family!r}")

tinyjev/families/__init__.py

@@ -0,0 +1,55 @@
+"""Model families: a prompt format plus the decision head that reads it.
+
+    pointer  one branch per question ending in a decide token, scored against each option's end token
+    marker   one sequence per candidate, pooled at its end token, with attention across the candidate set
+
+
+    enc = family.encode(record)                 # -> Encoded(prefix, rows, questions)
+    hs  = backbone.hidden_rows(enc.prefix, enc.rows, family.pad_token_id)
+    out = family.logits(hs, enc)                # -> {question id: np.ndarray logits}
+
+Heads run in numpy (fp32); they are tiny and this keeps every backend numerically identical.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List
+
+
+@dataclass
+class Encoded:
+    prefix: List[int]                       # tokens shared by every row (the state)
+    rows: List[List[int]]                   # per-row suffix tokens, appended to the prefix
+    questions: List[Dict[str, Any]] = field(default_factory=list)   # family-specific readout info
+
+
+def make(name: str, root, manifest: dict):
+    aliases = {"nanojev": "marker", "kev": "pointer"}      # layouts written before the rename
+    name = aliases.get(name, name)
+    if name == "marker":
+        from .marker import MarkerFamily as F
+    elif name == "pointer":
+        from .pointer import PointerFamily as F
+    else:
+        raise ValueError(f"unknown family {name!r}; expected 'pointer' or 'marker'")
+    return F(root, manifest)
+
+
+def load_head(root) -> Dict[str, Any]:
+    """The fp32 decision head from head.safetensors, as numpy; the backbone is never touched."""
+    import numpy as np
+    from pathlib import Path
+    from safetensors import safe_open
+    out = {}
+    with safe_open(str(Path(root) / "head.safetensors"), framework="numpy") as f:
+        for key in f.keys():
+            out[key] = np.asarray(f.get_tensor(key), dtype=np.float32)
+    return out
+
+
+def softmax(z):
+    import numpy as np
+    z = np.asarray(z, dtype=np.float64)
+    z = z - z.max()
+    e = np.exp(z)
+    return (e / e.sum()).astype(np.float64)

tinyjev/families/marker.py

@@ -0,0 +1,152 @@
+"""NanoJev: one row per candidate, EOS pooling, scalar scorer + set-attention over choice options.
+
+Prompt bytes and the head follow upstream `predict_toy_decisions` / `train_toy_decisions`
+(MIT, OpenJev contributors). Parity is checked against upstream's own CUDA predictions.
+"""
+from __future__ import annotations
+
+import json
+import math
+from typing import Any, Dict, List
+
+import numpy as np
+
+from . import Encoded, load_head, softmax
+
+QUESTION_TYPES = {"boolean", "choice", "score"}
+SET_DIM, SET_HEADS = 128, 4
+
+
+def _nonempty(v) -> bool:
+    return isinstance(v, str) and bool(v.strip())
+
+
+def validate_question(where: str, q: dict):
+    if not isinstance(q, dict) or set(q) - {"id", "type", "instructions", "criteria"}:
+        raise ValueError(f"{where}: question may only hold type, instructions, criteria")
+    typ = q.get("type")
+    if typ not in QUESTION_TYPES or not _nonempty(q.get("instructions")):
+        raise ValueError(f"{where}: invalid type or instructions")
+    c = q.get("criteria")
+    if typ == "boolean":
+        if "criteria" in q and (not isinstance(c, dict) or set(c) - {"false", "true"}
+                                or not all(_nonempty(v) for v in c.values())):
+            raise ValueError(f"{where}: boolean criteria may only hold non-empty false/true")
+    elif typ == "choice":
+        if not isinstance(c, dict) or not 2 <= len(c) <= 255 or not all(
+                _nonempty(k) and _nonempty(v) for k, v in c.items()):
+            raise ValueError(f"{where}: choice criteria must be 2-255 non-empty id: description pairs")
+    else:
+        if not isinstance(c, list) or not 2 <= len(c) <= 10 or not all(_nonempty(v) for v in c):
+            raise ValueError(f"{where}: score criteria must be an ordered list of 2-10 non-empty levels")
+
+
+def candidates(q: dict):
+    typ = q["type"]
+    if typ == "boolean":
+        return ["false", "true"], ["The proposition is true."]
+    if typ == "choice":
+        ids = list(q["criteria"])
+        return ids, [f"{k}: {q['criteria'][k]}" for k in ids]
+    return [str(i) for i in range(len(q["criteria"]))], list(q["criteria"])
+
+
+def layer_norm(x, w, b, eps=1e-5):
+    mu = x.mean(-1, keepdims=True)
+    var = ((x - mu) ** 2).mean(-1, keepdims=True)
+    return (x - mu) / np.sqrt(var + eps) * w + b
+
+
+class MarkerFamily:
+    name = "marker"
+
+    def __init__(self, root, manifest: dict):
+        from tokenizers import Tokenizer
+        self.tok = Tokenizer.from_file(str(root / "tokenizer.json"))
+        tk = manifest["tokenizer"]
+        self.eos_token_id = int(tk["eos_token_id"])
+        self.pad_token_id = int(tk.get("pad_token_id", tk["eos_token_id"]))
+        self.max_length = int(manifest.get("max_length", 8192))
+        self.set_head = manifest["head"].get("set_head", "attention")
+        self.w = load_head(root)
+
+    def _enc(self, text: str) -> List[int]:
+        return self.tok.encode(text, add_special_tokens=False).ids
+
+    # ---- prompt ----
+    def encode(self, record: dict) -> Encoded:
+        state = record["state"]
+        if not isinstance(state, (str, dict, list)) or not state:
+            raise ValueError(f"{record['id']}: state must be a non-empty string, object or array")
+        prefix = self._enc(f"State:\n{state}\n")          # str() of dict/list, as upstream trained
+        rows, questions = [], []
+        for q in record["questions"]:
+            where = f"{record['id']}:{q['id']}"
+            validate_question(where, q)
+            typ = q["type"]
+            ids, texts = candidates(q)
+            head = f"Question type: {typ}\nQuestion:\n{q['instructions']}\n"
+            if typ == "boolean" and "criteria" in q:
+                for key, label in (("false", "False"), ("true", "True")):
+                    if key in q["criteria"]:
+                        head += f"{label} criterion: {q['criteria'][key]}\n"
+            head_ids = self._enc(head)
+            start = len(rows)
+            for t in texts:
+                rows.append(head_ids + self._enc(f"Candidate:\n{t}\nDecision:") + [self.eos_token_id])
+            longest = len(prefix) + max(len(r) for r in rows[start:])
+            if longest > self.max_length:
+                raise ValueError(f"{where}: candidate path is {longest} tokens, over max_length={self.max_length}")
+            questions.append({"id": q["id"], "type": typ, "keys": ids,
+                              "rows": list(range(start, len(rows)))})
+        return Encoded(prefix=prefix, rows=rows, questions=questions)
+
+    # ---- head ----
+    def _set_attention(self, u: np.ndarray) -> np.ndarray:
+        w = self.w
+        k_, d = u.shape
+        qkv = u @ w["set_attention.in_proj_weight"].T + w["set_attention.in_proj_bias"]
+        q, k, v = np.split(qkv, 3, axis=-1)
+        hd = d // SET_HEADS
+        heads = lambda t: t.reshape(k_, SET_HEADS, hd).transpose(1, 0, 2)   # [H,K,hd]
+        q, k, v = heads(q), heads(k), heads(v)
+        s = (q @ k.transpose(0, 2, 1)) / math.sqrt(hd)
+        s = s - s.max(-1, keepdims=True)
+        a = np.exp(s); a /= a.sum(-1, keepdims=True)
+        mixed = (a @ v).transpose(1, 0, 2).reshape(k_, d)
+        return mixed @ w["set_attention.out_proj.weight"].T + w["set_attention.out_proj.bias"]
+
+    def logits(self, hidden_rows: List[np.ndarray], enc: Encoded) -> Dict[str, np.ndarray]:
+        w, out = self.w, {}
+        for q in enc.questions:
+            leaves = np.stack([hidden_rows[i][-1] for i in q["rows"]]).astype(np.float32)
+            h = layer_norm(leaves, w["norm.weight"], w["norm.bias"])
+            z = (h @ w["scalar.weight"].T + w["scalar.bias"]).reshape(-1)
+            if q["type"] == "choice" and self.set_head == "attention":
+                log_k = np.full((h.shape[0], 1), math.log(h.shape[0]), dtype=np.float32)
+                u = np.concatenate([h, log_k], -1) @ w["set_project.weight"].T + w["set_project.bias"]
+                delta = (np.tanh(u + self._set_attention(u)) @ w["set_output.weight"].T
+                         + w["set_output.bias"]).reshape(-1)
+                z = z + delta
+            if q["type"] == "boolean":
+                z = np.array([0.0, float(z[0])], dtype=np.float32)
+            out[q["id"]] = z.astype(np.float32)
+        return out
+
+    # ---- answers ----
+    @staticmethod
+    def answer(q: dict, probs: List[float]) -> dict:
+        ids = q["keys"]
+        if len(probs) != len(ids) or not all(math.isfinite(p) and 0 <= p <= 1 for p in probs) \
+                or abs(math.fsum(probs) - 1.0) > 1e-5:
+            raise ValueError("model produced invalid probabilities")
+        best = max(range(len(ids)), key=probs.__getitem__)
+        a = {"type": q["type"], "probabilities": dict(zip(ids, probs))}
+        if q["type"] == "boolean":
+            a.update(p_true=probs[1], value=bool(best))
+        elif q["type"] == "choice":
+            a.update(choice=ids[best], value=ids[best])
+        else:
+            s = math.fsum(i * p for i, p in enumerate(probs))
+            a.update(score=s, level=best, value=s)
+        return a

tinyjev/families/pointer.py

@@ -0,0 +1,131 @@
+"""Kev: one row per question, delimiter tokens, pointer head (<decide> against each </opt>).
+
+Prompt layout, delimiter choice and the head follow `kev.model` / `kev.api` (Apache-2.0,
+Jared Palmer). Each question runs as its own causal row continuing from the state, which
+kev.model documents as exactly equivalent to its packed block-causal form.
+"""
+from __future__ import annotations
+
+import math
+import re
+from typing import Any, Dict, List
+
+import numpy as np
+
+from . import Encoded, load_head
+
+SPECIAL = ["<|fim_prefix|>", "<|fim_middle|>", "<|box_start|>", "<|box_end|>", "<|fim_suffix|>"]
+_SPECIAL_RE = re.compile(r"<\|([A-Za-z0-9_]+)\|>")
+MAX_OPTIONS = 255
+
+
+def render(v, indent: int = 0) -> str:
+    pad = "  " * indent
+    if v is None:
+        return ""
+    if isinstance(v, (str, int, float, bool)):
+        return str(v)
+    if isinstance(v, list):
+        return "\n".join(f"{pad}- {render(x, indent + 1).lstrip()}" for x in v)
+    return "\n".join(f"{pad}{k}:\n{render(x, indent + 1)}" if isinstance(x, (dict, list))
+                     else f"{pad}{k}: {render(x)}" for k, x in v.items())
+
+
+def option_text(name: str, desc) -> str:
+    return name if desc is None or desc == "" else f"{name}: {render(desc)}"
+
+
+class PointerFamily:
+    name = "pointer"
+
+    def __init__(self, root, manifest: dict):
+        from tokenizers import Tokenizer
+        self.tok = Tokenizer.from_file(str(root / "tokenizer.json"))
+        ids = [self.tok.token_to_id(t) for t in SPECIAL]
+        if any(i is None for i in ids):
+            raise ValueError("tokenizer lacks the Qwen delimiter tokens Kev relies on")
+        self.state_id, self.q_id, self.o_id, self.c_id, self.d_id = ids
+        tk = manifest["tokenizer"]
+        self.pad_token_id = int(tk.get("pad_token_id", 0))
+        self.max_state = int(manifest.get("max_state", 8192))
+        self.max_branch = int(manifest.get("max_branch", 8192))
+        self.temperature = float(manifest["head"].get("temperature", 1.0))
+        self.dp = int(manifest["head"].get("head_dim", 256))
+        self.w = load_head(root)
+
+    def user_tokens(self, text: str) -> List[int]:
+        # caller text can never forge a delimiter: <|name|> becomes <¦name¦> before tokenizing
+        return self.tok.encode(_SPECIAL_RE.sub(r"<¦\1¦>", text), add_special_tokens=False).ids
+
+    # ---- prompt ----
+    def encode(self, record: dict) -> Encoded:
+        state_tokens = self.user_tokens(render(record["state"]))
+        if len(state_tokens) + 1 > self.max_state:
+            raise ValueError(f"{record['id']}: state exceeds {self.max_state} tokens")
+        prefix = [self.state_id] + state_tokens
+        rows, questions = [], []
+        for q in record["questions"]:
+            where = f"{record['id']}:{q['id']}"
+            typ = q["type"]
+            c = q.get("criteria")
+            if typ == "boolean":
+                c = c or {}
+                if not isinstance(c, dict) or set(c) - {"false", "true"}:
+                    raise ValueError(f"{where}: boolean criteria may only hold false/true")
+                keys, opts = ["false", "true"], [option_text("no", c.get("false")), option_text("yes", c.get("true"))]
+            elif typ == "choice":
+                if not isinstance(c, dict) or not 1 <= len(c) <= MAX_OPTIONS:
+                    raise ValueError(f"{where}: choice criteria must hold 1..{MAX_OPTIONS} options")
+                keys, opts = list(c), [option_text(k, v) for k, v in c.items()]
+            elif typ == "score":
+                if not isinstance(c, list) or not 2 <= len(c) <= MAX_OPTIONS:
+                    raise ValueError(f"{where}: score criteria must be a list of 2..{MAX_OPTIONS} levels")
+                keys, opts = [str(i) for i in range(len(c))], [render(x) for x in c]
+            else:
+                raise ValueError(f"{where}: unsupported question type {typ!r}")
+            instr = [self.q_id] + self.user_tokens(render(q["instructions"]))
+            spans = [[self.o_id] + self.user_tokens(o) + [self.c_id] for o in opts]
+            branch = instr + [t for sp in spans for t in sp] + [self.d_id]
+            if len(branch) > self.max_branch - len(prefix):
+                raise ValueError(f"{where}: branch too long ({len(branch)} tokens)")
+            ends, cursor = [], len(instr)
+            for sp in spans:
+                cursor += len(sp)
+                ends.append(cursor - 1)
+            rows.append(branch)
+            questions.append({"id": q["id"], "type": typ, "keys": keys, "row": len(rows) - 1,
+                              "decide": len(prefix) + len(branch) - 1,
+                              "opts": [len(prefix) + e for e in ends],
+                              "legend": dict(zip(keys, opts)) if typ == "score" else None})
+        return Encoded(prefix=prefix, rows=rows, questions=questions)
+
+    # ---- head ----
+    def logits(self, hidden_rows: List[np.ndarray], enc: Encoded) -> Dict[str, np.ndarray]:
+        w, out = self.w, {}
+        for q in enc.questions:
+            h = hidden_rows[q["row"]]
+            qv = h[q["decide"]] @ w["q.weight"].T + w["q.bias"]
+            kv = h[q["opts"]] @ w["k.weight"].T + w["k.bias"]
+            z = (kv @ qv) / math.sqrt(self.dp)
+            out[q["id"]] = (z / self.temperature).astype(np.float32)
+        return out
+
+    # ---- answers (System One shape, as kev.api.to_answers) ----
+    @staticmethod
+    def answer(q: dict, probs: List[float]) -> dict:
+        r2 = lambda x: round(float(x), 2)
+        if q["type"] == "boolean":
+            return {"type": "noul", "noul": r2(probs[1]), "value": bool(probs[1] >= 0.5),
+                    "p_true": probs[1], "probabilities": dict(zip(q["keys"], probs))}
+        if q["type"] == "choice":
+            K = len(probs)
+            best = max(range(K), key=probs.__getitem__)
+            conf = 1.0 if K == 1 else (max(probs) - 1 / K) / (1 - 1 / K)
+            return {"type": "choice", "choice": q["keys"][best], "value": q["keys"][best],
+                    "confidence": r2(conf), "probabilities": dict(zip(q["keys"], probs))}
+        L = len(probs)
+        mode = max(range(L), key=probs.__getitem__)
+        score = sum(i * p for i, p in enumerate(probs))
+        conf = 1.0 - sum(p * abs(i - mode) for i, p in enumerate(probs)) / (L - 1)
+        return {"type": "score", "score": r2(score), "value": score, "level": mode, "legend": q["legend"],
+                "probabilities": {str(i): v for i, v in enumerate(probs)}, "confidence": r2(conf)}

tinyjev/registry.py

@@ -0,0 +1,13 @@
+"""Known checkpoints. One entry: our model. Other projects' models are converted locally with
+`tinyjev convert` — we do not republish their weights."""
+MODELS = {
+    "tinyjev-0.6b": {"repo": "AnkitAI/tinyjev-0.6b", "family": "pointer", "params": "0.6B",
+                     "what": "Qwen3-0.6B-Base + pointer head, 596M"},
+}
+
+
+def resolve(name: str):
+    """alias -> (repo, None); anything else -> (name, None): a local path or a Hub repo id."""
+    if name in MODELS:
+        return MODELS[name]["repo"], None
+    return name, None