compare_baselines_tinystories.py

import os
import sys
import time
import math
import random
import subprocess
from collections import defaultdict, Counter

import torch
import torch.nn as nn
from torch.nn import functional as F

# Set seeds
torch.manual_seed(42)
random.seed(42)

CORPUS_PATH = "benchmarks/data/tinystories_sample.txt"
DZETA_MODEL_PATH = "benchmarks/models/dzeta_tinystories_v3.dzeta.bin"
DZETA_BIN = "./dzeta_inspect_model.exe"

PROMPTS = [
    "Once upon a time",
    "Lily had a little",
    "The little robot",
    "One day, Tim",
    "A funny monkey"
]

def load_corpus():
    with open(CORPUS_PATH, "r", encoding="utf-8") as f:
        text = f.read()
    return text

def tokenize(text):
    import re
    return re.findall(r"\w+|[^\w\s]", text.lower())

print("Loading corpus...", flush=True)
corpus_text = load_corpus()
tokens = tokenize(corpus_text)
print(f"Total tokens in corpus: {len(tokens)}", flush=True)
vocab = sorted(list(set(tokens)))
vocab_size = len(vocab)
print(f"Unique vocab size: {vocab_size}", flush=True)

w2i = {w: i for i, w in enumerate(vocab)}
i2w = {i: w for i, w in enumerate(vocab)}

# -------------------------------------------------------------
# 1. MARKOV CHAIN BASELINE (Trigram with Backoff)
# -------------------------------------------------------------
class MarkovTrigram:
    def __init__(self):
        self.trigrams = defaultdict(Counter)
        self.bigrams = defaultdict(Counter)
        self.unigrams = Counter()

    def train(self, tokens):
        print("[Markov] Training trigram model...", flush=True)
        t0 = time.time()
        for i in range(len(tokens) - 2):
            w1, w2, w3 = tokens[i], tokens[i+1], tokens[i+2]
            self.trigrams[(w1, w2)][w3] += 1
            self.bigrams[w1][w2] += 1
            self.unigrams[w1] += 1
        self.unigrams[tokens[-2]] += 1
        self.unigrams[tokens[-1]] += 1
        elapsed = time.time() - t0
        print(f"[Markov] Trained in {elapsed:.4f}s", flush=True)
        return elapsed

    def generate(self, prompt, length=18, temperature=0.7):
        prompt_tokens = tokenize(prompt)
        curr = list(prompt_tokens)
        for _ in range(length):
            candidates = None
            if len(curr) >= 2:
                key = (curr[-2], curr[-1])
                if key in self.trigrams:
                    candidates = self.trigrams[key]
            if not candidates and len(curr) >= 1:
                key = curr[-1]
                if key in self.bigrams:
                    candidates = self.bigrams[key]
            if not candidates:
                candidates = self.unigrams

            words, counts = zip(*candidates.items())
            probs = torch.tensor(counts, dtype=torch.float)
            if temperature > 0:
                probs = probs ** (1.0 / temperature)
                probs = probs / probs.sum()
                next_word = words[torch.multinomial(probs, 1).item()]
            else:
                next_word = words[torch.argmax(probs).item()]
            curr.append(next_word)
        return " ".join(curr[len(prompt_tokens):])

# -------------------------------------------------------------
# 2. nanoGPT TRANSFORMER BASELINE
# -------------------------------------------------------------
class Head(nn.Module):
    def __init__(self, head_size, n_embd, block_size):
        super().__init__()
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))

    def forward(self, x):
        B, T, C = x.shape
        k = self.key(x)
        q = self.query(x)
        wei = q @ k.transpose(-2, -1) * (C ** -0.5)
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
        wei = F.softmax(wei, dim=-1)
        v = self.value(x)
        out = wei @ v
        return out

class MultiHeadAttention(nn.Module):
    def __init__(self, num_heads, head_size, n_embd, block_size):
        super().__init__()
        self.heads = nn.ModuleList([Head(head_size, n_embd, block_size) for _ in range(num_heads)])
        self.proj = nn.Linear(head_size * num_heads, n_embd)

    def forward(self, x):
        out = torch.cat([h(x) for h in self.heads], dim=-1)
        out = self.proj(out)
        return out

class FeedForward(nn.Module):
    def __init__(self, n_embd):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),
            nn.GELU(),
            nn.Linear(4 * n_embd, n_embd),
        )

    def forward(self, x):
        return self.net(x)

class Block(nn.Module):
    def __init__(self, n_embd, n_head, block_size):
        super().__init__()
        head_size = n_embd // n_head
        self.sa = MultiHeadAttention(n_head, head_size, n_embd, block_size)
        self.ffwd = FeedForward(n_embd)
        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)

    def forward(self, x):
        x = x + self.sa(self.ln1(x))
        x = x + self.ffwd(self.ln2(x))
        return x

class NanoGPT(nn.Module):
    def __init__(self, vocab_size, n_embd=128, block_size=64, n_layer=4, n_head=4):
        super().__init__()
        self.block_size = block_size
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
        self.position_embedding_table = nn.Embedding(block_size, n_embd)
        self.blocks = nn.Sequential(*[Block(n_embd, n_head, block_size) for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        tok_emb = self.token_embedding_table(idx)
        pos_emb = self.position_embedding_table(torch.arange(T, device=idx.device))
        x = tok_emb + pos_emb
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)

        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            logits = logits.view(B*T, C)
            targets = targets.view(B*T)
            loss = F.cross_entropy(logits, targets)

        return logits, loss

    def generate(self, idx, max_new_tokens, temperature=0.8):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.block_size:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :] / temperature
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

