proto_frame/proto_phase0.py

"""Frame prototype Phase 0 — instruments + full extraction (v0.1).

1. T5 walk table of C (instrument-first: entry loss needs whole-walk
   census; only the wordpiece walk existed).
2. Boundary-context profiles per state (prev/next byte distributions
   at attested sites) for the byte-structural codebook init.
3. Extraction v2: ALL sites, no cap. Unmasked anchor readouts
   (byte-anchored first subtoken + last subtoken, k per side) AND
   masked-span readouts (bert: [MASK]*k, read first mask; t5:
   <extra_id_0> replacing span, read sentinel in encoder) — the
   reading-vs-guessing control. CUDA, <18GB.
"""
import json
import sys

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

hlog.set_verbosity_error()
CODEX = r"E:\mirel\data\bytelex\codex_v1.txt"
WORDS = r"E:\mirel\data\bytelex\words_of_C.json"
OUTDIR = r"E:\mirel\data\bytelex\proto_frame"
SEPS = set(b" \t\n\r.,;:!?\"'()[]{}")
L_BERT, L_T5 = 8, 6
DEV = "cuda"
BATCH = 256

states = json.load(open(WORDS, encoding="utf-8"))
sid_of = {s["text"]: i for i, s in enumerate(states)}
tkB = AutoTokenizer.from_pretrained("bert-base-uncased")
tkA = AutoTokenizer.from_pretrained("google/flan-t5-small")

# ---- 1. t5 walk table of C
t5_walk = []
for s in states:
    ids = tkA(s["text"], add_special_tokens=False)["input_ids"]
    toks = tkA.convert_ids_to_tokens(ids)
    t5_walk.append({"text": s["text"], "count": s["count"],
                    "k": len(ids), "seg": "|".join(toks),
                    "ids": ids,
                    "whole": len(ids) == 1})
n_whole = sum(w["whole"] for w in t5_walk)
with open(rf"{OUTDIR}\t5_walk_of_C.json", "w", encoding="utf-8") as f:
    json.dump(t5_walk, f, indent=0)
print(f"[P0] t5 walk of C: {n_whole}/999 whole", flush=True)

# ---- sites (ALL, no cap) + 2. boundary-context profiles
lines = open(CODEX, "rb").read().decode("ascii").split("\n")
sites = []
prev_ctx = np.zeros((999, 256))
next_ctx = np.zeros((999, 256))
for li, ln in enumerate(lines):
    if not ln:
        continue
    raw = ln.encode("ascii") + b" "
    lo = None
    for j, ch in enumerate(raw):
        if ch in SEPS:
            if lo is not None:
                sid = sid_of.get(ln[lo:j])
                if sid is not None:
                    sites.append((li, lo, j, sid))
                    prev_ctx[sid, raw[lo - 1] if lo else 32] += 1
                    next_ctx[sid, ch] += 1
                lo = None
        elif lo is None:
            lo = j
np.savez_compressed(rf"{OUTDIR}\ctx_profiles.npz",
                    prev_ctx=prev_ctx, next_ctx=next_ctx)
print(f"[P0] {len(sites)} sites (uncapped), ctx profiles saved",
      flush=True)

mB = AutoModel.from_pretrained("bert-base-uncased").to(DEV).eval()
mT = T5EncoderModel.from_pretrained("google/flan-t5-small").to(DEV).eval()
embT = mT.get_input_embeddings()
SENT = tkA.convert_tokens_to_ids("<extra_id_0>")
MASK = tkB.mask_token_id

N = len(sites)
H_B = np.zeros((N, 2, 768), dtype=np.float16)     # first,last
H_T = np.zeros((N, 2, 512), dtype=np.float32)
M_B = np.zeros((N, 768), dtype=np.float16)        # masked readout
M_T = np.zeros((N, 512), dtype=np.float32)
E_T = np.zeros((N, 512), dtype=np.float32)        # t5 input emb (first)
KK = np.zeros((N, 2), dtype=np.int16)             # k_B, k_T

# ---- unmasked pass, batched by line
by_line = {}
for k, (li, lo, hi, sid) in enumerate(sites):
    by_line.setdefault(li, []).append(k)
