Update app.py

app.py

@@ -1,18 +1,27 @@
 # app.py
-import math, json, random, time, threading, io, os
+import math, json, random, time, threading
 from dataclasses import dataclass, asdict
 from typing import List, Tuple, Dict, Any, Optional
+from functools import lru_cache
+
 import numpy as np
 import plotly.graph_objs as go
 import gradio as gr
+import pandas as pd
+
+# New deps for proxy fitness
+import torch
+import torch.nn as nn
+import torch.optim as optim
+
+# Local utils (add this file next to app.py)
+from data_utils import load_piqa, load_hellaswag, hash_vectorize
 
 # =========================
 # UX THEME & STYLES
 # =========================
 CUSTOM_CSS = """
-:root {
-  --radius-2xl: 20px;
-}
+:root { --radius-2xl: 20px; }
 .gradio-container {max-width: 1400px !important}
 #header-card {border-radius: var(--radius-2xl); box-shadow: 0 6px 24px rgba(0,0,0,0.08)}
 #viz-card, #right-card, #table-card {border-radius: var(--radius-2xl); box-shadow: 0 6px 24px rgba(0,0,0,0.06)}
@@ -84,26 +93,121 @@ def crossover(a: Genome, b: Genome, rng: random.Random) -> Genome:
     )
 
 # =========================
-# FITNESS HOOK (Phase 1: fast surrogate)
-# Swap this later for real PIQA/HellaSwag evaluation
+# PROXY FITNESS (Phase 2a)
 # =========================
 def rastrigin(x: np.ndarray) -> float:
     A, n = 10.0, x.shape[0]
     return A * n + np.sum(x**2 - A * np.cos(2 * math.pi * x))
 
+class TinyMLP(nn.Module):
+    """Small MLP whose capacity depends on the genome (so evolution matters)."""
+    def __init__(self, in_dim: int, genome: Genome):
+        super().__init__()
+        h1 = max(64, int(0.25 * genome.d_model))
+        h2 = max(32, int(genome.ffn_mult * 32))
+        self.net = nn.Sequential(
+            nn.Linear(in_dim, h1), nn.ReLU(),
+            nn.Linear(h1, h2), nn.ReLU(),
+            nn.Linear(h2, 1)
+        )
+    def forward(self, x):
+        return self.net(x).squeeze(-1)
+
+@lru_cache(maxsize=4)
+def _cached_dataset(name: str):
+    if name.startswith("PIQA"):
+        return load_piqa(subset=800, seed=42)
+    if name.startswith("HellaSwag"):
+        return load_hellaswag(subset=800, seed=42)
+    return None  # Demo uses surrogate
+
+def _train_eval_proxy(genome: Genome, dataset_name: str, explore: float, device: str = "cpu") -> Optional[float]:
+    data = _cached_dataset(dataset_name)
+    if data is None:
+        return None
+    Xtr_txt, ytr, Xva_txt, yva = data
+
+    # Hash vectorize to fixed dimension (fast, no tokenizer)
+    nfeat = 4096
+    Xtr = hash_vectorize(Xtr_txt, n_features=nfeat, seed=1234)
+    Xva = hash_vectorize(Xva_txt, n_features=nfeat, seed=5678)
+
+    # to torch tensors
+    Xtr_t = torch.from_numpy(Xtr)
+    ytr_t = torch.from_numpy(ytr.astype(np.float32))
+    Xva_t = torch.from_numpy(Xva)
+    yva_t = torch.from_numpy(yva.astype(np.float32))
+
+    model = TinyMLP(nfeat, genome).to(device)
+    opt = optim.AdamW(model.parameters(), lr=2e-3)
+    lossf = nn.BCEWithLogitsLoss()
+
+    # small, fast loop
+    model.train()
+    steps = 120
+    bs = 256
+    N = Xtr_t.size(0)
+    for _ in range(steps):
+        idx = torch.randint(0, N, (bs,))
+        xb = Xtr_t[idx].to(device)
+        yb = ytr_t[idx].to(device)
+        logits = model(xb)
+        loss = lossf(logits, yb)
+        opt.zero_grad()
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        opt.step()
+
+    # eval
+    model.eval()
+    with torch.no_grad():
+        logits = model(Xva_t.to(device))
+        probs = torch.sigmoid(logits).cpu().numpy()
+
+    # Turn rows into accuracy
+    if dataset_name.startswith("PIQA"):
+        # rows in pairs [A,B]; label vector marks which row is positive
+        probs = probs.reshape(-1, 2)
+        yva2 = yva.reshape(-1, 2)
+        pred = (probs[:, 0] > probs[:, 1]).astype(np.int64)
+        truth = (yva2[:, 0] == 1).astype(np.int64)  # 1 means first row is correct
+        acc = float((pred == truth).mean())
+    else:
+        # HellaSwag: groups of 4; pick argmax
+        probs = probs.reshape(-1, 4)
+        yva2 = yva.reshape(-1, 4)
+        pred = probs.argmax(axis=1)
+        truth = yva2.argmax(axis=1)
+        acc = float((pred == truth).mean())
+
+    # Fitness = error + tiny parsimony + small exploration noise (minimize)
+    parsimony = 0.00000002 * (genome.d_model**2 * genome.n_layers) + 0.0001 * genome.memory_tokens
+    noise = np.random.normal(scale=0.01 * max(0.0, min(1.0, explore)))
+    fitness = (1.0 - acc) + parsimony + noise
+    return float(max(0.0, min(1.5, fitness)))
+
 def fitness_hook(genome: Genome, dataset: str, explore: float) -> float:
-    """
-    Phase 1 (demo, fast):
-    - Build vector v in [-1,1] from genome params and score via Rastrigin.
-    - Add small parsimony penalty and exploration noise.
-    Phase 2 (real):
-    - Replace with tiny train/eval steps on chosen dataset (PIQA/HellaSwag/WikiText-ppl).
-    """
-    v = genome.vector() * 2 - 1  # [-1,1]
-    base = rastrigin(v)
-    parsimony = 0.001 * (genome.d_model + 50*genome.n_layers + 20*genome.n_heads + 100*genome.memory_tokens)
-    noise = np.random.normal(scale=0.05 * max(0.0, min(1.0, explore)))
-    return float(base + parsimony + noise)
+    """Selects the correct fitness path based on dropdown."""
+    if dataset == "Demo (Surrogate)":
+        v = genome.vector() * 2 - 1
+        base = rastrigin(v)
+        parsimony = 0.001 * (genome.d_model + 50*genome.n_layers + 20*genome.n_heads + 100*genome.memory_tokens)
+        noise = np.random.normal(scale=0.05 * max(0.0, min(1.0, explore)))
+        return float(base + parsimony + noise)
+
+    if dataset.startswith("PIQA"):
+        fit = _train_eval_proxy(genome, "PIQA", explore)
+        if fit is not None:
+            return fit
+
+    if dataset.startswith("HellaSwag"):
+        fit = _train_eval_proxy(genome, "HellaSwag", explore)
+        if fit is not None:
+            return fit
+
+    # fallback to surrogate if something went wrong
+    v = genome.vector() * 2 - 1
+    return float(rastrigin(v))
 
 # =========================
 # PROJECTION & VIZ
@@ -160,8 +264,7 @@ def approx_params(g: Genome) -> int:
     # per-layer ~ (4 + 2*ffn_mult) * d_model^2
     per_layer = (4.0 + 2.0 * float(g.ffn_mult)) * (g.d_model ** 2)
     total = per_layer * g.n_layers
-    # tiny bump for memory tokens pathways (illustrative only)
-    total += 1000 * g.memory_tokens
+    total += 1000 * g.memory_tokens  # tiny bump for memory pathways (illustrative)
     return int(total)
 
 # =========================
@@ -303,13 +406,13 @@ def poll_state():
         )
     else:
         stats_md = "Waiting… click **Start Evolution**."
-    import pandas as pd
     df = pd.DataFrame(top)
     return sphere, history, stats_md, df
 
 def export_snapshot():
+    from json import dumps
     with runner.lock:
-        payload = json.dumps(runner.state, default=lambda o: o, indent=2)
+        payload = dumps(runner.state, default=lambda o: o, indent=2)
     path = "evo_snapshot.json"
     with open(path, "w", encoding="utf-8") as f:
         f.write(payload)
@@ -334,7 +437,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=CUSTOM_CSS) as demo:
                     label="Dataset",
                     choices=["Demo (Surrogate)", "PIQA (Phase 2)", "HellaSwag (Phase 2)", "WikiText Perplexity (Phase 2)"],
                     value="Demo (Surrogate)",
-                    info="Demo is instant. Phase 2 datasets will do tiny train/eval steps per genome."
+                    info="Demo is instant. PIQA/HellaSwag run a tiny CPU MLP proxy for real dataset fitness."
                 )
                 pop = gr.Slider(8, 80, value=24, step=2, label="Population size")
                 gens = gr.Slider(5, 200, value=60, step=1, label="Max generations")
@@ -350,7 +453,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css=CUSTOM_CSS) as demo:
 
             with gr.Group(elem_id="right-card"):
                 stats_md = gr.Markdown("Waiting…")
-
                 export_btn = gr.Button("Export Snapshot (JSON)")
                 export_file = gr.File(label="Download snapshot", visible=False)