Blame for src/mulattice/protocol.py

"""Shared evaluation/training protocol for the capacity, forgetting and real-data goals.

Single source of truth: the acceptance tests (tests/test_capacity.py,
tests/test_forgetting.py, tests/test_realtext.py) and the training scripts
(experiments/train_capacity.py, experiments/train_forgetting.py,
experiments/realdata_v2.py) all build sequences and measure accuracy through
this module, so "train until the test passes" cannot drift from what is tested.

Synthetic task (capacity + forgetting):
    vocab  keys [0, NK)  values [NK, NK+NV)  SEP = NK+NV+1   V = NK+NV+2
    sequence per row (filler f tokens before each pair; f=0 -> dense MQAR):
        [f k1 v1] [f k2 v2] ... [f kP vP]            store pass
        [f ku1 vu1'] ... [f kuM vuM']                update pass (overwrites)
        [q1 SEP] [q2 SEP] ... [qQ SEP]               query pass
    A query on an updated key must return the UPDATED value (last write wins).
    The answer is read at the SEP token of each query (causal).

Real-data task: wikitext-103 mined key->bigram-successor pairs embedded in a
real corpus filler stream (same mining as experiments/eval_realdata.py).
"""
from __future__ import annotations

import os
import time
from dataclasses import dataclass

import torch

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.join(HERE, "..")

# ---------------------------------------------------------------------------
# synthetic generator
# ---------------------------------------------------------------------------


@dataclass
class SynthCfg:
    NK: int = 16384          # distinct key tokens available
    NV: int = 16384          # distinct value tokens available
    d_model: int = 64
    n_heads: int = 2
    d_k: int = 32
    d_v: int = 32
    addr_scale: float = 64.0  # (2*64+3)^2 = 17,161 cells for 8,192 keys
    retention_init: float = 9.0

    @property
    def V(self) -> int:
        return self.NK + self.NV + 2

    @property
    def SEP(self) -> int:
        return self.NK + self.NV + 1


