probes/probe_memory_causality.py

"""Memory-readout causality probe.

Runs the same current-vs-stale items through the live CEDL memory path and a
memory-off ablation. The paired comparison estimates whether contextual memory
readout causally improves current-token selection.

Usage:
  python probes/probe_memory_causality.py \
    --checkpoint pytorch_model.bin \
    --sidecar cedl_config.json \
    --n-items 50
"""

import argparse
import json
import os
import sys

import numpy as np
import torch
import torch.nn as nn

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

import CEDL
import data_v4c_pairs as v4c


MEMORY_GATE_ATTR = "v" + "6_lambda_head"


class ZeroLambdaHead(nn.Module):
    """Forces the memory mixture gate near zero."""
    def forward(self, h):
        return torch.full(
            h.shape[:-1] + (1,), -100.0,
            device=h.device, dtype=h.dtype,
        )


def load_sidecar_constructor_kwargs(sidecar_path):
    """Read public CEDL config and return constructor keyword arguments."""
    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(mem.get("source", "contextual_memory_state"))
    source_map = {
        "contextual_memory_state": "q_mem",
        "decoder_state": "h_d",
        "expanded_state": "h_e",
        "attractor_state": "q_attractor",
        "q_mem": "q_mem",
    }
    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_map.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 model_forward_logits(m, ids_b):
    """Call m(ids_b) and return the logits tensor [B, T, V] regardless of
    whether forward returns logits-only or (logits, aux_loss)."""
    out = m(ids_b)
    if isinstance(out, tuple):
        logits = out[0]
    else:
        logits = out
    return logits


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--checkpoint", type=str, default="pytorch_model.bin")
    p.add_argument("--sidecar", type=str, default="cedl_config.json")
    p.add_argument("--n-items", type=int, default=50)
    p.add_argument("--seed", type=int, default=0)
    args = p.parse_args()

    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] n_items={args.n_items}")

    model_kwargs = load_sidecar_constructor_kwargs(args.sidecar)
    print(f"[setup] memory_head={model_kwargs['lambda_head']}  "
          f"hidden={model_kwargs['lambda_head_hidden']}  "
          f"bias_init={model_kwargs['lambda_head_bias_init']}  "
          f"sparsity_weight={model_kwargs['wt_sparsity_weight']}  "
          f"sparsity_target={model_kwargs['wt_sparsity_target']}  "
          f"memory_source={model_kwargs['memory_query_source']}  "
          f"readout_mode={model_kwargs['memory_readout_mode']}")

    m = CEDL.build_model("CEDL", vocab=50257, max_seq=1024, **model_kwargs)
    m = m.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 = m.load_state_dict(msd, strict=True)
    print(f"[load] strict OK  (missing={len(res.missing_keys)}  unexpected={len(res.unexpected_keys)})")

    m.feedback_alpha.fill_(1.0)
    if hasattr(m, "sl_alpha"):
        m.sl_alpha.fill_(1.0)

    from transformers import GPT2TokenizerFast
    tok = GPT2TokenizerFast.from_pretrained("gpt2")
    print(f"\n[v4c] generating {args.n_items} items (seed={args.seed})...")
    items = v4c.generate(tok, n=args.n_items, seed=args.seed)
    print(f"[v4c] generated {len(items)} raw items")

    saved_head = getattr(m, MEMORY_GATE_ATTR)
    zero_head = ZeroLambdaHead().to(device)

    margin_active = []
    margin_off = []
    pred_correct_active = []
    pred_correct_off = []
    lam_at_ans_active = []
    n_skipped = 0
    skip_reasons = {"neutral_control": 0, "no_cur_stale": 0,
                    "bad_length": 0, "same_token": 0}

    for it_idx, it in enumerate(items):
        if it.family == "neutral_control":
            n_skipped += 1; skip_reasons["neutral_control"] += 1; continue
        if not it.current or not it.stale:
            n_skipped += 1; skip_reasons["no_cur_stale"] += 1; continue
        ol = getattr(it, "original_length", 0)
        if ol <= 1 or ol > 1024:
            n_skipped += 1; skip_reasons["bad_length"] += 1; continue
        cur_t = int(it.ids[it.current[0][0]])
        stale_t = int(it.ids[it.stale[0][0]])
        if cur_t == stale_t:
            n_skipped += 1; skip_reasons["same_token"] += 1; continue
        ans_p = ol - 1

        ids_b = torch.tensor([it.ids[:ol]], device=device, dtype=torch.long)

        setattr(m, MEMORY_GATE_ATTR, saved_head)
        with torch.no_grad():
            logits_a = model_forward_logits(m, ids_b)
        cur_lp_a = float(logits_a[0, ans_p, cur_t].item())
        stl_lp_a = float(logits_a[0, ans_p, stale_t].item())
        pred_a = int(logits_a[0, ans_p].argmax().item())
        margin_active.append(cur_lp_a - stl_lp_a)
        pred_correct_active.append(pred_a == cur_t)

        with torch.no_grad():
            h_C = m.c_stage(ids_b, feedback=None)
            h_E, h_E_sparse = m.e_stage(h_C, feedback=None)
            v_vec, _ = m.salience(h_E) if m.use_salience else (None, None)
            h_D, _ = m.d_stage(h_E, h_E_sparse, v_vec=v_vec)
            lam_logit_a = saved_head(h_D[:, ans_p:ans_p + 1, :])
            lam_a = torch.sigmoid(lam_logit_a).clamp(1e-4, 1.0 - 1e-4)
            lam_at_ans_active.append(float(lam_a.item()))

        setattr(m, MEMORY_GATE_ATTR, zero_head)
        with torch.no_grad():
            logits_o = model_forward_logits(m, ids_b)
        cur_lp_o = float(logits_o[0, ans_p, cur_t].item())
        stl_lp_o = float(logits_o[0, ans_p, stale_t].item())
        pred_o = int(logits_o[0, ans_p].argmax().item())
        margin_off.append(cur_lp_o - stl_lp_o)
        pred_correct_off.append(pred_o == cur_t)

    setattr(m, MEMORY_GATE_ATTR, saved_head)

    n = len(margin_active)
    print(f"\n[items] used={n}  skipped={n_skipped}  reasons={skip_reasons}")
    if n == 0:
        print("No valid items — aborting.")
        return

    ma = np.array(margin_active)
    mo = np.array(margin_off)
    delta = ma - mo

    print(f"\n[λ at answer row, ACTIVE path]  "
          f"mean={np.mean(lam_at_ans_active):.4f}  "
          f"std={np.std(lam_at_ans_active):.4f}  "
          f"min={min(lam_at_ans_active):.4f}  "
          f"max={max(lam_at_ans_active):.4f}")

    print(f"\n[current − stale logit margin at answer position]")
    print(f"  ACTIVE:    mean={ma.mean():+.3f}  std={ma.std():.3f}  "
          f"median={np.median(ma):+.3f}")
    print(f"  BANK-OFF:  mean={mo.mean():+.3f}  std={mo.std():.3f}  "
          f"median={np.median(mo):+.3f}")
    print(f"  Δ (act−off):  mean={delta.mean():+.3f}  std={delta.std():.3f}  "
          f"median={np.median(delta):+.3f}")

    from scipy import stats as sst
    t_stat, p_val = sst.ttest_rel(ma, mo)
    print(f"  paired t = {t_stat:+.3f}  p = {p_val:.4f}  N={n}")
    d_paired = float(delta.mean() / max(delta.std(), 1e-9))
    print(f"  Cohen's d (paired) = {d_paired:+.3f}")

    acc_a = float(np.mean(pred_correct_active))
    acc_o = float(np.mean(pred_correct_off))
    print(f"\n[top-1 accuracy: argmax(logits[ans_p]) == cur_t]")
    print(f"  ACTIVE:    {acc_a*100:5.1f}%   ({sum(pred_correct_active)}/{n})")
    print(f"  BANK-OFF:  {acc_o*100:5.1f}%   ({sum(pred_correct_off)}/{n})")
    print(f"  Δ:        {(acc_a-acc_o)*100:+5.1f} pp")

    print(f"\n[first 10 items — paired breakdown]")
    print(f"  {'#':>3}  {'m_A':>7}  {'m_O':>7}  {'Δ':>7}  "
          f"{'pred_A':>6}  {'pred_O':>6}")
    for i in range(min(10, n)):
        print(f"  {i:>3}  {ma[i]:>+7.3f}  {mo[i]:>+7.3f}  {delta[i]:>+7.3f}  "
              f"{'OK' if pred_correct_active[i] else 'XX':>6}  "
              f"{'OK' if pred_correct_off[i] else 'XX':>6}")

    print(f"\n{'='*64}")
    print(f"Memory-readout causality — verdict")
    print(f"{'='*64}")
    if t_stat > 2.0 and (acc_a - acc_o) > 0.10:
        verdict = (
            f"ACTIVE > BANK-OFF  (t={t_stat:+.2f}, Δacc=+{(acc_a-acc_o)*100:.1f}pp, "
            f"d={d_paired:+.2f}).\n"
            f"  → Bank is causally useful for V4c current-vs-stale.\n"
            f"  → PROCEED to full probe suite + manuscript update."
        )
    elif abs(t_stat) < 1.0 and abs(acc_a - acc_o) < 0.05:
        verdict = (
            f"ACTIVE ≈ BANK-OFF  (t={t_stat:+.2f}, Δacc={(acc_a-acc_o)*100:+.1f}pp, "
            f"d={d_paired:+.2f}).\n"
            f"  → λ head routes correctly but BANK READOUT IS NOT CAUSALLY USEFUL.\n"
            f"  → If mem_head_bank is already untied and alternate sources still\n"
            f"    fail, test direct-source readout before more bank replay."
        )
    elif t_stat < -2.0:
        verdict = (
            f"BANK-OFF > ACTIVE  (t={t_stat:+.2f}, Δacc={(acc_a-acc_o)*100:+.1f}pp, "
            f"d={d_paired:+.2f}).\n"
            f"  → BANK ACTIVELY HURTS V4c prediction. Diagnose the bank readout\n"
            f"    bottleneck before any manuscript work."
        )
    else:
        verdict = (
            f"INTERMEDIATE  (t={t_stat:+.2f}, Δacc={(acc_a-acc_o)*100:+.1f}pp, "
            f"d={d_paired:+.2f}).\n"
            f"  → Mixed signal. Run full probe suite to see broader pattern; if\n"
            f"    contrastive / stale-vs-current probes show clear lift, the bank\n"
            f"    is helpful on average even if the V4c-margin metric is noisy."
        )
    print(f"  {verdict}")
    print(f"{'='*64}")


if __name__ == "__main__":
    main()