Upload jev_toy/model.py with huggingface_hub

jev_toy/model.py

@@ -0,0 +1,198 @@
+"""
+jev_toy/model.py
+
+A faithful-from-the-interface, toy-scale "System One" model.
+
+WHAT THIS IS (honest framing)
+-----------------------------
+The real Jev model (TypeSafe AI, 2026) is proprietary and unpublished. Its
+internals are NOT public. What IS public and repeatedly demonstrated:
+  - input  = a `state` (text/JSON) + a set of TYPED `questions`
+  - output = per-question typed, probabilistic decisions
+      * choice  -> probability distribution over named options
+      * score   -> a continuous score + the distribution underneath
+      * noul    -> a probability that a yes/no statement is true
+  - every question on one state is answered in PARALLEL (no token-by-token
+    autoregressive generation), and probabilities are claimed CALIBRATED.
+
+What we build here is OUR OWN design to reproduce that exact interface at toy
+scale. Every internal (encoder, heads, loss, calibration) is our choice, NOT a
+claim about Jev's internals.
+
+ARCHITECTURE (ours) --- built so "parallel" is real, not cosmetic
+-----------------------------------------------------------------
+Two-tower design driven by the proven property: one state, many questions.
+
+    state          question(s)
+      |                |
+      v                v
+  StateEncoder    QuestionEncoder (SHARED weights)
+      |                |
+   pooled h_s       pooled h_q
+      |                |
+      +----- concat --+
+      |        |
+      v        v
+   typed heads (noul / choice / score)
+
+The key fact we mirror: the state is encoded ONCE. Every question in the
+request consumes that same pooled state vector, each through its own head, in
+PARALLEL. Adding more questions to a request does not re-run the state encoder
+--- matching the documented behaviour ("evaluate every question in parallel").
+
+State and question are encoded with the SAME transformer weights (a shared
+encoder). We feed tokens once at sampling; instr-time the two streams are just
+concatenated batches through one transformer.
+
+Heads:
+  noul   : P(statement true)   = sigmoid(logit)
+  choice : softmax over K logits -> distribution + confidence
+  score  : logistic(logit)*range -> bounded real value in [lo, hi]
+
+Calibration is a separate post-training step (calibrate.py), NOT baked into the
+heads, because temperature scaling needs held-out data.
+
+Build knobs (d_model, n_layers, vocab) are small so this trains on a laptop CPU
+in minutes. Larger values + cuda device -> the same code scales to a GPU.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+@dataclass
+class SystemOneConfig:
+    vocab_size: int = 8192
+    d_model: int = 128
+    n_layers: int = 3
+    n_heads: int = 4
+    d_ff: int = 256
+    max_seq_len: int = 256
+    pad_token_id: int = 0
+    num_choice_heads: int = 8
+    score_range: tuple[float, float] = (1.0, 5.0)  # logistic maps logit -> this interval
+
+
+def _rope_cache(seq_len: int, dim: int, device: torch.device):
+    inv = 1.0 / (10000.0 ** (torch.arange(0, dim, 2).float() / dim))
+    pos = torch.arange(seq_len).float().unsqueeze(1)
+    return (pos * inv.unsqueeze(0)).to(device)
+
+
+def _apply_rope(x: torch.Tensor, cache: torch.Tensor):
+    B, T, H, D = x.shape
+    x = x.float().view(B, T, H, D // 2, 2)
+    x0, x1 = x.unbind(-1)
+    c = cache[:T].view(1, T, 1, D // 2)
+    cos, sin = c.cos(), c.sin()
+    out = torch.stack([x0 * cos - x1 * sin, x1 * cos + x0 * sin], -1).view(B, T, H, D)
+    return out.to(x.dtype)
+
+
+class _Attn(nn.Module):
+    def __init__(self, cfg: SystemOneConfig):
+        super().__init__()
+        self.h = cfg.n_heads
+        self.dh = cfg.d_model // cfg.n_heads
+        self.wq = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
+        self.wk = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
+        self.wv = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
+        self.wo = nn.Linear(cfg.d_model, cfg.d_model)
+        self.rope = None
+
+    def forward(self, x):
+        B, T, D = x.shape
+        if self.rope is None or self.rope.device != x.device:
+            self.rope = _rope_cache(x.shape[1], self.dh, x.device)
+        q = self.wq(x).view(B, T, self.h, self.dh).transpose(1, 2)
+        k = self.wk(x).view(B, T, self.h, self.dh).transpose(1, 2)
+        v = self.wv(x).view(B, T, self.h, self.dh).transpose(1, 2)
+        q = _apply_rope(q, self.rope)
+        k = _apply_rope(k, self.rope)
+        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.dh)
+        # No causal mask: a decision model attends to the whole context.
+        h = F.softmax(att, dim=-1) @ v
+        return self.wo(h.transpose(1, 2).reshape(B, T, D))
+
+
+class _Block(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(cfg.d_model)
+        self.attn = _Attn(cfg)
+        self.ln2 = nn.LayerNorm(cfg.d_model)
+        self.ff = nn.Sequential(
+            nn.Linear(cfg.d_model, cfg.d_ff), nn.GELU(), nn.Linear(cfg.d_ff, cfg.d_model)
+        )
+
+    def forward(self, x):
+        x = x + self.attn(self.ln1(x))
+        x = x + self.ff(self.ln2(x))
+        return x
+
+
+class _Encoder(nn.Module):
+    """Shared transformer encoder for state AND question text."""
+
+    def __init__(self, cfg: SystemOneConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model, padding_idx=cfg.pad_token_id)
+        self.blocks = nn.ModuleList([_Block(cfg) for _ in range(cfg.n_layers)])
+        self.ln = nn.LayerNorm(cfg.d_model)
+
+    def forward(self, ids, mask):
+        x = self.embed(ids) * mask.unsqueeze(-1).float()
+        # position comes from RoPE inside attention
+        for b in self.blocks:
+            x = b(x)
+        x = self.ln(x)
+        pooled = (x * mask.unsqueeze(-1).float()).sum(1) / mask.sum(1, keepdim=True).clamp_min(1)
+        return pooled
+
+
+class SystemOneModel(nn.Module):
+    """state -> one pooled vector; each question -> its own typed head."""
+
+    def __init__(self, cfg: SystemOneConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.encoder = _Encoder(cfg)
+        # Cross-tower fusion per head.
+        fusion_dim = cfg.d_model * 2
+        self.merge = nn.Linear(fusion_dim, cfg.d_model)
+        self.noul_head = nn.Linear(cfg.d_model, 1)
+        self.score_head = nn.Linear(cfg.d_model, 1)
+        self.choice_head = nn.Linear(cfg.d_model, cfg.num_choice_heads)
+
+    def encode_state(self, s_ids, s_mask):
+        """Single forward pass over the state -> [B, D]. Called ONCE per request."""
+        return self.encoder(s_ids, s_mask)
+
+    def answer(self, h_state, q_ids, q_mask, q_type):
+        """
+        Answer a batch of questions off a shared state vector.
+        h_state: [B_state, D] (one row per state)
+        q_ids/q_mask: question tokens [B_q, T]; B_q may be B_state * n_questions
+        q_type: list/str per row: 'noul' | 'choice' | 'score'
+        Returns a list of per-row decision dicts (heads applied per type).
+        """
+        h_q = self.encoder(q_ids, q_mask)  # [B_q, D]
+        # broadcast state rows to question rows
+        h_s = h_state  # expected already expanded to B_q
+        h = F.gelu(self.merge(torch.cat([h_s, h_q], dim=-1)))
+        logits = {
+            "noul": self.noul_head(h).squeeze(-1),
+            "score": self.score_head(h).squeeze(-1),
+            "choice": self.choice_head(h),
+        }
+        return logits, h
+
+    def n_params(self):
+        return sum(p.numel() for p in self.parameters())