probes/probe_memory_source_readout.py

#!/usr/bin/env python3
"""Probe whether a memory source vector is directly decodable.

This diagnostic freezes the trained CEDL checkpoint and trains a temporary
linear vocab readout from one selected source vector:

  h_d, h_e, q_attractor, or q_mem

It bypasses the external bank bottleneck:

    source -> direct_head -> vocab

instead of:

    source -> bank_q_proj -> 256-slot mem_keys/mem_vals -> mem_head_bank -> vocab

If direct_head generalizes on B2 while the bank path fails, the bottleneck is
the learned 256-slot bank readout. If direct_head also fails, the selected
source does not robustly encode the current-token target under this synthetic
held-out probe.
"""

import argparse
import json
import os
import sys
from typing import Dict, List, Tuple

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, "/content")

import CEDL
import data_v4c_pairs as v4c


SOURCES = ("h_d", "h_e", "q_attractor", "q_mem")
MEMORY_GATE_ATTR = "v" + "6_lambda_head"
STORE_SOURCES_ATTR = "store_v" + "6_bank_sources"
NEEDS_SOURCES_METHOD = "_v" + "6_needs_dstage_bank_sources"
SOURCE_INPUT_METHOD = "_v" + "6_bank_query_input"

SOURCE_ALIASES = {
    "contextual_memory_state": "q_mem",
    "decoder_state": "h_d",
    "expanded_state": "h_e",
    "attractor_state": "q_attractor",
    "q_mem": "q_mem",
    "h_d": "h_d",
    "h_e": "h_e",
    "q_attractor": "q_attractor",
}


def load_sidecar_constructor_kwargs(
    sidecar_path: str,
    source: str | None,
) -> Dict[str, object]:
    with open(sidecar_path) as f:
        sc = json.load(f)
    mem = sc.get("memory_readout")
    if not isinstance(mem, dict):
        raise ValueError("Expected cedl_config.json with a memory_readout block.")
    source_name = str(source or mem.get("source", "contextual_memory_state"))
    return dict(
        lambda_head=bool(mem.get("lambda_head", True)),
        lambda_head_hidden=int(mem.get("lambda_head_hidden", 160)),
        lambda_head_bias_init=float(mem.get("lambda_head_bias_init", -7.0)),
        lambda_head_w_init_std=float(
            mem.get("lambda_head_w_init_std", 0.05)),
        bce_objective=(
            mem.get("selection_objective") == "binary_answer_background"),
        sel_weight=1.0,
        bg_weight=1.0,
        bg_target=float(mem.get("background_target", 0.01)),
        wt_sparsity_weight=float(mem.get("sparsity_weight", 0.05)),
        wt_sparsity_target=float(mem.get("sparsity_target", 0.05)),
        memory_head_enabled=bool(mem.get("enabled", True)),
        memory_ce_weight=float(mem.get("memory_ce_weight", 1.0)),
        memory_pair_ce_weight=float(mem.get("pair_ce_weight", 5.0)),
        memory_query_source=SOURCE_ALIASES.get(source_name, source_name),
        memory_readout_mode="direct",
        source_adapter=bool(mem.get("source_adapter", True)),
        context_adapter=bool(mem.get("context_adapter", True)),
        specialist_noinject=bool(mem.get("no_injection", True)),
    )


def load_model(args, device: torch.device):
    source = None if args.source == "sidecar" else args.source
    model_kwargs = load_sidecar_constructor_kwargs(args.sidecar, source)
    model = CEDL.build_model("CEDL", vocab=50257, max_seq=1024, **model_kwargs)
    model = model.to(device).eval()

    state = torch.load(args.checkpoint, map_location="cpu", weights_only=True)
    msd = state["model"] if isinstance(state, dict) and "model" in state else state
    if any(k.startswith("_orig_mod.") for k in msd):
        msd = {k.replace("_orig_mod.", ""): v for k, v in msd.items()}
    res = model.load_state_dict(msd, strict=True)
    print(f"[load] strict OK  missing={len(res.missing_keys)}  "
          f"unexpected={len(res.unexpected_keys)}")

    model.feedback_alpha.fill_(1.0)
    if hasattr(model, "sl_alpha"):
        model.sl_alpha.fill_(1.0)
    for p in model.parameters():
        p.requires_grad_(False)
    return model