def synth_batch(
    B: int,
    n_pairs: int,
    n_queries: int,
    cfg: SynthCfg,
    filler: int = 0,
    overwrite_frac: float = 0.0,
    generator: torch.Generator | None = None,
    qpos: torch.Tensor | None = None,
    updated_idx: torch.Tensor | None = None,
    kw: dict | None = None,
):
    """Build a capacity/forgetting batch. See module docstring for the layout.

    Returns dict with:
      tokens      (B, L) long          the sequence
      ans_pos     (B, Q) long          SEP positions; answer read there
      gold        (B, Q) long          value token id that must be produced
      was_updated (B, Q) bool          query key was overwritten -> updated gold
      key_pos     (B, P) long          first-pass key token positions (spread loss)
      upd_key_pos (B, M) long | None   update-pass key token positions
      n_pairs, n_upd, L
    """
    g = generator
    NK, NV = cfg.NK, cfg.NV
    if n_pairs > NK:
        raise ValueError("n_pairs > key vocab")
    n_upd = int(round(n_pairs * overwrite_frac))
    if n_pairs + n_upd > NV:
        raise ValueError("not enough distinct values")
    dev = "cpu"  # build on CPU; caller moves to device

    # distinct keys per row; `include_keys` are guaranteed members of the stored
    # set (placed FIRST), and the rest are sampled from the complement, so the
    # batch is consistent by construction -- no duplicate stores of the forced
    # keys with conflicting values (the token-surgery approach produced exactly
    # that and wrecked the model: census 9 -> 713 errors, polish_weakkeys v1).
    include_keys = None
    if "include_keys" in (kw or {}):
        include_keys = kw["include_keys"]
        if include_keys.dim() == 1:
            include_keys = include_keys.unsqueeze(0).expand(B, -1)
    if include_keys is not None:
        nk_inc = include_keys.shape[1]
        assert n_pairs >= nk_inc
        rows = []
        for r in range(B):
            rest = torch.randperm(NK, generator=g)
            rest = rest[torch.isin(rest, include_keys[r], invert=True)][: n_pairs - nk_inc]
            rows.append(torch.cat([include_keys[r], rest]))
        keys = torch.stack(rows)
    else:
        keys = torch.stack([torch.randperm(NK, generator=g)[:n_pairs] for _ in range(B)])
    # distinct values: n_pairs originals + n_upd updates, all different in a row
    vals = torch.stack([torch.randperm(NV, generator=g)[: n_pairs + n_upd] for _ in range(B)])
    orig_vals, upd_vals = vals[:, :n_pairs], vals[:, n_pairs:]

    if updated_idx is None:
        # which stored keys get overwritten (row-specific)
        updated_idx = torch.stack([torch.randperm(n_pairs, generator=g)[:n_upd] for _ in range(B)])
    if qpos is None:
        # DISTINCT query keys (sample without replacement).  Sampling with
        # replacement creates duplicate queries at large Q (measured: 9 of 512
        # at Q=512/NK=16,384) -- not part of the published protocol semantics
        # ("every query has a unique answer", mqar.py); distinct sampling is
        # the clean generalization of the published small-Q generator.
        qpos = torch.stack([torch.randperm(n_pairs, generator=g)[:n_queries] for _ in range(B)])
    else:
        if qpos.dim() == 1:
            qpos = qpos.unsqueeze(0).expand(B, -1).contiguous()
        n_queries = qpos.shape[1]      # explicit query set defines Q
    qkeys = torch.gather(keys, 1, qpos)
    upd_mask = torch.zeros(B, n_pairs, dtype=torch.bool)
    upd_mask.scatter_(1, updated_idx, True)
    was_updated = torch.gather(upd_mask, 1, qpos)
    # latest value for every stored key: original, replaced where updated
    latest = orig_vals.clone()
    latest.scatter_(1, updated_idx, upd_vals)
    gold = torch.gather(latest, 1, qpos) + NK

    f = filler
    block = f + 2
    L = (n_pairs + n_upd) * block + 2 * n_queries

    tokens = torch.zeros(B, L, dtype=torch.long)
    # filler ids: uniform from the complement of the row's stored keys
    if f > 0:
        need = f * (n_pairs + n_upd)
        filler_ids = torch.zeros(B, need, dtype=torch.long)
        keyset = torch.zeros(B, NK, dtype=torch.bool)
        keyset.scatter_(1, keys, True)
        done = torch.zeros(B, need, dtype=torch.bool)
        while not bool(done.all()):
            cand = torch.randint(0, NK, (B, need), generator=g)
            ok = ~torch.gather(keyset, 1, cand) & ~done
            filler_ids = torch.where(ok, cand, filler_ids)
            done = done | ok
    # store pass
    for i in range(n_pairs):
        s = i * block
        if f:
            tokens[:, s:s + f] = filler_ids[:, i * f:(i + 1) * f]
        tokens[:, s + f] = keys[:, i]
        tokens[:, s + f + 1] = orig_vals[:, i] + NK
    # update pass: same key, new value, later in the sequence
    upd_tokens = torch.gather(keys, 1, updated_idx)
    for j in range(n_upd):
        s = (n_pairs + j) * block
        if f:
            tokens[:, s:s + f] = filler_ids[:, (n_pairs + j) * f:(n_pairs + j + 1) * f]
        tokens[:, s + f] = upd_tokens[:, j]
        tokens[:, s + f + 1] = upd_vals[:, j] + NK
    # query pass
    S2 = (n_pairs + n_upd) * block
    for q in range(n_queries):
        tokens[:, S2 + 2 * q] = qkeys[:, q]
        tokens[:, S2 + 2 * q + 1] = cfg.SEP

    key_pos = (torch.arange(n_pairs) * block + f).unsqueeze(0).expand(B, -1).contiguous()
    upd_key_pos = ((n_pairs + torch.arange(n_upd)) * block + f).unsqueeze(0).expand(B, -1).contiguous() \
        if n_upd else None
    ans_pos = (S2 + 2 * torch.arange(n_queries) + 1).unsqueeze(0).expand(B, -1).contiguous()
    return {
        "tokens": tokens, "ans_pos": ans_pos, "gold": gold, "was_updated": was_updated,
        "key_pos": key_pos, "upd_key_pos": upd_key_pos, "n_pairs": n_pairs,
        "n_upd": n_upd, "L": L,
    }