line_ids = sorted(by_line)
tokcacheB, tokcacheT = {}, {}
with torch.no_grad():
    for bs in range(0, len(line_ids), BATCH):
        chunk = line_ids[bs:bs + BATCH]
        texts = [lines[li] for li in chunk]
        eb = tkB(texts, return_offsets_mapping=True, padding=True,
                 return_tensors="pt")
        et = tkA(texts, return_offsets_mapping=True, padding=True,
                 return_tensors="pt")
        hb = mB(input_ids=eb["input_ids"].to(DEV),
                attention_mask=eb["attention_mask"].to(DEV),
                output_hidden_states=True).hidden_states[L_BERT].cpu()
        ht = mT(input_ids=et["input_ids"].to(DEV),
                attention_mask=et["attention_mask"].to(DEV),
                output_hidden_states=True).hidden_states[L_T5].cpu()
        em = embT(et["input_ids"].to(DEV)).cpu()
        for r, li in enumerate(chunk):
            offB = eb["offset_mapping"][r].tolist()
            offT = et["offset_mapping"][r].tolist()
            idsB = eb["input_ids"][r].tolist()
            idsT = et["input_ids"][r].tolist()
            tokcacheB[li] = (idsB, offB)
            tokcacheT[li] = (idsT, offT)
            for k in by_line[li]:
                _, lo, hi, sid = sites[k]
                ixB = [i for i, (s, t) in enumerate(offB)
                       if t > s and s < hi and t > lo]
                ixT = [i for i, (s, t) in enumerate(offT)
                       if t > s and s < hi and t > lo]
                if not ixB or not ixT:
                    KK[k] = (0, 0)
                    continue
                H_B[k, 0] = hb[r, ixB[0]].numpy()
                H_B[k, 1] = hb[r, ixB[-1]].numpy()
                H_T[k, 0] = ht[r, ixT[0]].numpy()
                H_T[k, 1] = ht[r, ixT[-1]].numpy()
                E_T[k] = em[r, ixT[0]].numpy()
                KK[k] = (len(ixB), len(ixT))
        if (bs // BATCH) % 5 == 0:
            print(f"[P0-unmasked] {bs}/{len(line_ids)} lines", flush=True)
print("[P0] unmasked pass done", flush=True)

# ---- masked pass: one sequence per SITE, id-spliced
def masked_batchB(ks):
    seqs, poss = [], []
    for k in ks:
        li, lo, hi, sid = sites[k]
        idsB, offB = tokcacheB[li]
        ix = [i for i, (s, t) in enumerate(offB)
              if t > s and s < hi and t > lo]
        pre = [idsB[i] for i, (s, t) in enumerate(offB)
               if t > s and t <= lo]
        post = [idsB[i] for i, (s, t) in enumerate(offB)
                if t > s and s >= hi]
        seqs.append([tkB.cls_token_id] + pre + [MASK] * max(len(ix), 1)
                    + post + [tkB.sep_token_id])
        poss.append(1 + len(pre))
    return seqs, poss


def masked_batchT(ks):
    seqs, poss = [], []
    for k in ks:
        li, lo, hi, sid = sites[k]
        ln = lines[li]
        src = ln[:lo] + "<extra_id_0>" + ln[hi:]
        ids = tkA(src, add_special_tokens=False)["input_ids"]
        try:
            p = ids.index(SENT)
        except ValueError:
            p = 0
        seqs.append(ids)
        poss.append(p)
    return seqs, poss


def run_masked(model, seqs, poss, layer, pad_id):
    mx = max(len(s) for s in seqs)
    ids = torch.full((len(seqs), mx), pad_id, dtype=torch.long)
    att = torch.zeros((len(seqs), mx), dtype=torch.long)
    for i, s in enumerate(seqs):
        ids[i, :len(s)] = torch.tensor(s)
        att[i, :len(s)] = 1
    with torch.no_grad():
        h = model(input_ids=ids.to(DEV), attention_mask=att.to(DEV),
                  output_hidden_states=True).hidden_states[layer].cpu()
    return h[torch.arange(len(seqs)), torch.tensor(poss)]

order = [k for k in range(N) if KK[k, 0] > 0]
for bs in range(0, len(order), BATCH):
    ks = order[bs:bs + BATCH]
    sq, ps = masked_batchB(ks)
    out = run_masked(mB, sq, ps, L_BERT, tkB.pad_token_id)
    for j, k in enumerate(ks):
        M_B[k] = out[j].numpy()
    sq, ps = masked_batchT(ks)
    out = run_masked(mT, sq, ps, L_T5, tkA.pad_token_id)
    for j, k in enumerate(ks):
        M_T[k] = out[j].numpy()
    if (bs // BATCH) % 20 == 0:
        print(f"[P0-masked] {bs}/{len(order)} sites", flush=True)

np.savez_compressed(
    rf"{OUTDIR}\frame_dump_v2.npz",
    H_B=H_B, H_T=H_T, M_B=M_B, M_T=M_T, E_T=E_T, KK=KK,
    sid=np.array([s[3] for s in sites], dtype=np.int32),
    line=np.array([s[0] for s in sites], dtype=np.int32),
    lo=np.array([s[1] for s in sites], dtype=np.int32),
    hi=np.array([s[2] for s in sites], dtype=np.int32))
meta = {"n_sites": N, "skipped": int(N - len(order)),
        "layers": {"bert": L_BERT, "t5": L_T5},
        "t5_whole": n_whole,
        "env": {"transformers": transformers.__version__,
                "torch": torch.__version__}}
with open(rf"{OUTDIR}\phase0_meta.json", "w", encoding="utf-8") as f:
    json.dump(meta, f, indent=1)
print(f"[P0] COMPLETE: {len(order)}/{N} sites, "
      f"vram peak {torch.cuda.max_memory_allocated()/2**30:.1f}GB",
      flush=True)