def make_batch(tokenizer, batch_size: int, max_seq: int, seed: int,
               split: str, hard_collision_frac: float):
    items = v4c.generate(
        tokenizer,
        n=batch_size,
        seed=seed,
        split=split,
        hard_collision_frac=hard_collision_frac,
        family_weights={
            "but_update": 0.50,
            "however_revision": 0.20,
            "temporal_update": 0.15,
            "paraphrased_equiv": 0.05,
            "neutral_control": 0.10,
        },
    )
    return CEDL._pad_v4c_items(items, max_seq), items


@torch.no_grad()
def answer_source(model, ids: torch.Tensor, items) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    h_C = model.c_stage(ids, feedback=None)
    h_E, h_E_sparse = model.e_stage(h_C, feedback=None)
    v_vec, _v_scalar = model.salience(h_E)

    prev = getattr(model.d_stage, STORE_SOURCES_ATTR, False)
    setattr(model.d_stage, STORE_SOURCES_ATTR,
            getattr(model, NEEDS_SOURCES_METHOD)())
    try:
        h_D, _ = model.d_stage(h_E, h_E_sparse, v_vec=v_vec)
    finally:
        setattr(model.d_stage, STORE_SOURCES_ATTR, prev)

    b_idx, p_idx, cur_tok, stale_tok = CEDL._v4c_collect_answer_quads(ids, items)
    if b_idx.numel() == 0:
        empty = torch.empty(0, model.d_model, device=ids.device)
        return empty, cur_tok, stale_tok
    src = getattr(model, SOURCE_INPUT_METHOD)(h_D, h_E, b_idx, p_idx)
    return src.detach(), cur_tok, stale_tok


def init_direct_head(model, init_from: str) -> nn.Linear:
    head = nn.Linear(model.d_model, 50257, bias=True).to(next(model.parameters()).device)
    with torch.no_grad():
        if init_from == "mem_head_bank" and hasattr(model, "mem_head_bank"):
            head.weight.copy_(model.mem_head_bank.weight)
            head.bias.copy_(model.mem_head_bank.bias)
        elif init_from == "tok_emb":
            head.weight.copy_(model.c_stage.tok_emb.weight)
            if getattr(model.l_stage.mem_head, "bias", None) is not None:
                head.bias.copy_(model.l_stage.mem_head.bias)
            else:
                head.bias.zero_()
        elif init_from == "random":
            nn.init.normal_(head.weight, std=0.02)
            head.bias.zero_()
        else:
            raise ValueError(f"unknown --init-from {init_from!r}")
    return head


def row_metrics(logits: torch.Tensor, cur_tok: torch.Tensor, stale_tok: torch.Tensor):
    logp = F.log_softmax(logits, dim=-1)
    top1 = logp.argmax(dim=-1)
    entropy = -(logp.exp() * logp).sum(dim=-1)
    cur_rank = (logp > logp.gather(1, cur_tok[:, None])).sum(dim=1) + 1
    stale_rank = (logp > logp.gather(1, stale_tok[:, None])).sum(dim=1) + 1
    margin = (
        logp.gather(1, cur_tok[:, None]).squeeze(1)
        - logp.gather(1, stale_tok[:, None]).squeeze(1)
    )
    return logp, top1, entropy, cur_rank, stale_rank, margin


@torch.no_grad()
def evaluate_replay(model, head, tokenizer, *, n_items: int, batch_size: int,
                    max_seq: int, seed: int, split: str) -> Dict[str, float]:
    rows: List[Dict[str, float]] = []
    produced = 0
    batch_seed = seed
    while produced < n_items:
        need = min(batch_size, n_items - produced)
        ids, items = make_batch(
            tokenizer, need, max_seq, batch_seed, split,
            hard_collision_frac=0.0,
        )
        ids = ids.to(next(model.parameters()).device)
        src, cur_tok, stale_tok = answer_source(model, ids, items)
        batch_seed += 1009
        if src.numel() == 0:
            continue
        logits = head(src)
        _logp, top1, entropy, cur_rank, stale_rank, margin = row_metrics(
            logits, cur_tok, stale_tok)
        for i in range(cur_tok.numel()):
            rows.append({
                "top1_cur": float(top1[i].item() == cur_tok[i].item()),
                "top1_stale": float(top1[i].item() == stale_tok[i].item()),
                "entropy": float(entropy[i].item()),
                "cur_rank": float(cur_rank[i].item()),
                "stale_rank": float(stale_rank[i].item()),
                "margin": float(margin[i].item()),
            })
        produced += need
    if not rows:
        raise RuntimeError("evaluation generated no valid V4c answer rows")
    return {
        "n": float(len(rows)),
        "top1_cur": float(np.mean([r["top1_cur"] for r in rows])),
        "top1_stale": float(np.mean([r["top1_stale"] for r in rows])),
        "entropy": float(np.mean([r["entropy"] for r in rows])),
        "cur_rank_mean": float(np.mean([r["cur_rank"] for r in rows])),
        "cur_rank_median": float(np.median([r["cur_rank"] for r in rows])),
        "stale_rank_mean": float(np.mean([r["stale_rank"] for r in rows])),
        "stale_rank_median": float(np.median([r["stale_rank"] for r in rows])),
        "margin": float(np.mean([r["margin"] for r in rows])),
    }