# ---------------------------------------------------------------------------
# model helpers
# ---------------------------------------------------------------------------


def build_model(cfg: SynthCfg, n_layers: int = 1):
    from mulattice.models import MQARModel, MixerCfg
    mcfg = MixerCfg(d_model=cfg.d_model, n_heads=cfg.n_heads, d_k=cfg.d_k, d_v=cfg.d_v,
                    addr_scale=cfg.addr_scale, learn_addr_scale=False, bound_address=True,
                    read_tie=False, retention_init=cfg.retention_init)
    return MQARModel(cfg.V, kind="mulattice", n_layers=n_layers, cfg=mcfg)


def logits_at(model, tokens: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
    """Forward the trunk, then apply the head ONLY at `positions` (B,Q).

    Avoids materialising (B,T,V) logits; at T=131K with V=32,770 that is 17GB.
    """
    x = model.emb(tokens)
    for blk in model.blocks:
        x = blk(x, read_mode="all")
    x = model.norm(x)
    B = x.shape[0]
    h = x[torch.arange(B, device=x.device).unsqueeze(1), positions]
    return model.head(h)                       # (B,Q,V)


@torch.no_grad()
def eval_batch(model, batch, device) -> dict:
    """Accuracy of `model` on a synth_batch, split by updated/untouched queries."""
    model.eval()
    pred = logits_at(model, batch["tokens"].to(device), batch["ans_pos"].to(device)).argmax(-1).cpu()
    gold = batch["gold"]
    ok = pred == gold
    res = {"acc": ok.float().mean().item(), "n": ok.numel(),
           "correct": int(ok.sum())}
    up = batch["was_updated"]
    if bool(up.any()):
        res["acc_updated"] = ok[up].float().mean().item()
    if bool((~up).any()):
        res["acc_untouched"] = ok[~up].float().mean().item()
    model.train()
    return res


def scan_info(model) -> str:
    parts = []
    for i, blk in enumerate(model.blocks):
        sc = blk.mixer.scan
        parts.append(f"L{i} cells={sc.last_n_cells} rounds={getattr(sc, 'last_rounds', '?')}"
                     f" path={getattr(sc, 'last_path', '?')}")
    return " ".join(parts)


def force_fast_scan(model):
    for blk in model.blocks:
        blk.mixer.scan.crossover = 1e9


def key_addresses(model, tokens, key_pos):
    """Write addresses (B,N,2) of the tokens at `key_pos` (B,N).

    Uses the mixer's LIVE addr_scale (it can be ramped mid-run), not the
    construction-time cfg value."""
    mx = model.blocks[0].mixer
    xc = mx.conv(model.blocks[0].n1(model.emb(tokens)))
    kaddr = mx.cfg.addr_scale * torch.tanh(mx.Wa(xc))
    ar = torch.arange(tokens.shape[0], device=tokens.device).unsqueeze(1)
    return kaddr[ar, key_pos]


def spread_loss(model, tokens, key_pos, cfg: SynthCfg, gate_subsample: int = 2048):
    """Kernel-D gated spread loss on the TRUE key-token write addresses.

    Same definition as mulattice.addressing.gated_spread_loss, but the
    pairwise write-overlap gate is computed on a subsample when N is large
    (it is (B,N,N); at N=8,192 that is 67M entries per row)."""
    addr = key_addresses(model, tokens, key_pos)
    B, N, _ = addr.shape
    addr_gate = addr[:, torch.randperm(N, device=addr.device)[:gate_subsample]] if N > gate_subsample else addr
    return _gated_spread(addr, addr_gate, float(model.blocks[0].mixer.cfg.addr_scale))


def _gated_spread(addr, addr_gate, scale):
    from mulattice.addressing import write_overlap
    from mulattice.kernel import phi_2d
    B, N, _ = addr.shape
    gate = (write_overlap(addr_gate) - 0.05).clamp(0.0, 1.0).detach()
    lam = 0.3 + 0.7 * gate
    w, z = phi_2d(addr)
    size = int(2 * scale + 3)
    off = int(scale) + 1
    gi = (z + off).clamp(0, size - 1)
    idx = (gi[..., 0] * size + gi[..., 1]).reshape(B, -1).to(torch.int64)
    mass = (w / w.sum(dim=(-2, -1), keepdim=True)).reshape(B, -1)
    hist = torch.zeros(B, size * size, device=addr.device, dtype=addr.dtype)
    for b in range(B):
        hist[b].index_add_(0, idx[b], mass[b])
    p = hist / hist.sum(dim=1, keepdim=True).clamp_min(1e-9)
    entropy = -(p * (p + 1e-12).log()).sum(dim=1)
    return (lam * (-entropy)).mean()


# ---------------------------------------------------------------------------
# vocab surgery: extend the proven 1024-key model to the full key vocabulary
# ---------------------------------------------------------------------------


def extend_from_checkpoint(new_model, ckpt_path: str, NK: int = 16384, NV: int = 16384,
                           old_NK: int = 1024, old_NV: int = 1024,
                           noise: float = 0.05, seed: int = 0):
    """Copy the proven model's mixer into `new_model`; extend emb/head rows.

    New key/value ids are initialised from a donor id (j mod 1024) plus small
    noise, so fresh ids start inside the working recall circuit instead of as
    random noise -- extending a working vocabulary rather than re-bootstrapping.
    From-scratch training at V>=16K historically never escaped the chance plateau
    (logs/vocab8192.log, logs/vocabcur.log).
    """
    ck = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    old = ck["model"]
    new = {k: (v.clone() if k not in ("emb.weight", "head.weight") else None)
           for k, v in new_model.state_dict().items()}
    # every non-emb/head parameter has identical shape -> straight copy
    for k in new:
        if new[k] is None:
            continue
        if k not in old or old[k].shape != new[k].shape:
            raise KeyError(f"cannot warm-start: parameter {k} missing/shape-mismatched")
        new[k] = old[k].clone()
    g = torch.Generator().manual_seed(seed)

    def build_rows(rows_old, std, off_new, span_new, off_old, span_old):
        out = torch.empty(span_new, rows_old.shape[1])
        donors = torch.arange(span_new) % span_old
        out[:] = rows_old[off_old + donors]
        out += torch.randn(out.shape, generator=g) * (noise * std)
        return out

    for pname in ("emb.weight", "head.weight"):
        rows_old, rows_new = old[pname], new_model.state_dict()[pname]
        std = rows_old.std().item()
        # keys: first old_NK ids keep their trained rows, the rest inherit donors
        rows_new[:old_NK] = rows_old[:old_NK]
        rows_new[old_NK:NK] = build_rows(rows_old, std, old_NK, NK - old_NK, 0, old_NK)
        # values (offset NK in the new layout, old_NK in the old one)
        rows_new[NK:NK + old_NV] = rows_old[old_NK:old_NK + old_NV]
        rows_new[NK + old_NV:NK + NV] = build_rows(rows_old, std, 0, NV - old_NV, old_NK, old_NV)
        # unused id + SEP
        rows_new[NK + NV] = rows_old[old_NK + old_NV]
        rows_new[NK + NV + 1] = rows_old[old_NK + old_NV + 1]
        new[pname] = rows_new
    new_model.load_state_dict(new)
    return new_model


# ---------------------------------------------------------------------------
# real-data protocol (same mining as experiments/eval_realdata.py)
# ---------------------------------------------------------------------------


class RealData:
    def __init__(self, corpus_path: str, VS: int = 4094):
        from collections import Counter
        self.VS = VS
        words = open(corpus_path, encoding="utf-8", errors="replace").read().lower().split()
        cnt = Counter(words)
        vocab = [w for w, _ in cnt.most_common(VS - 1)]
        wid = {w: i for i, w in enumerate(vocab)}
        big = Counter()
        for a, b in zip(words, words[1:]):
            if a in wid and b in wid:
                big[(wid[a], wid[b])] += 1
        succ = {}
        for (a, b), c in big.items():
            if c >= 4 and (a not in succ or succ[a][1] < c):
                succ[a] = (b, c)
        kv = {}
        ranked = sorted(succ.keys(), key=lambda k: -cnt[vocab[k]])
        lo, hi = int(len(ranked) * 0.02), int(len(ranked) * 0.95)
        for k in ranked[lo:hi]:
            v = succ[k][0]
            if v != k and cnt[vocab[v]] >= 10:
                kv[k] = v
        self.KVLIST = torch.tensor(sorted(kv.keys()))
        self.KVARR = torch.tensor([kv[int(k)] for k in self.KVLIST])
        # Zipfian position weights (cnt^0.75, as in realdata256k.py sample_pairs):
        # concentrating on frequent keys is what let the old run transition in
        # <400 steps; uniform sampling left rare-key embeddings gradient-starved
        # and the model stuck at 0.35 (measured 2026-11-02, realdata_v2 first run).
        w = torch.tensor([float(cnt[vocab[int(k)]]) ** 0.75 for k in self.KVLIST])
        self.KEY_W = w / w.sum()
        self.FILLER = torch.tensor([wid[w] for w in words if w in wid], dtype=torch.long)
        self.KVOCAB = VS
        self.SEP = 2 * VS
        self.V = 2 * VS + 2
        self.n_pairs = len(self.KVLIST)

    def sample_positions(self, B: int, P: int, g: torch.Generator, zipf: bool) -> torch.Tensor:
        """(B,P) distinct KVLIST positions; uniform or Zipf-weighted per row."""
        if not zipf:
            return torch.stack([torch.randperm(self.n_pairs, generator=g)[:P] for _ in range(B)])
        rows = []
        for _ in range(B):
            cand = torch.multinomial(self.KEY_W, 2 * P, replacement=True, generator=g).tolist()
            seen, out, j = set(), [], 0
            while len(out) < P and j < len(cand):
                k = cand[j]; j += 1
                if k not in seen:
                    seen.add(k); out.append(k)
            if len(out) < P:                      # exhausted weighted draws: fill uniformly
                rest = [i for i in torch.randperm(self.n_pairs, generator=g).tolist()
                        if i not in seen][: P - len(out)]
                out += rest
            rows.append(out)
        return torch.tensor(rows)

    def batch(self, P: int, filler: int, Q: int, seed: int, B: int = 1,
              overwrite_frac: float = 0.0, qpos_idx: torch.Tensor | None = None,
              zipf: bool = False):
        """Real-data store/query batch. qpos_idx: (Q,) indices into the P stored
        pairs (allows depth-stratified query sets); overwrite re-presents chosen
        keys with a DIFFERENT value token (another pair's value)."""
        g = torch.Generator().manual_seed(seed)
        P = min(P, self.n_pairs)
        perms = self.sample_positions(B, P, g, zipf)
        keys = self.KVLIST[perms]                       # (B,P) key token ids
        vals = self.KVARR[perms]                        # (B,P) value token ids
        n_upd = int(round(P * overwrite_frac))
        upd_idx = None
        upd_vals = None
        if n_upd:
            upd_idx = torch.stack([torch.randperm(P, generator=g)[:n_upd] for _ in range(B)])
            # update value = another pair's value; redraw the rare self-collisions
            shift = torch.stack([torch.randperm(P, generator=g)[:n_upd] for _ in range(B)])
            upd_vals = torch.gather(vals, 1, shift)
            same = upd_vals == torch.gather(vals, 1, upd_idx)
            while bool(same.any()):
                # rare: key's successor equals donor's successor -> bump by one token id
                flat = same.reshape(-1)
                upd_vals.reshape(-1)[flat] = (upd_vals.reshape(-1)[flat] + 1) % self.VS
                same = upd_vals == torch.gather(vals, 1, upd_idx)
        if qpos_idx is None:
            qpos_idx = torch.stack([torch.randperm(P, generator=g)[:Q] for _ in range(B)])
        elif qpos_idx.dim() == 1:
            qpos_idx = qpos_idx.unsqueeze(0).expand(B, -1).contiguous()
        Q = qpos_idx.shape[1]          # the actual query count (clamped to P)
        qkeys = torch.gather(keys, 1, qpos_idx)

        f = filler
        block = f + 2
        L = (P + n_upd) * block + 2 * Q
        tokens = torch.zeros(B, L, dtype=torch.long)
        fptr = torch.randint(0, max(1, len(self.FILLER) - f * (P + n_upd) - 1), (1,),
                             generator=g).item()
        for i in range(P):
            s = i * block
            if f:
                tokens[:, s:s + f] = self.FILLER[fptr:fptr + f]
                fptr += f
            tokens[:, s + f] = keys[:, i]
            tokens[:, s + f + 1] = vals[:, i] + self.KVOCAB
        if n_upd:
            upd_keys = torch.gather(keys, 1, upd_idx)
            for j in range(n_upd):
                s = (P + j) * block
                if f:
                    tokens[:, s:s + f] = self.FILLER[fptr:fptr + f]
                    fptr += f
                tokens[:, s + f] = upd_keys[:, j]
                tokens[:, s + f + 1] = upd_vals[:, j] + self.KVOCAB
        S2 = (P + n_upd) * block
        latest = vals.clone()
        if n_upd:
            latest.scatter_(1, upd_idx, upd_vals)
        gold = torch.gather(latest, 1, qpos_idx) + self.KVOCAB
        for q in range(Q):
            tokens[:, S2 + 2 * q] = qkeys[:, q]
            tokens[:, S2 + 2 * q + 1] = self.SEP
        key_pos = (torch.arange(P) * block + f).unsqueeze(0).expand(B, -1).contiguous()
        upd_key_pos = ((P + torch.arange(n_upd)) * block + f).unsqueeze(0).expand(B, -1).contiguous() \
            if n_upd else None
        # ANSWER AT THE QUERY-KEY POSITION (even offsets), matching the original
        # realdata256k.py training and eval_realdata.py: `labels[L-2*(n-j)]` is the
        # query-key token, not the SEP after it.  (The SYNTHETIC task answers at
        # the SEP -- mqar.py `ans_pos = arange(1, 2*n_queries, 2)` -- two different
        # conventions, both from the previous session; the warm checkpoint knows
        # the query-key one.  Measured: same tokens score 0.875 under the key
        # convention and 0.000 under the SEP convention.)
        ans_pos = (S2 + 2 * torch.arange(Q)).unsqueeze(0).expand(B, -1).contiguous()
        was_upd = torch.zeros(B, Q, dtype=torch.bool)
        if n_upd:
            m = torch.zeros(B, P, dtype=torch.bool)
            m.scatter_(1, upd_idx, True)
            was_upd = torch.gather(m, 1, qpos_idx)
        return {"tokens": tokens, "ans_pos": ans_pos, "gold": gold, "was_updated": was_upd,
                "key_pos": key_pos, "upd_key_pos": upd_key_pos, "n_pairs": P, "n_upd": n_upd,
                "L": L}


def stratified_qpos(P: int, Q: int, strata: int = 4, seed: int = 0):
    """Query indices stratified over store depth: each stratum gets Q/strata
    queries drawn from pairs stored at that depth fraction of the sequence."""
    g = torch.Generator().manual_seed(seed)
    per = Q // strata
    out = []
    bounds = torch.linspace(0, P, strata + 1).long()
    for s in range(strata):
        lo, hi = int(bounds[s]), int(bounds[s + 1])
        out.append(lo + torch.randperm(hi - lo, generator=g)[:per])
    return torch.cat(out)