ATOM (atom-ai): code snapshot (100fae0 + batched-train), README with measured numbers, curated checkpoints (part 3)
969c687 verified | """(1a) Byte n-gram with interpolated Witten-Bell backoff, n=1..8. | |
| Val = the exact ATOM 2048-transition val (4x512, context reset at each window start). | |
| Two training sets: 'full' (whole corpus, same as ATOM which saw val ~4.5x) and | |
| 'heldout' (corpus with every occurrence-range of the 4 val windows cut out, +-4KB). | |
| Also samples the 3 chat prompts (dialogue wrap, T=0.8, top-k 12, <=100 bytes, stop at \n\n). | |
| """ | |
| import json, math, random, sys, time | |
| from collections import defaultdict | |
| from pathlib import Path | |
| sys.path.insert(0, str(Path(__file__).resolve().parent)) | |
| from common import OFFSETS, PROMPTS, ROOT, N, val_windows_bytes | |
| OUT = Path(__file__).resolve().parent | |
| def heldout_corpus(): | |
| parts = [] | |
| for f in ("data/corpus_fr_medium.txt", "data/corpus_fr_medium_s21.txt"): | |
| raw = (ROOT / f).read_bytes() | |
| cuts = sorted((o - 4096, o + 4 * N + 4096) for p, o in OFFSETS if p == f) | |
| pos = 0 | |
| for a, b in cuts: | |
| parts.append(raw[pos:max(a, pos)]); pos = b | |
| parts.append(raw[pos:]) | |
| return b"".join(parts) | |
| class WB: | |
| def __init__(self, data: bytes, n: int): | |
| self.n = n | |
| self.c = [defaultdict(lambda: defaultdict(int)) for _ in range(n)] | |
| for i in range(len(data)): | |
| for k in range(n): | |
| if i - k < 0: break | |
| self.c[k][data[i - k:i]][data[i]] += 1 | |
| self.tot = [{h: sum(d.values()) for h, d in ck.items()} for ck in self.c] | |
| def dist(self, hist: bytes): | |
| p = [1.0 / 256] * 256 | |
| for k in range(self.n): | |
| if k > len(hist): break | |
| h = hist[len(hist) - k:] if k else b"" | |
| d = self.c[k].get(h) | |
| if not d: break | |
| t = self.tot[k][h]; u = len(d) | |
| lam = t / (t + u) | |
| q = [(1 - lam) * x for x in p] | |
| for b, cnt in d.items(): | |
| q[b] += lam * cnt / t | |
| p = q | |
| return p | |
| def ce(self, windows): | |
| tot, n = 0.0, 0 | |
| for w in windows: | |
| for i in range(len(w) - 1): | |
| p = self.dist(w[max(0, i + 1 - (self.n - 1)): i + 1]) | |
| tot += -math.log(max(p[w[i + 1]], 1e-12)); n += 1 | |
| return tot / n | |
| def sample(self, prompt: str, rng, T=0.8, k=12, max_len=100): | |
| ctx = f"Utilisateur: {prompt}\nAssistant: ".encode() | |
| out = bytearray() | |
| while len(out) < max_len: | |
| p = self.dist((ctx + out)[-(self.n - 1):] if self.n > 1 else b"") | |
| logits = [math.log(max(x, 1e-12)) / T for x in p] | |
| top = sorted(range(256), key=lambda b: -logits[b])[:k] | |
| m = max(logits[b] for b in top) | |
| ws = [math.exp(logits[b] - m) for b in top] | |
| out.append(rng.choices(top, ws)[0]) | |
| if out.endswith(b"\n\n"): break | |
| return out.decode("utf-8", "replace") | |
| def main(): | |
| wins = val_windows_bytes() | |
| full = (ROOT / "data/corpus_fr_medium.txt").read_bytes() + (ROOT / "data/corpus_fr_medium_s21.txt").read_bytes() | |
| held = heldout_corpus() | |
| res = {"val": "ATOM 2048-transition val (eval_windows.py), nats/byte", "full_bytes": len(full), "heldout_bytes": len(held), "rows": []} | |
| for name, data in (("full", full), ("heldout", held)): | |
| for n in (1, 2, 3, 4, 5, 6, 8): | |
| t0 = time.time(); m = WB(data, n); ce = m.ce(wins) | |
| row = {"train": name, "n": n, "val_ce": round(ce, 4), "sec": round(time.time() - t0, 1)} | |
| if n in (3, 5, 8) and name == "heldout": | |
| rng = random.Random(0) | |
| row["samples"] = {p: [m.sample(p, rng) for _ in range(3)] for p in PROMPTS} | |
| res["rows"].append(row); print(json.dumps(row, ensure_ascii=False), flush=True) | |
| (OUT / "ngram_results.json").write_text(json.dumps(res, indent=1, ensure_ascii=False)) | |
| if __name__ == "__main__": | |
| main() | |