| """Add words from EXAMPLES, no gradient: a new word's frozen code = sign of the average readout direction the |
| trained model produces just before the word, over up to K example sentences from the TRAINING text; prior = log of the |
| word's training frequency. Evaluated on HELD-OUT text: total nats per character over whole passages (words merged where |
| they occur) vs no new words. Also reports tokens saved.""" |
| import sys, numpy as np, torch, torch.nn.functional as F |
| import new_words as nw |
| from transformers import AutoTokenizer |
| ARM, CKPT = (sys.argv[1], sys.argv[2]) if len(sys.argv) > 2 else ("B", "finals/B/step_07820.pt") |
| NW, K, CTX = 200, 50, 128 |
| tok = AutoTokenizer.from_pretrained("google/gemma-3-270m") |
| counts = np.load("data/unigram_counts.npy"); val = np.load("data/val.npy", mmap_mode="r") |
| tr = np.load("data/train_000.npy", mmap_mode="r")[:100_000_000] |
| words = nw.corpus_words(tok, val, NW, 5) |
| wp = [tok(" " + w, add_special_tokens=False).input_ids for w in words] |
| m, tag = nw.load(ARM, CKPT, counts) |
| T0, S0 = m.table.clone(), m.seal.detach().clone(); V = T0.shape[0] |
| def occ(p, arr, limit): |
| k = len(p); h = np.flatnonzero(arr[:len(arr) - k + 1] == p[0]) |
| for j in range(1, k): h = h[arr[h + j] == p[j]] |
| return h, len(h) |
| codes, priors = [], [] |
| with torch.no_grad(): |
| for p in wp: |
| h, n = occ(p, tr, K) |
| h = [x for x in h if x >= CTX][:K] |
| if not h: |
| v = torch.zeros(512, device="cuda") |
| else: |
| ctx = torch.tensor(np.stack([tr[x - CTX:x] for x in h]).astype(np.int64), device="cuda") |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| z = m.out(m.hidden(ctx)[:, -1]).float() |
| v = z.mean(0) |
| c = torch.sign(v); c[c == 0] = 1; codes.append(c) |
| priors.append(np.log((n * 20.5 + 1) / (counts.sum() + 262144))) |
| codes = torch.stack(codes); priors = torch.tensor(priors, dtype=torch.float32, device="cuda") |
| NSEQ, L = 600, 512 |
| seqs = [np.asarray(val[i * L:(i + 1) * L]).tolist() for i in range(NSEQ)] |
| chars = sum(len(tok.decode(s)) for s in seqs) |
| order = sorted(range(len(wp)), key=lambda i: -len(wp[i])) |
| def merge(s, ids): |
| out, i = [], 0 |
| while i < len(s): |
| for j in order: |
| p = wp[j] |
| if s[i:i + len(p)] == p: out.append(ids[j]); i += len(p); break |
| else: out.append(s[i]); i += 1 |
| return out |
| @torch.no_grad() |
| def total(seq_list): |
| t = 0.0 |
| for s in seq_list: |
| x = torch.tensor([[1] + s], device="cuda") |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| lg = m.logits(m.hidden(x[:, :-1])).float() |
| t += float(F.cross_entropy(lg[0], x[0, 1:], reduction="sum")) |
| return t |
| m.table, m.seal = T0, S0 |
| base = total(seqs) |
| ids = list(range(V, V + len(wp))) |
| merged = [merge(s, ids) for s in seqs] |
| saved = sum(map(len, seqs)) - sum(map(len, merged)) |
| for name, cd, pr in (("examples code + frequency prior", codes, priors), |
| ("piece-mean code + frequency prior", None, priors)): |
| if cd is None: |
| E = nw.gemma_embedding(); vv = torch.stack([E[p].mean(0) for p in wp]).cuda() |
| cd = torch.sign((vv - m.table_mu) @ m.table_W); cd[cd == 0] = 1 |
| m.table = torch.cat([T0, cd]); m.seal = torch.cat([S0, pr]) |
| t = total(merged) |
| print(f"{tag} | {name}: {t / chars:.4f} vs {base / chars:.4f} nats/char ({(t - base) / chars:+.4f}, {(t - base) / base * 100:+.2f}%) | " |
| f"{saved} tokens saved of {sum(map(len, seqs))}", flush=True) |
|
|
