ATOM (atom-ai): code snapshot (100fae0 + batched-train), README with measured numbers, curated checkpoints
9f1076a verified | """Shared measurement helpers for docs/DECODE_FIX_AND_BPTT.md (reuse the root_cause tools).""" | |
| import random, re, sys | |
| from pathlib import Path | |
| import numpy as np | |
| import torch, torch.nn.functional as F | |
| HERE = Path(__file__).resolve().parent | |
| ROOT = HERE.parents[2] | |
| RC = ROOT / "docs/artifacts/root_cause" | |
| for p in (str(ROOT), str(ROOT / "tools"), str(RC), str(ROOT / "docs/artifacts/field_leak_spectrum")): | |
| if p not in sys.path: | |
| sys.path.insert(0, p) | |
| from chat_atom_native import load_model, generate_reply # noqa: E402 | |
| from common import PROMPTS, val_windows_bytes # noqa: E402 (root_cause) | |
| import eval_windows as EW # noqa: E402 (the 2048-transition val) | |
| SP = set(b" ,.-'?!:;") | |
| def val2048(m): | |
| """normal + field-alone byte CE on the exact eval_windows.py val (4 x 512).""" | |
| normal, alone = [], [] | |
| for w in EW.windows(): | |
| normal += EW.run(m, w["pairs"], False); alone += EW.run(m, w["pairs"], True) | |
| return {"normal_byte_ce": float(np.mean(normal)), "field_alone_byte_ce": float(np.mean(alone))} | |
| def run_logp(m, pk): | |
| m.eval(); m.reset_state(reset_atomizer=False) | |
| return [torch.log_softmax(m.forward_packet(p)["surface"]["byte_logits"][0], -1) for p in pk[:-1]] | |
| def perturbation_kl(m, lags=(1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64), per_window=6, seed=0): | |
| """Same protocol as root_cause/diag_atom_more.py: flip one ASCII letter at p in a val window, | |
| KL(next-byte | orig || perturbed) L bytes later; normal and field-alone (last_atom off).""" | |
| wins = val_windows_bytes(); rng = random.Random(seed); sens = {} | |
| for mode in ("normal", "field_alone"): | |
| m.surface.last_atom_force_off = mode == "field_alone" | |
| kl = {L: [] for L in lags} | |
| for w in wins: | |
| cands = [p for p in range(64, 440) if 97 <= w[p] <= 122] | |
| for p in rng.sample(cands, per_window): | |
| wp = bytearray(w); wp[p] = rng.choice([c for c in range(97, 123) if c != w[p]]) | |
| seg, segp = bytes(w[: p + 66]), bytes(wp[: p + 66]) | |
| a = run_logp(m, m.atomizer.encode_bytes(seg, reset=True)) | |
| b = run_logp(m, m.atomizer.encode_bytes(segp, reset=True)) | |
| for L in lags: | |
| i = p + L - 1 | |
| kl[L].append(float((a[i].exp() * (a[i] - b[i])).sum())) | |
| sens[mode] = {int(L): float(np.mean(v)) for L, v in kl.items()} | |
| m.surface.last_atom_force_off = False | |
| return sens | |
| def grad_by_lag(m, K=None, lags=(0, 1, 2, 4, 8, 16, 32, 64), T=200): | |
| """||d loss_T / d atom_{T-L}|| (r,phi,omega,E,kappa) in train mode on val window 1, | |
| root_cause protocol; K = field_bptt to set (None = keep the model's).""" | |
| old_k = m.field_bptt | |
| if K is not None: | |
| m.set_field_bptt(K) | |
| m.train(); drop = m.surface.last_atom_dropout; m.surface.last_atom_dropout = 0.0 | |
| stash = []; orig = m.compiler.forward | |
| def hooked(features, atom_count=0): | |
| atom, op, conf = orig(features, atom_count) | |
| for n in ("r", "phi", "omega", "E", "kappa"): | |
| getattr(atom, n).retain_grad() | |
| stash.append(atom); return atom, op, conf | |
| m.compiler.forward = hooked | |
| try: | |
| wins = val_windows_bytes() | |
| m.reset_state(reset_atomizer=True); pk = m.atomizer.encode_bytes(wins[1][: T + 1], reset=True) | |
| for q in pk[: T - 1]: | |
| m.forward_packet(q) | |
| out = m.forward_packet(pk[T - 1]) | |
| loss = F.cross_entropy(out["surface"]["byte_logits"], torch.tensor([pk[T].payload[0]])) | |
| loss.backward() | |
| finally: | |
| m.compiler.forward = orig | |
| m.surface.last_atom_dropout = drop | |
| m.zero_grad(set_to_none=True) | |
| m.set_field_bptt(old_k) | |
| m.eval() | |
| return {int(L): sum(float(getattr(stash[T - 1 - L], n).grad.norm()) if getattr(stash[T - 1 - L], n).grad is not None else 0.0 | |
| for n in ("r", "phi", "omega", "E", "kappa")) for L in lags} | |
| def chat_samples(m, seeds=range(8), gate=True): | |
| """diag_decode.py protocol: generate_packets(prompt) seeds 0..7, T .8, top-k 12, 96 pkts/100 bytes.""" | |
| out = {} | |
| for p in PROMPTS: | |
| out[p] = [] | |
| for s in seeds: | |
| torch.manual_seed(s) | |
| raw = m.generate_packets(p, max_packets=96, temperature=0.8, top_k=12, deterministic=False, | |
| max_length=100, prefer_printable=True, reset=True, merge_enabled=False, | |
| speech_gate=gate, anti_repeat=gate) | |
| out[p].append(raw.decode("utf-8", "replace")) | |
| return out | |
| def doubled_rate(samples): | |
| txt = "".join(s for v in samples.values() for s in v) | |
| return {"chars": len(txt), "punct": sum(1 for c in txt if c in ",.-'?!:;"), | |
| "doubled_punct_pairs": len(re.findall(r"([,.\-'?!:;])\1", txt)), | |
| "space_runs_ge3": len(re.findall(r" +", txt))} | |
| def probe_chats(m, seeds=(1, 2, 3)): | |
| """tools/chat_atom_native.py --no-role-prime --max-packets 96 --max-length 100 -T .8 -k 12; | |
| unseeded = one draw per prompt (as the probe runs), plus chat_seeds_probe seeds.""" | |
| rows = [] | |
| for p in PROMPTS: | |
| rows.append({"seed": None, "prompt": p, "resp": generate_reply(m, p, max_packets=96, temperature=0.8, top_k=12, | |
| max_length=100, role_prime=False)}) | |
| for s in seeds: | |
| for p in PROMPTS: | |
| torch.manual_seed(s) | |
| rows.append({"seed": s, "prompt": p, "resp": generate_reply(m, p, max_packets=96, temperature=0.8, top_k=12, | |
| max_length=100, role_prime=False)}) | |
| return rows | |
| _JUDGE = {} | |
| def judge_nll(samples): | |
| """GRU f100 judge (val .10) NLL/byte on the first reply line (root_cause/judge.py protocol).""" | |
| from judge import gru_nll | |
| from gru_baseline import GRULM | |
| if "m" not in _JUDGE: | |
| j = GRULM(512, 2); j.load_state_dict(torch.load(RC / "gru_f100.pt")); j.eval(); _JUDGE["m"] = j | |
| vals = [] | |
| for p, vs in samples.items(): | |
| for v in vs: | |
| first = v.split("\n")[0] + ("\n" if "\n" in v else "") | |
| x = gru_nll(_JUDGE["m"], p, first) | |
| if x is not None: | |
| vals.append(x) | |
| return float(np.mean(vals)) if vals else None | |
| def prompt_identical_frac(samples): | |
| vals = list(samples.values()) | |
| return sum(len({vals[k][i] for k in range(len(vals))}) == 1 for i in range(len(vals[0]))) / len(vals[0]) | |