proto_frame/proto_v02_entry.py

"""v0.2 sequence entry: position-specific linear maps G_0..G_3 emit
the length-matched embedding sequence for ANY state (k<=4), opening
entry from 149 whole-walk states to the full inventory. Same
generalization discipline: fit on train states, probe with UNSEEN
states; counterfactual restricted to SAME-k pairs (length confound
excluded); random-codebook control retained (Gate 4)."""
import json
import sys
import zlib

sys.path.insert(0, r"E:\mirel\geolip-bytelex")
sys.stdout.reconfigure(encoding="utf-8", errors="replace")
import numpy as np
import torch
import torch.nn.functional as F
import transformers
from transformers import AutoTokenizer, T5EncoderModel
import transformers.utils.logging as hlog

hlog.set_verbosity_error()
D = r"E:\mirel\data\bytelex\proto_frame"
DEV = "cuda"
SEED = zlib.crc32(b"frame-proto-v01") & 0xFFFFFFFF
L_T5 = 6
KMAX = 4
torch.manual_seed(SEED)
from geolip.bytelex.frame import (apply_whitening, fit_whitening,
                                  split_sites)

dump = np.load(rf"{D}\frame_dump_v2.npz")
anch = np.load(rf"{D}\frame_anchors_v01.npz")
walk = json.load(open(rf"{D}\t5_walk_of_C.json", encoding="utf-8"))
sid = dump["sid"].astype(np.int64)
KK = dump["KK"]
ok = KK[:, 0] > 0
sp = split_sites(sid[ok], seed=SEED)
gix = {k: np.flatnonzero(ok)[v] for k, v in sp.items()}

kk = np.array([w["k"] for w in walk])
ids_of = [w["ids"] for w in walk]
usable = np.flatnonzero((kk >= 1) & (kk <= KMAX))
KMAX = int(kk[usable].max())
print(f"[v02e] observed KMAX={KMAX}, k census: "
      f"{np.bincount(kk[usable]).tolist()}", flush=True)
rngw = np.random.default_rng(SEED + 41)
perm = rngw.permutation(len(usable))
gfit = usable[perm[:int(0.6 * len(usable))]]
gprobe = usable[perm[int(0.6 * len(usable)):]]
print(f"[v02e] usable {len(usable)}/999 (k<=4) -> fit {len(gfit)} / "
      f"probe {len(gprobe)}", flush=True)

tkA = AutoTokenizer.from_pretrained("google/flan-t5-small")
mT = T5EncoderModel.from_pretrained("google/flan-t5-small").to(DEV)
mT.eval()
for p in mT.parameters():
    p.requires_grad_(False)
emb = mT.get_input_embeddings().weight.detach()
E_C = torch.tensor(anch["E_C"], dtype=torch.float32, device=DEV)
E_Cn = F.normalize(E_C, dim=-1)
W_S = torch.tensor(anch["W_S"], dtype=torch.float32, device=DEV)
H_T = dump["H_T"][:, 0].astype(np.float64)
muT, wTw = fit_whitening(H_T[gix["fit"]])
tmu = torch.tensor(muT, dtype=torch.float32, device=DEV)
tw = torch.tensor(wTw, dtype=torch.float32, device=DEV)


def frame_readout(h):
    return F.normalize(((h - tmu) @ tw) @ W_S, dim=-1)


def fit_seq_g(E_use):
    """G_i: 256->512 per piece position, lstsq on fit states with
    k > i. Returns (list of G, per-position held-out residual)."""
    Gs, res = [], []
    for i in range(KMAX):
        fs = [int(u) for u in gfit if kk[u] > i]
        ps = [int(u) for u in gprobe if kk[u] > i]
        a = E_use[torch.tensor(fs, device=DEV)].double().cpu().numpy()
        b = emb[torch.tensor([ids_of[u][i] for u in fs], device=DEV)
                ].double().cpu().numpy()
        g, *_ = np.linalg.lstsq(a, b, rcond=None)
        if ps:
            ah = E_use[torch.tensor(ps, device=DEV)
                       ].double().cpu().numpy()
            bh = emb[torch.tensor([ids_of[u][i] for u in ps],
                                  device=DEV)].double().cpu().numpy()
            res.append(round(float(np.linalg.norm(ah @ g - bh)
                                   / np.linalg.norm(bh)), 4))
        Gs.append(torch.tensor(g, dtype=torch.float32, device=DEV))
    return Gs, res


