| """Adding words after training, without retraining (the main test). |
| |
| A frozen-table model (B, C, B2) can take a NEW vocabulary row at any time: the new word gets a frozen 512-bit code made by |
| the same transform as the table (Gemma 3 270M embedding space -> centre -> the saved 512 principal directions -> sign), |
| here from the mean of the Gemma embeddings of the pieces the word is currently split into, and a frequency prior at the |
| 75th percentile of the existing rows (the SHADOW-50M v1.1 rule). Nothing is trained. A learned table (A) has no way to |
| add a row without training, so A keeps reading the word piece by piece. |
| |
| Two measurements, each for: A pieces | arm pieces | arm with the new token |
| 1 in-context recall: "My name is <w>. ... What is my name? My name is" -> is <w> produced? (greedy; pieces must all |
| be right; the new token must be the single next token) |
| 2 natural text: words that Gemma splits into 2+ pieces, at their real occurrences in held-out FineWeb-Edu text: |
| nats per word occurrence (sum of piece losses) vs nats for the single new token in the same context |
| |
| python new_words.py --arm B --ckpt finals/B/step_07820.pt [--tag B] [--ref-ckpt finals/A/step_07820.pt] |
| """ |
| import argparse, json, pathlib, random, re, time |
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
|
|
| SEQ = 1024 |
|
|
|
|
| def gemma_embedding(): |
| from safetensors import safe_open |
| from huggingface_hub import hf_hub_download |
| with safe_open(hf_hub_download("google/gemma-3-270m", "model.safetensors"), "pt") as f: |
| return f.get_tensor([k for k in f.keys() if k.endswith("embed_tokens.weight")][0]).float()[:262144] |
|
|
|
|
| def load(arm, ckpt, counts): |
| import train |
| m = train.Model(arm, counts).cuda().eval() |
| st = torch.load(ckpt, map_location="cuda", weights_only=False) |
| m.load_state_dict(st["model"]) |
| return m, st.get("tag", arm) |
|
|
|
|
| def make_names(tok, n, seed=0): |
| """made-up names that the tokenizer splits into 2-4 pieces (fresh: never in any training text)""" |
| rng = random.Random(seed); syl = ["ka", "zor", "mel", "vit", "ran", "dul", "pex", "sho", "lin", "qua", "ber", "tis", |
| "mon", "gra", "fel", "yun", "dra", "vok", "ser", "nim", "tal", "rux", "jeb", "wen"] |
| out = set() |
| while len(out) < n: |
| w = "".join(rng.choice(syl) for _ in range(rng.randint(2, 3))).capitalize() |
| k = len(tok(" " + w, add_special_tokens=False).input_ids) |
| if 2 <= k <= 4: |
| out.add(w) |
| return sorted(out) |
|
|
|
|
| def corpus_words(tok, val, n_words, min_count): |
| """real English words that the tokenizer splits into 2+ pieces, by frequency in held-out text""" |
| text = tok.decode(val[:3_000_000].tolist()) |
| cnt = {} |
| for w in re.findall(r"(?<= )[a-z]{6,}", text): |
| cnt[w] = cnt.get(w, 0) + 1 |
| cands = [w for w, c in sorted(cnt.items(), key=lambda x: -x[1]) if c >= min_count] |
| out = [] |
| for w in cands: |
| if len(tok(" " + w, add_special_tokens=False).input_ids) >= 2: |
| out.append(w) |
| if len(out) >= n_words: |
| break |
| return out |
|
|
|
|
| @torch.no_grad() |
| def add_rows(m, E, piece_lists, prior_pct=75): |
| """append one frozen code + prior per new word; returns the new ids""" |
| V = m.table.shape[0] |
| vecs = torch.stack([E[p].mean(0) for p in piece_lists]).cuda() |
| codes = torch.sign((vecs - m.table_mu) @ m.table_W); codes[codes == 0] = 1 |
| prior = torch.quantile(m.seal.detach().float(), prior_pct / 100) |
| m.table = torch.cat([m.table, codes]) |
| seal = torch.cat([m.seal.detach(), prior.repeat(len(piece_lists))]) |
| if isinstance(m.seal, torch.nn.Parameter): |
| del m.seal |
| m.seal = seal |
| return list(range(V, V + len(piece_lists))) |
|
|
|
|
| @torch.no_grad() |
| def logp_next(m, ids): |
| x = torch.tensor([ids[-SEQ:]], device="cuda") |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| return F.log_softmax(m.logits(m.hidden(x))[0, -1].float(), -1) |
|
|
|
|
| @torch.no_grad() |
| def greedy(m, ids, k): |
| out = [] |
| for _ in range(k): |
| nxt = int(logp_next(m, ids + out).argmax()); out.append(nxt) |
| return out |
|
|
|
|
| RECALL = ["My name is {w}. I live in a small town near the sea. What is my name? My name is", |
| "The new ship was called the {w}. Everyone in the harbour admired the", |
| "Our team leader is {w}. Tomorrow the report will be presented by", |
| "The rare mineral {w} was found in the cave. Scientists are now studying"] |
|
|
|
|
| def recall_test(m, tok, sep, words, new_ids=None): |
| hits = 0; tot = 0 |
| for i, w in enumerate(words): |
| pieces = tok(" " + w, add_special_tokens=False).input_ids |
| for t in RECALL: |
| pre, post = t.split("{w}") |
| a = tok(pre.rstrip(), add_special_tokens=False).input_ids if pre.strip() else [] |
| b = tok(post, add_special_tokens=False).input_ids |
| word = [new_ids[i]] if new_ids else pieces |
| ids = [sep] + a + word + b |
| got = greedy(m, ids, len(word)) |
| hits += int(got == word); tot += 1 |
| return hits / tot |
|
|
|
|
| def text_test(m, val, words_pieces, new_ids=None, max_occ=40, ctx=256): |
| """nats per word occurrence at real occurrences in held-out text""" |
| arr = np.asarray(val[:20_000_000]); res = [] |
| for i, p in enumerate(words_pieces): |
| p = np.array(p); k = len(p) |
| hits = np.where(np.all(np.stack([arr[j:len(arr) - k + j + 1] == p[j] for j in range(k)]), 0))[0] |
| hits = [h for h in hits if h >= ctx][:max_occ] |
| for h in hits: |
| context = arr[h - ctx:h].tolist() |
| if new_ids is None: |
| lp = 0.0; seq = context[:] |
| for t in p.tolist(): |
| lp += float(logp_next(m, seq)[t]); seq.append(t) |
| else: |
| lp = float(logp_next(m, context)[new_ids[i]]) |
| res.append(-lp) |
| return float(np.mean(res)), len(res) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--arm", required=True, choices=["B", "C"]); ap.add_argument("--ckpt", required=True) |
| ap.add_argument("--tag", default=""); ap.add_argument("--ref-ckpt", default="", help="arm A checkpoint (pieces only)") |
| ap.add_argument("--data", default="data"); ap.add_argument("--names", type=int, default=60) |
| ap.add_argument("--words", type=int, default=100) |
| a = ap.parse_args() |
| from transformers import AutoTokenizer |
| tok = AutoTokenizer.from_pretrained("google/gemma-3-270m"); sep = tok.eos_token_id |
| counts = np.load(pathlib.Path(a.data) / "unigram_counts.npy") |
| val = np.load(pathlib.Path(a.data) / "val.npy", mmap_mode="r") |
| names = make_names(tok, a.names) |
| words = corpus_words(tok, val, a.words, min_count=5) |
| wp = [tok(" " + w, add_special_tokens=False).input_ids for w in words] |
| np_ = [tok(" " + w, add_special_tokens=False).input_ids for w in names] |
| R = {"names": names, "words": words, "date": time.strftime("%F %H:%M")} |
| if a.ref_ckpt: |
| A, _ = load("A", a.ref_ckpt, counts) |
| R["A_pieces"] = {"recall": recall_test(A, tok, sep, names), "text_nats_per_word": text_test(A, val, wp)} |
| print("A pieces:", R["A_pieces"], flush=True); del A; torch.cuda.empty_cache() |
| m, tag = load(a.arm, a.ckpt, counts); tag = a.tag or tag |
| R[f"{tag}_pieces"] = {"recall": recall_test(m, tok, sep, names), "text_nats_per_word": text_test(m, val, wp)} |
| print(f"{tag} pieces:", R[f"{tag}_pieces"], flush=True) |
| E = gemma_embedding() |
| ids_n = add_rows(m, E, np_); ids_w = add_rows(m, E, wp) |
| R[f"{tag}_new_tokens"] = {"recall": recall_test(m, tok, sep, names, ids_n), |
| "text_nats_per_word": text_test(m, val, wp, ids_w), "rows_added": len(ids_n) + len(ids_w)} |
| print(f"{tag} new tokens:", R[f"{tag}_new_tokens"], flush=True) |
| pathlib.Path("reports").mkdir(exist_ok=True) |
| json.dump(R, open(f"reports/new_words_{tag}.json", "w"), indent=1) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
