jev_toy/model.py

"""
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 encode_questions(self, q_ids, q_mask):
        """Encode ALL questions in ONE shared-encoder pass (true parallelism)."""
        return self.encoder(q_ids, q_mask)  # [B_q, D]

    def answer_state(self, h_state, h_q, q_indices):
        """
        Answer N questions given a SINGLE shared state vector and all their
        pre-encoded question vectors. This is the "one state, many questions"
        path: h_state [1,D], h_q [B_q,D], q_indices picks which question rows.
        """
        h = F.gelu(self.merge(torch.cat([h_state, h_q], dim=-1)))
        return {
            "noul": self.noul_head(h).squeeze(-1),
            "score": self.score_head(h).squeeze(-1),
            "choice": self.choice_head(h),
        }

    def n_params(self):
        return sum(p.numel() for p in self.parameters())