G_byte, res_byte = fit_seq_g(E_C)
E_rand = F.normalize(torch.randn_like(E_C), dim=-1)
G_rand, res_rand = fit_seq_g(E_rand)
print(f"[v02e] held-out residuals byte={res_byte} rand={res_rand}",
      flush=True)

lines = open(r"E:\mirel\data\bytelex\codex_v1.txt",
             "rb").read().decode("ascii").split("\n")
line_tok = {}
def enc_line(li):
    if li not in line_tok:
        e = tkA(lines[li], return_offsets_mapping=True)
        line_tok[li] = (e["input_ids"], e["offset_mapping"])
    return line_tok[li]


probe_by_k = {}
for u in gprobe:
    probe_by_k.setdefault(int(kk[u]), []).append(int(u))
ev_cand = [k for k in gix["eval"]
           if int(kk[sid[k]]) in probe_by_k
           and len(probe_by_k[int(kk[sid[k]])]) > 1
           and KK[k, 1] == int(kk[sid[k]])]
rng = np.random.default_rng(SEED + 31)
ev_cand = rng.permutation(ev_cand)[:1200]
print(f"[v02e] probe sites: {len(ev_cand)}", flush=True)


def probe(Gs, E_use, tag):
    flips = stuck = n = 0
    div = 0.0
    nb = 0
    rngp = np.random.default_rng(SEED + 37)
    BS = 48
    for i in range(0, len(ev_cand), BS):
        seqs, spans, alts, poss = [], [], [], []
        for k in ev_cand[i:i + BS]:
            li = int(dump["line"][k])
            ids, off = enc_line(li)
            lo, hi = int(dump["lo"][k]), int(dump["hi"][k])
            ix = [j for j, (s, t) in enumerate(off)
                  if t > s and s < hi and t > lo]
            u = int(sid[k])
            cand = probe_by_k[int(kk[u])]
            up = cand[int(rngp.integers(0, len(cand)))]
            while up == u:
                up = cand[int(rngp.integers(0, len(cand)))]
            if len(ix) != int(kk[u]):
                continue
            seqs.append(ids)
            spans.append(ix)
            alts.append(up)
            poss.append(ix[0])
        if not seqs:
            continue
        mx = max(len(s) for s in seqs)
        idt = torch.full((len(seqs), mx), tkA.pad_token_id,
                         dtype=torch.long)
        att = torch.zeros((len(seqs), mx), dtype=torch.long)
        for j, s in enumerate(seqs):
            idt[j, :len(s)] = torch.tensor(s)
            att[j, :len(s)] = 1
        idt, att = idt.to(DEV), att.to(DEV)
        with torch.no_grad():
            clean = mT(input_ids=idt, attention_mask=att,
                       output_hidden_states=True).hidden_states[L_T5]
            x = emb[idt].clone()
            for j, (ix, up) in enumerate(zip(spans, alts)):
                for pi, p in enumerate(ix):
                    x[j, p] = E_use[up] @ Gs[pi]
            subh = mT(inputs_embeds=x, attention_mask=att,
                      output_hidden_states=True).hidden_states[L_T5]
            div += float((((subh - clean) ** 2).sum(-1) * att).sum()
                         / att.sum() / (clean ** 2).sum(-1).mean())
            nb += 1
            zr = frame_readout(subh[torch.arange(len(seqs)),
                                    torch.tensor(poss, device=DEV)])
            pred = (zr @ E_Cn.T).argmax(-1).cpu().numpy()
            flips += int((pred == np.array(alts)).sum())
            stuck += int((pred == np.array(
                [sid[k] for k in ev_cand[i:i + BS]][:len(seqs)])).sum())
            n += len(seqs)
    return {"n": n, "flip_to_injected": round(flips / max(n, 1), 4),
            "stuck_on_true": round(stuck / max(n, 1), 4),
            "rel_divergence": round(div / max(nb, 1), 4)}


res = {"byte_codebook": probe(G_byte, E_C, "byte"),
       "random_codebook_control": probe(G_rand, E_rand, "rand"),
       "heldout_residuals": {"byte": res_byte, "random": res_rand},
       "g_split": {"usable": int(len(usable)), "fit": int(len(gfit)),
                   "probe": int(len(gprobe))},
       "_env": {"transformers": transformers.__version__,
                "torch": torch.__version__}}
with open(rf"{D}\v02_entry_ledger.json", "w", encoding="utf-8") as f:
    json.dump(res, f, indent=1)
print(json.dumps(res, indent=1), flush=True)
print("[v02e] SEQUENCE ENTRY COMPLETE", flush=True)