import os import math import json from collections import Counter import torch import torch.nn as nn from torch.nn import functional as F class AttoLM(nn.Module): """ Neural n-gram language model designed for extreme intelligence density. Architecture: embed last K chars → flatten → linear mix → tanh → project via weight-tied embedding → logits. Every single parameter directly participates in prediction. Total params = embd_dim × (vocab_size + context_len × embd_dim) """ def __init__(self, vocab_size, embd_dim, context_len): super().__init__() self.vocab_size = vocab_size self.embd_dim = embd_dim self.context_len = context_len self.embedding = nn.Embedding(vocab_size, embd_dim) self.mix = nn.Linear(context_len * embd_dim, embd_dim, bias=False) # output projection is weight-tied: logits = h @ embedding.weight.T # so lm_head costs 0 extra parameters self._init_weights() def _init_weights(self): nn.init.normal_(self.embedding.weight, mean=0.0, std=0.5) nn.init.normal_(self.mix.weight, mean=0.0, std=0.1) def forward(self, idx): # idx: (batch, context_len) — indices of last K characters x = self.embedding(idx) # (B, K, D) x = x.reshape(x.size(0), -1) # (B, K*D) x = self.mix(x) # (B, D) x = torch.tanh(x) # free nonlinearity logits = x @ self.embedding.weight.T # (B, vocab_size) — weight tying return logits def count_params(self): return sum(p.numel() for p in self.parameters()) # ────────────────────────────────────────────────────────────────────── # Exact parameter budgets: # params = embd_dim × (vocab_size + context_len × embd_dim) # # atto-64: embd=2, vocab=26, ctx=3 → 2×(26 + 3×2) = 2×32 = 64 # atto-128: embd=2, vocab=50, ctx=7 → 2×(50 + 7×2) = 2×64 = 128 # atto-256: embd=4, vocab=32, ctx=8 → 4×(32 + 8×4) = 4×64 = 256 # atto-512: embd=4, vocab=64, ctx=16 → 4×(64 + 16×4) = 4×128 = 512 # atto-1024: embd=8, vocab=64, ctx=8 → 8×(64 + 8×8) = 8×128 = 1024 # ────────────────────────────────────────────────────────────────────── CONFIGS = { "atto-64": dict(embd_dim=2, vocab_size=26, context_len=3), "atto-128": dict(embd_dim=2, vocab_size=50, context_len=7), "atto-256": dict(embd_dim=4, vocab_size=32, context_len=8), "atto-512": dict(embd_dim=4, vocab_size=64, context_len=16), "atto-1024": dict(embd_dim=8, vocab_size=64, context_len=8), } def build_vocab(text, vocab_size): counts = Counter(text) common = [ch for ch, _ in counts.most_common(vocab_size)] chars = sorted(common) stoi = {ch: i for i, ch in enumerate(chars)} itos = {i: ch for i, ch in enumerate(chars)} return chars, stoi, itos def encode_text(text, stoi): return [stoi[c] for c in text if c in stoi] def make_windows(data, ctx): """Pre-build all (input, target) pairs as a single tensor. No Python loops at train time.""" n = len(data) - ctx # indices [0..n-1], each window is data[i:i+ctx] → data[i+ctx] idx = torch.arange(n).unsqueeze(1) + torch.arange(ctx + 1).unsqueeze(0) # (n, ctx+1) windows = data[idx] # (n, ctx+1) return windows[:, :ctx], windows[:, ctx] # x: (n, ctx), y: (n,) def train_model(name, cfg, text): print(f"\n{'='*60}") print(f" Training {name}") print(f" embd_dim={cfg['embd_dim']}, vocab={cfg['vocab_size']}, ctx={cfg['context_len']}") print(f"{'='*60}") chars, stoi, itos = build_vocab(text, cfg["vocab_size"]) data = torch.tensor(encode_text(text, stoi), dtype=torch.long) n = int(0.9 * len(data)) train_data = data[:n] val_data = data[n:] ctx = cfg["context_len"] # pre-build all sliding windows — makes batching a single index op train_x, train_y = make_windows(train_data, ctx) val_x, val_y = make_windows(val_data, ctx) model = AttoLM(cfg["vocab_size"], cfg["embd_dim"], cfg["context_len"]) p_count = model.count_params() print(f" Parameters: {p_count}") # scale iters to model capacity — tiny models converge fast max_iters = min(5000 + p_count * 10, 20000) batch_size = 512 lr_max = 3e-2 lr_min = 1e-4 eval_every = max(max_iters // 10, 1) optimizer = torch.optim.AdamW(model.parameters(), lr=lr_max, weight_decay=0.01) best_val_loss = float("inf") best_state = None import time t0 = time.time() for i in range(max_iters): lr = lr_min + 0.5 * (lr_max - lr_min) * (1 + math.cos(math.pi * i / max_iters)) for pg in optimizer.param_groups: pg["lr"] = lr # fast batch: just index into pre-built windows ix = torch.randint(0, len(train_x), (batch_size,)) logits = model(train_x[ix]) loss = F.cross_entropy(logits, train_y[ix]) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() if i % eval_every == 0: model.eval() with torch.no_grad(): vix = torch.randint(0, len(val_x), (min(2048, len(val_x)),)) val_loss = F.cross_entropy(model(val_x[vix]), val_y[vix]).item() model.train() print(f" iter {i:6d} | train {loss.item():.4f} | val {val_loss:.4f} | lr {lr:.2e}") if val_loss < best_val_loss: best_val_loss = val_loss best_state = {k: v.clone() for k, v in model.state_dict().items()} elapsed = time.time() - t0 print(f" Trained in {elapsed:.1f}s ({max_iters} iters)") if best_state is not None: model.load_state_dict(best_state) # quick sample model.eval() with torch.no_grad(): seed = [stoi.get(c, 0) for c in " the"][-ctx:] while len(seed) < ctx: seed = [0] + seed generated = list(seed) for _ in range(120): inp = torch.tensor(generated[-ctx:], dtype=torch.long).unsqueeze(0) logits = model(inp) probs = F.softmax(logits / 0.8, dim=-1) nxt = torch.multinomial(probs, 1).item() generated.append(nxt) sample = "".join(itos.get(c, "?") for c in generated[ctx:]) print(f" Sample: {sample}") # export os.makedirs("models", exist_ok=True) model_json = { "name": name, "architecture": "atto-neural-ngram", "params": p_count, "config": { "vocab_size": cfg["vocab_size"], "embd_dim": cfg["embd_dim"], "context_len": cfg["context_len"], }, "vocab": {str(i): ch for i, ch in itos.items()}, "weights": { "embedding": model.embedding.weight.detach().tolist(), "mix": model.mix.weight.detach().tolist(), }, } path = f"models/{name}.json" with open(path, "w") as f: json.dump(model_json, f) print(f" Saved → {path} ({os.path.getsize(path)} bytes)") if __name__ == "__main__": with open("hamlet.txt", "r", encoding="utf-8") as f: text = f.read() for name, cfg in CONFIGS.items(): train_model(name, cfg, text)