def print_metrics(label: str, metrics: Dict[str, float]):
    print(f"[{label}] n={int(metrics['n'])}  "
          f"top1_cur={metrics['top1_cur'] * 100:5.1f}%  "
          f"top1_stale={metrics['top1_stale'] * 100:5.1f}%  "
          f"entropy={metrics['entropy']:.3f}  "
          f"cur_rank_mean={metrics['cur_rank_mean']:.0f}  "
          f"cur_rank_med={metrics['cur_rank_median']:.0f}  "
          f"stale_rank_med={metrics['stale_rank_median']:.0f}  "
          f"margin={metrics['margin']:+.3f}")


def token_decode(tok, token_id):
    s = tok.decode([int(token_id)])
    return repr(s) if len(s) < 16 else repr(s[:13]) + "..."


class ConstantLambdaHead(nn.Module):
    def __init__(self, logit_value):
        super().__init__()
        self.logit_value = float(logit_value)

    def forward(self, h):
        return torch.full(
            h.shape[:-1] + (1,), self.logit_value,
            device=h.device, dtype=h.dtype,
        )


def model_forward_logits(m, ids_b):
    out = m(ids_b)
    return out[0] if isinstance(out, tuple) else out


@torch.no_grad()
def evaluate_b2_direct(model, head, tokenizer, *, n_items: int, seed: int):
    device = next(model.parameters()).device
    items = v4c.generate(tokenizer, n=n_items, seed=seed)
    records = []
    saved_head = getattr(model, MEMORY_GATE_ATTR)
    setattr(model, MEMORY_GATE_ATTR, ConstantLambdaHead(-100.0).to(device))
    try:
        for it_idx, it in enumerate(items):
            if it.family == "neutral_control":
                continue
            if not it.current or not it.stale:
                continue
            ol = getattr(it, "original_length", 0)
            if ol <= 1 or ol > 1024:
                continue
            cur_t = int(it.ids[it.current[0][0]])
            stale_t = int(it.ids[it.stale[0][0]])
            if cur_t == stale_t:
                continue
            ans_p = ol - 1
            ids_b = torch.tensor([it.ids[:ol]], device=device, dtype=torch.long)

            trunk_logits = model_forward_logits(model, ids_b)[0, ans_p]
            src, cur_tok, stale_tok = answer_source(model, ids_b, [it])
            if src.numel() == 0:
                continue
            direct_logits = head(src)[0]

            results = {"item": it_idx, "cur_t": cur_t, "stale_t": stale_t}
            for path, row in (("trunk", trunk_logits), ("direct", direct_logits)):
                logp = F.log_softmax(row, dim=-1)
                top5_p, top5_t = torch.topk(logp, k=5)
                results[path] = {
                    "top5_t": top5_t.tolist(),
                    "top5_lp": top5_p.tolist(),
                    "top1_t": int(top5_t[0].item()),
                    "ent": float(-(logp.exp() * logp).sum().item()),
                    "cur_rank": int((logp > logp[cur_t]).sum().item()) + 1,
                    "stl_rank": int((logp > logp[stale_t]).sum().item()) + 1,
                    "cur_lp": float(logp[cur_t].item()),
                    "stl_lp": float(logp[stale_t].item()),
                }
            records.append(results)
    finally:
        setattr(model, MEMORY_GATE_ATTR, saved_head)

    n = len(records)
    print(f"\n[b2-direct] generated={len(items)}  used={n}")
    if n == 0:
        return {}

    for path in ("trunk", "direct"):
        top1_cur = sum(1 for r in records if r[path]["top1_t"] == r["cur_t"])
        top1_stl = sum(1 for r in records if r[path]["top1_t"] == r["stale_t"])
        top1_other = n - top1_cur - top1_stl
        ents = np.array([r[path]["ent"] for r in records])
        cur_ranks = np.array([r[path]["cur_rank"] for r in records])
        stl_ranks = np.array([r[path]["stl_rank"] for r in records])
        cur_lps = np.array([r[path]["cur_lp"] for r in records])
        stl_lps = np.array([r[path]["stl_lp"] for r in records])
        print(f"[{path:>6}] top1: cur={top1_cur:>2}/{n} ({top1_cur/n*100:5.1f}%)  "
              f"stl={top1_stl:>2}/{n} ({top1_stl/n*100:5.1f}%)  "
              f"other={top1_other:>2}/{n} ({top1_other/n*100:5.1f}%)")
        print(f"         entropy: mean={ents.mean():.3f}  std={ents.std():.3f}")
        print(f"         cur_rank: mean={cur_ranks.mean():.0f}  median={np.median(cur_ranks):.0f}  "
              f"min={cur_ranks.min()}  max={cur_ranks.max()}")
        print(f"         stl_rank: mean={stl_ranks.mean():.0f}  median={np.median(stl_ranks):.0f}  "
              f"min={stl_ranks.min()}  max={stl_ranks.max()}")
        print(f"         logp(cur): mean={cur_lps.mean():.3f}  logp(stl): mean={stl_lps.mean():.3f}\n")

    direct_eq_trunk = sum(
        1 for r in records if r["direct"]["top1_t"] == r["trunk"]["top1_t"])
    direct_in_trunk_top5 = sum(
        1 for r in records if r["direct"]["top1_t"] in r["trunk"]["top5_t"])
    trunk_in_direct_top5 = sum(
        1 for r in records if r["trunk"]["top1_t"] in r["direct"]["top5_t"])
    print("[cross-path]")
    print(f"  direct top1 == trunk top1:      {direct_eq_trunk:>2}/{n} ({direct_eq_trunk/n*100:5.1f}%)")
    print(f"  direct top1 in trunk top5:      {direct_in_trunk_top5:>2}/{n} ({direct_in_trunk_top5/n*100:5.1f}%)")
    print(f"  trunk top1 in direct top5:      {trunk_in_direct_top5:>2}/{n} ({trunk_in_direct_top5/n*100:5.1f}%)")

    print("\n[first 10 - qualitative direct vs trunk top-3 at answer position]")
    for r in records[:10]:
        cur_s = token_decode(tokenizer, r["cur_t"])
        stl_s = token_decode(tokenizer, r["stale_t"])
        print(f"\n  --- item {r['item']}  cur={cur_s}  stale={stl_s} ---")
        for path in ("trunk", "direct"):
            top3_t = r[path]["top5_t"][:3]
            top3_lp = r[path]["top5_lp"][:3]
            top3_decoded = [
                (token_decode(tokenizer, t), f"{lp:+.2f}")
                for t, lp in zip(top3_t, top3_lp)
            ]
            print(f"  {path:>6}: top3={top3_decoded}  "
                  f"cur_rank={r[path]['cur_rank']}  stl_rank={r[path]['stl_rank']}")

    direct_top1_cur = sum(
        1 for r in records if r["direct"]["top1_t"] == r["cur_t"]) / n
    trunk_top1_cur = sum(
        1 for r in records if r["trunk"]["top1_t"] == r["cur_t"]) / n
    print(f"\n{'='*64}")
    print("Direct-source readout verdict")
    print(f"{'='*64}")
    if direct_top1_cur >= trunk_top1_cur:
        print(f"  DIRECT >= TRUNK ({direct_top1_cur*100:.1f}% vs {trunk_top1_cur*100:.1f}%).")
        print("  -> Source is decodable; 256-slot bank readout is the bottleneck.")
    elif direct_top1_cur >= 0.5 * trunk_top1_cur:
        print(f"  DIRECT PARTIAL ({direct_top1_cur*100:.1f}% vs trunk {trunk_top1_cur*100:.1f}%).")
        print("  -> Source has signal, but direct readout is still not robust enough.")
    else:
        print(f"  DIRECT FAIL ({direct_top1_cur*100:.1f}% vs trunk {trunk_top1_cur*100:.1f}%).")
        print("  -> Selected source is not robustly decodable under B2.")
    print(f"{'='*64}")

    return {
        "n": n,
        "direct_top1_cur": direct_top1_cur,
        "trunk_top1_cur": trunk_top1_cur,
    }