def train_nanogpt(tokens, vocab_size, epochs=3, batch_size=32, block_size=64):
    print("[nanoGPT] Initializing model...", flush=True)
    model = NanoGPT(vocab_size=vocab_size, n_embd=128, block_size=block_size, n_layer=4, n_head=4)
    param_count = sum(p.numel() for p in model.parameters())
    print(f"[nanoGPT] Total parameters: {param_count:,}", flush=True)

    token_ids = torch.tensor([w2i[w] for w in tokens], dtype=torch.long)
    n_samples = len(token_ids) - block_size - 1

    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-2)

    steps_per_epoch = n_samples // (batch_size * block_size)
    total_steps = steps_per_epoch * epochs
    print(f"[nanoGPT] Training {epochs} epochs ({total_steps} steps, batch={batch_size}, block={block_size})...", flush=True)

    t0 = time.time()
    model.train()
    step = 0
    final_loss = 0.0
    for epoch in range(epochs):
        perm = torch.randperm(n_samples)
        for i in range(0, n_samples - batch_size, batch_size):
            batch_indices = perm[i:i+batch_size]
            x = torch.stack([token_ids[idx:idx+block_size] for idx in batch_indices])
            y = torch.stack([token_ids[idx+1:idx+block_size+1] for idx in batch_indices])

            logits, loss = model(x, y)
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            optimizer.step()

            step += 1
            final_loss = loss.item()
            if step % 50 == 0 or step == total_steps:
                print(f"[nanoGPT] step {step}/{total_steps} (epoch {epoch+1}/{epochs}) loss={final_loss:.4f}", flush=True)
            if step >= total_steps:
                break
        if step >= total_steps:
            break

    elapsed = time.time() - t0
    print(f"[nanoGPT] Training completed in {elapsed:.2f}s, final loss={final_loss:.4f}", flush=True)
    return model, elapsed, param_count, final_loss

# -------------------------------------------------------------
# 3. DZETA AGI BASELINE
# -------------------------------------------------------------
def run_dzeta(prompts):
    print("[DZETA] Running dzeta_inspect_model on prompts...", flush=True)
    cmd = [DZETA_BIN, "--model", DZETA_MODEL_PATH]
    for p in prompts:
        cmd.extend(["--prompt", p])
    res = subprocess.run(cmd, capture_output=True, text=True, check=True)
    outputs = {}
    current_prompt = None
    for line in res.stdout.splitlines():
        if line.startswith("prompt="):
            current_prompt = line[len("prompt="):].strip()
        elif line.startswith("prompt_output=") and current_prompt:
            outputs[current_prompt] = line[len("prompt_output="):].strip()
            current_prompt = None
    return outputs

# -------------------------------------------------------------
# MAIN BENCHMARK RUN
# -------------------------------------------------------------
def main():
    # 1. Markov
    markov = MarkovTrigram()
    markov_time = markov.train(tokens)

    # 2. nanoGPT
    nanogpt, gpt_time, gpt_params, gpt_final_loss = train_nanogpt(tokens, vocab_size, epochs=3, batch_size=32, block_size=64)

    # 3. DZETA
    dzeta_outputs = run_dzeta(PROMPTS)

    print("\n" + "="*80)
    print("EMPIRICAL SHOWDOWN: DZETA AGI vs MARKOV vs nanoGPT (TinyStories 1000, 3 epochs)")
    print("="*80)

    results = []
    for p in PROMPTS:
        # Markov generation
        m_out = markov.generate(p, length=18, temperature=0.7)

        # nanoGPT generation
        p_tokens = tokenize(p)
        p_ids = [w2i.get(w, 0) for w in p_tokens]
        if not p_ids:
            p_ids = [0]
        inp = torch.tensor([p_ids], dtype=torch.long)
        nanogpt.eval()
        with torch.no_grad():
            out_ids = nanogpt.generate(inp, max_new_tokens=18, temperature=0.8)[0].tolist()
        gpt_words = [i2w.get(idx, "<unk>") for idx in out_ids[len(p_ids):]]
        gpt_out = " ".join(gpt_words)

        # DZETA generation
        dz_out = dzeta_outputs.get(p, "N/A")

        results.append({
            "prompt": p,
            "markov": m_out,
            "nanogpt": gpt_out,
            "dzeta": dz_out
        })

    # Print table
    for r in results:
        print(f"\n--- PROMPT: \"{r['prompt']}\" ---")
        print(f"[Markov] : {r['markov']}")
        print(f"[nanoGPT]: {r['nanogpt']}")
        print(f"[DZETA]  : {r['dzeta']}")

    # Overlap metric (Attractor Collapse)
    def compute_overlap(outputs):
        import itertools
        def to_words(text):
            return set(tokenize(text))
        sets = [to_words(o) for o in outputs]
        scores = []
        for s1, s2 in itertools.combinations(sets, 2):
            if not s1 or not s2:
                continue
            inter = len(s1 & s2)
            scores.append(inter / min(len(s1), len(s2)))
        return sum(scores) / max(1, len(scores)) if scores else 0.0

    m_overlap = compute_overlap([r["markov"] for r in results])
    gpt_overlap = compute_overlap([r["nanogpt"] for r in results])
    dz_overlap = compute_overlap([r["dzeta"] for r in results])

    print("\n" + "="*80)
    print("SUMMARY COMPARISON METRICS")
    print("="*80)
    print(f"Markov  : Train Time = {markov_time:.2f}s | Overlap = {m_overlap:.4f}")
    print(f"nanoGPT : Train Time = {gpt_time:.2f}s | Params = {gpt_params:,} | Loss = {gpt_final_loss:.4f} | Overlap = {gpt_overlap:.4f}")
    print(f"DZETA   : Train Time = ~463s (20 threads) | Storage = 335 MB (4305 oscs) | Loss = 3.05e-6 | Overlap = {dz_overlap:.4f}")

if __name__ == "__main__":
    main()