def save_outputs(head, args, before, after, b2):
    os.makedirs(args.out_dir, exist_ok=True)
    stem = f"CEDL_direct_source_{args.source}"
    head_path = os.path.join(args.out_dir, f"{stem}_head.pt")
    metrics_path = os.path.join(args.out_dir, f"{stem}_metrics.json")
    torch.save({
        "direct_head": head.state_dict(),
        "checkpoint": args.checkpoint,
        "sidecar": args.sidecar,
        "source": args.source,
        "init_from": args.init_from,
        "steps": args.steps,
        "lr": args.lr,
        "batch_size": args.batch_size,
        "seed": args.seed,
    }, head_path)
    with open(metrics_path, "w") as f:
        json.dump({"before": before, "after": after, "b2": b2}, f, indent=2)
    print(f"[save] direct_head={head_path}")
    print(f"[save] metrics={metrics_path}")


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--checkpoint", required=True)
    p.add_argument("--sidecar", required=True)
    p.add_argument("--out-dir", default="outputs/memory_source_readout")
    p.add_argument("--source", choices=("sidecar",) + SOURCES, default="sidecar")
    p.add_argument("--init-from", choices=("mem_head_bank", "tok_emb", "random"),
                   default="mem_head_bank")
    p.add_argument("--steps", type=int, default=2000)
    p.add_argument("--batch-size", type=int, default=64)
    p.add_argument("--max-seq", type=int, default=128)
    p.add_argument("--lr", type=float, default=5e-4)
    p.add_argument("--weight-decay", type=float, default=0.0)
    p.add_argument("--seed", type=int, default=7234)
    p.add_argument("--split", choices=("all", "train", "test"), default="all")
    p.add_argument("--hard-collision-frac", type=float, default=0.2)
    p.add_argument("--eval-items", type=int, default=256)
    p.add_argument("--b2-items", type=int, default=50)
    p.add_argument("--b2-seed", type=int, default=0)
    p.add_argument("--log-every", type=int, default=100)
    args = p.parse_args()

    torch.manual_seed(args.seed)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"[setup] device={device}")
    print(f"[setup] checkpoint={args.checkpoint}")
    print(f"[setup] sidecar={args.sidecar}")
    print(f"[setup] out_dir={args.out_dir}")
    print(f"[setup] source={args.source}")
    print(f"[setup] init_from={args.init_from}")

    from transformers import GPT2TokenizerFast
    tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")

    model = load_model(args, device)
    print(f"[setup] resolved_source={model_kwargs['memory_query_source']}")
    head = init_direct_head(model, args.init_from)
    trainable = sum(p.numel() for p in head.parameters())
    print(f"[freeze] trainable params={trainable:,} (direct_head only)")

    before = evaluate_replay(
        model, head, tokenizer, n_items=args.eval_items,
        batch_size=args.batch_size, max_seq=args.max_seq,
        seed=args.seed + 17, split=args.split,
    )
    print_metrics("before", before)

    opt = torch.optim.AdamW(
        head.parameters(), lr=args.lr, weight_decay=args.weight_decay,
        betas=(0.9, 0.95),
    )
    seen_rows = 0
    ema_loss = None
    for step in range(1, args.steps + 1):
        ids, items = make_batch(
            tokenizer, args.batch_size, args.max_seq,
            args.seed + step * 7919, args.split, args.hard_collision_frac,
        )
        ids = ids.to(device)
        src, cur_tok, _stale_tok = answer_source(model, ids, items)
        if src.numel() == 0:
            continue
        logits = head(src)
        loss = F.cross_entropy(logits, cur_tok)
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(head.parameters(), max_norm=1.0)
        opt.step()

        seen_rows += int(cur_tok.numel())
        loss_val = float(loss.detach().item())
        ema_loss = loss_val if ema_loss is None else 0.98 * ema_loss + 0.02 * loss_val
        if step == 1 or step % args.log_every == 0:
            with torch.no_grad():
                logp = F.log_softmax(logits, dim=-1)
                top1_cur = (logp.argmax(dim=-1) == cur_tok).float().mean().item()
                entropy = (-(logp.exp() * logp).sum(dim=-1)).mean().item()
            print(f"[train] step={step}/{args.steps} rows={seen_rows} "
                  f"loss={loss_val:.4f} ema={ema_loss:.4f} "
                  f"batch_top1_cur={top1_cur * 100:5.1f}% "
                  f"entropy={entropy:.3f}", flush=True)

    after = evaluate_replay(
        model, head, tokenizer, n_items=args.eval_items,
        batch_size=args.batch_size, max_seq=args.max_seq,
        seed=args.seed + 17, split=args.split,
    )
    print_metrics("after", after)
    b2 = evaluate_b2_direct(
        model, head, tokenizer, n_items=args.b2_items, seed=args.b2_seed)
    save_outputs(head, args, before, after, b2)


if __name__ == "__main__":
    main()