Update app.py

app.py

@@ -9,27 +9,25 @@ import plotly.graph_objs as go
 import gradio as gr
 import pandas as pd
 
-# New deps for proxy fitness
+# Proxy fitness deps
 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
+# UX THEME & STYLES (cleaner, pro)
 # =========================
 CUSTOM_CSS = """
-: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)}
-#stats {display:flex; gap:16px; flex-wrap:wrap}
-.stat {flex:1; min-width:180px; background:#0b1220; color:white; border-radius:16px; padding:14px 16px}
-.stat .k {font-size:14px; opacity:0.8}
-.stat .v {font-size:22px; font-weight:700}
-.gr-button {border-radius:14px}
+:root { --radius-2xl: 18px; }
+.gradio-container {max-width: 1320px !important}
+#header-card, #viz-card, #right-card, #table-card {
+  border-radius: var(--radius-2xl);
+  box-shadow: 0 6px 24px rgba(0,0,0,0.06);
+}
+.gr-button {border-radius: 12px}
+#stats-md {font-size: 15px;}
 """
 
 # =========================
@@ -45,9 +43,9 @@ class Genome:
     dropout: float
     species: int = 0
     fitness: float = float("inf")
+    acc: Optional[float] = None  # accuracy when dataset is PIQA/HS
 
     def vector(self) -> np.ndarray:
-        # Normalized structural vector (0..1)
         return np.array([
             self.d_model / 1024.0,
             self.n_layers / 24.0,
@@ -77,7 +75,7 @@ def mutate(g: Genome, rng: random.Random, rate: float) -> Genome:
     if rng.random() < rate: g.memory_tokens = rng.choice([0, 4, 8, 16])
     if rng.random() < rate: g.dropout = rng.choice([0.0, 0.05, 0.1, 0.15])
     if rng.random() < rate * 0.5: g.species = rng.randrange(5)
-    g.fitness = float("inf")
+    g.fitness = float("inf"); g.acc = None
     return g
 
 def crossover(a: Genome, b: Genome, rng: random.Random) -> Genome:
@@ -89,7 +87,8 @@ def crossover(a: Genome, b: Genome, rng: random.Random) -> Genome:
         memory_tokens = a.memory_tokens if rng.random()<0.5 else b.memory_tokens,
         dropout = a.dropout if rng.random()<0.5 else b.dropout,
         species = a.species if rng.random()<0.5 else b.species,
-        fitness = float("inf")
+        fitness = float("inf"),
+        acc = None
     )
 
 # =========================
@@ -100,7 +99,6 @@ def rastrigin(x: np.ndarray) -> float:
     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))
@@ -110,29 +108,27 @@ class TinyMLP(nn.Module):
             nn.Linear(h1, h2), nn.ReLU(),
             nn.Linear(h2, 1)
         )
-    def forward(self, x):
-        return self.net(x).squeeze(-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
+    if name.startswith("PIQA"): return load_piqa(subset=800, seed=42)
+    if name.startswith("HellaSwag"): return load_hellaswag(subset=800, seed=42)
+    return None
 
-def _train_eval_proxy(genome: Genome, dataset_name: str, explore: float, device: str = "cpu") -> Optional[float]:
+def _train_eval_proxy(genome: Genome, dataset_name: str, explore: float, device: str = "cpu") -> Tuple[float, Optional[float]]:
+    """Returns (fitness, accuracy or None)."""
     data = _cached_dataset(dataset_name)
     if data is None:
-        return None
-    Xtr_txt, ytr, Xva_txt, yva = data
+        # Demo path handled elsewhere
+        v = genome.vector() * 2 - 1
+        return float(rastrigin(v)), None
 
-    # Hash vectorize to fixed dimension (fast, no tokenizer)
+    Xtr_txt, ytr, Xva_txt, yva = data
     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)
@@ -142,129 +138,119 @@ def _train_eval_proxy(genome: Genome, dataset_name: str, explore: float, device:
     opt = optim.AdamW(model.parameters(), lr=2e-3)
     lossf = nn.BCEWithLogitsLoss()
 
-    # small, fast loop
     model.train()
-    steps = 120
-    bs = 256
+    steps, bs = 120, 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()
+        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
+        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)
         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)
+        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)))
+    return float(max(0.0, min(1.5, fitness))), float(acc)
 
-def fitness_hook(genome: Genome, dataset: str, explore: float) -> float:
-    """Selects the correct fitness path based on dropdown."""
+def evaluate_genome(genome: Genome, dataset: str, explore: float) -> Tuple[float, Optional[float]]:
     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)
-
+        return float(base + parsimony + noise), None
     if dataset.startswith("PIQA"):
-        fit = _train_eval_proxy(genome, "PIQA", explore)
-        if fit is not None:
-            return fit
-
+        return _train_eval_proxy(genome, "PIQA", explore)
     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
+        return _train_eval_proxy(genome, "HellaSwag", explore)
+    # fallback
     v = genome.vector() * 2 - 1
-    return float(rastrigin(v))
+    return float(rastrigin(v)), None
 
 # =========================
-# PROJECTION & VIZ
+# PROJECTION & VIZ (bigger, transparent sphere, rich hover)
 # =========================
 def sphere_project(points: np.ndarray) -> np.ndarray:
-    # Fixed random projection 6D -> 3D then normalize to unit sphere
     rng = np.random.RandomState(42)
     W = rng.normal(size=(points.shape[1], 3)).astype(np.float32)
     Y = points @ W
     norms = np.linalg.norm(Y, axis=1, keepdims=True) + 1e-8
-    return Y / norms
+    return (Y / norms) * 1.15  # slightly larger radius
 
 def make_sphere_figure(points3d: np.ndarray, genomes: List[Genome], gen_idx: int) -> go.Figure:
     species = np.array([g.species for g in genomes])
-    tooltip = [
-        json.dumps({k:v for k,v in asdict(g).items() if k!="fitness"}) + f"\nfitness={g.fitness:.3f}"
-        for g in genomes
-    ]
+    # Prepare hover with all fields
+    custom = np.array([[g.d_model, g.n_layers, g.n_heads, g.ffn_mult, g.memory_tokens, g.dropout,
+                        g.species, g.fitness, (g.acc if g.acc is not None else -1.0)]
+                       for g in genomes], dtype=np.float32)
 
     scatter = go.Scatter3d(
         x=points3d[:,0], y=points3d[:,1], z=points3d[:,2],
         mode='markers',
-        marker=dict(size=6, color=species, opacity=0.9),
-        text=tooltip, hoverinfo='text'
+        marker=dict(size=7, color=species, opacity=0.95),
+        customdata=custom,
+        hovertemplate=(
+            "d_model=%{customdata[0]:.0f}<br>"
+            "layers=%{customdata[1]:.0f} · heads=%{customdata[2]:.0f}<br>"
+            "ffn_mult=%{customdata[3]:.1f} · mem=%{customdata[4]:.0f} · drop=%{customdata[5]:.2f}<br>"
+            "species=%{customdata[6]:.0f}<br>"
+            "fitness=%{customdata[7]:.4f}<br>"
+            "accuracy=%{customdata[8]:.3f}<extra></extra>"
+        )
     )
 
-    # Sphere mesh
-    u = np.linspace(0, 2*np.pi, 48)
-    v = np.linspace(0, np.pi, 24)
-    xs = np.outer(np.cos(u), np.sin(v))
-    ys = np.outer(np.sin(u), np.sin(v))
-    zs = np.outer(np.ones_like(u), np.cos(v))
-    sphere = go.Surface(x=xs, y=ys, z=zs, opacity=0.15, showscale=False)
+    # Faint sphere
+    u = np.linspace(0, 2*np.pi, 64)
+    v = np.linspace(0, np.pi, 32)
+    r = 1.15
+    xs = r*np.outer(np.cos(u), np.sin(v))
+    ys = r*np.outer(np.sin(u), np.sin(v))
+    zs = r*np.outer(np.ones_like(u), np.cos(v))
+    sphere = go.Surface(x=xs, y=ys, z=zs, opacity=0.06, showscale=False)
 
     layout = go.Layout(
         title=f"Evo Sphere — Generation {gen_idx}",
         scene=dict(xaxis=dict(visible=False), yaxis=dict(visible=False), zaxis=dict(visible=False)),
         margin=dict(l=0, r=0, t=40, b=0),
-        showlegend=False
+        showlegend=False,
+        height=680
     )
     return go.Figure(data=[sphere, scatter], layout=layout)
 
-def make_history_figure(history: List[Tuple[int,float]]) -> go.Figure:
+def make_history_figure(history: List[Tuple[int,float,float]], metric: str) -> go.Figure:
+    # history items: (gen, best_fitness, best_acc or NaN)
     xs = [h[0] for h in history]
-    ys = [h[1] for h in history]
+    if metric == "Accuracy":
+        ys = [h[2] if (h[2] == h[2]) else None for h in history]  # keep None for Demo
+        title, ylab = "Best Accuracy per Generation", "Accuracy"
+    else:
+        ys = [h[1] for h in history]
+        title, ylab = "Best Fitness per Generation", "Fitness (lower is better)"
     fig = go.Figure(data=[go.Scatter(x=xs, y=ys, mode="lines+markers")])
-    fig.update_layout(title="Best Fitness per Generation", xaxis_title="Generation",
-                      yaxis_title="Fitness (lower is better)",
-                      margin=dict(l=30,r=10,t=40,b=30))
+    fig.update_layout(title=title, xaxis_title="Generation", yaxis_title=ylab,
+                      margin=dict(l=30,r=10,t=40,b=30), height=360)
     return fig
 
 def approx_params(g: Genome) -> int:
-    # Very rough estimate ignoring embeddings/vocab:
-    # 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
-    total += 1000 * g.memory_tokens  # tiny bump for memory pathways (illustrative)
+    total = per_layer * g.n_layers + 1000 * g.memory_tokens
     return int(total)
 
 # =========================
@@ -277,7 +263,7 @@ class EvoRunner:
         self.stop_flag = False
         self.state: Dict[str, Any] = {}
 
-    def run(self, dataset, pop_size, generations, mutation_rate, explore, exploit, seed, pace_ms):
+    def run(self, dataset, pop_size, generations, mutation_rate, explore, exploit, seed, pace_ms, metric_choice):
         rng = random.Random(int(seed))
         self.stop_flag = False
         self.running = True
@@ -285,15 +271,16 @@ class EvoRunner:
         pop: List[Genome] = [random_genome(rng) for _ in range(pop_size)]
         # initial eval
         for g in pop:
-            g.fitness = fitness_hook(g, dataset, explore)
+            fit, acc = evaluate_genome(g, dataset, explore)
+            g.fitness, g.acc = fit, acc
 
-        history: List[Tuple[int,float]] = []
+        history: List[Tuple[int,float,float]] = []  # (gen, best_fitness, best_acc or NaN)
         best_overall: Optional[Genome] = None
 
         for gen in range(1, generations+1):
             if self.stop_flag: break
 
-            # Selection: tournament size depends on exploitation
+            # Selection (tournament)
             k = max(2, int(2 + exploit * 5))
             parents = []
             for _ in range(pop_size):
@@ -303,16 +290,16 @@ class EvoRunner:
             # Reproduce
             children = []
             for i in range(0, pop_size, 2):
-                a = parents[i]
-                b = parents[(i+1) % pop_size]
+                a = parents[i]; b = parents[(i+1) % pop_size]
                 child1 = mutate(crossover(a,b,rng), rng, mutation_rate)
                 child2 = mutate(crossover(b,a,rng), rng, mutation_rate)
                 children.extend([child1, child2])
             children = children[:pop_size]
 
-            # Evaluate kids
+            # Evaluate children
             for c in children:
-                c.fitness = fitness_hook(c, dataset, explore)
+                fit, acc = evaluate_genome(c, dataset, explore)
+                c.fitness, c.acc = fit, acc
 
             # Elitism
             elite_n = max(1, pop_size // 10)
@@ -326,18 +313,19 @@ class EvoRunner:
             if best_overall is None or best.fitness < best_overall.fitness:
                 best_overall = best
 
-            history.append((gen, best.fitness))
+            history.append((gen, best.fitness, (best.acc if best.acc is not None else float("nan"))))
 
             # Viz snapshot
             P = np.stack([g.vector() for g in pop], axis=0)
             P3 = sphere_project(P)
             sphere_fig = make_sphere_figure(P3, pop, gen)
-            hist_fig = make_history_figure(history)
+            hist_fig = make_history_figure(history, metric_choice)
             top = sorted(pop, key=lambda x: x.fitness)[: min(12, len(pop))]
             top_table = [
                 {
                     "gen": gen,
                     "fitness": round(t.fitness, 4),
+                    "accuracy": (None if t.acc is None else round(float(t.acc), 4)),
                     "d_model": t.d_model,
                     "layers": t.n_layers,
                     "heads": t.n_heads,
@@ -357,7 +345,8 @@ class EvoRunner:
                     "top": top_table,
                     "best": best_card,
                     "gen": gen,
-                    "dataset": dataset
+                    "dataset": dataset,
+                    "metric": metric_choice
                 }
 
             time.sleep(max(0.0, pace_ms/1000.0))
@@ -369,16 +358,15 @@ class EvoRunner:
         t = threading.Thread(target=self.run, args=args, kwargs=kwargs, daemon=True)
         t.start()
 
-    def stop(self):
-        self.stop_flag = True
+    def stop(self): self.stop_flag = True
 
 runner = EvoRunner()
 
 # =========================
-# GRADIO UI CALLBACKS
+# UI CALLBACKS
 # =========================
-def start_evo(dataset, pop, gens, mut, explore, exploit, seed, pace_ms):
-    runner.start(dataset, int(pop), int(gens), float(mut), float(explore), float(exploit), int(seed), int(pace_ms))
+def start_evo(dataset, pop, gens, mut, explore, exploit, seed, pace_ms, metric_choice):
+    runner.start(dataset, int(pop), int(gens), float(mut), float(explore), float(exploit), int(seed), int(pace_ms), metric_choice)
     return (gr.update(interactive=False), gr.update(interactive=True))
 
 def stop_evo():
@@ -389,16 +377,19 @@ def poll_state():
     with runner.lock:
         s = runner.state.copy()
     sphere = s.get("sphere", go.Figure())
-    history = s.get("history", go.Figure())
+    history = s.get("history", go.Figure())  # already built by runner
     best = s.get("best", {})
     gen = s.get("gen", 0)
     dataset = s.get("dataset", "Demo (Surrogate)")
     top = s.get("top", [])
+    # Stats text
     if best:
+        acc_txt = "—" if best.get("accuracy") is None else f"{best.get('accuracy'):.3f}"
         stats_md = (
             f"**Dataset:** {dataset}  \n"
             f"**Generation:** {gen}  \n"
             f"**Best fitness:** {best.get('fitness','–')}  \n"
+            f"**Best accuracy:** {acc_txt}  \n"
             f"**Config:** d_model={best.get('d_model')} · layers={best.get('layers')} · "
             f"heads={best.get('heads')} · ffn_mult={best.get('ffn_mult')} · mem={best.get('mem')} · "
             f"dropout={best.get('dropout')}  \n"
@@ -419,14 +410,13 @@ def export_snapshot():
     return path
 
 # =========================
-# BUILD UI
+# BUILD UI (bigger sphere, metric toggle)
 # =========================
 with gr.Blocks(theme=gr.themes.Soft(), css=CUSTOM_CSS) as demo:
     with gr.Column(elem_id="header-card"):
         gr.Markdown(
-            "# Evo Playground — Live Evolving Transformer Architectures\n"
-            "Watch the population **mutate, recombine, and converge** in real time. "
-            "Choose a dataset and search behavior; the 3D sphere shows the architecture landscape (species = colors)."
+            "# Evo Playground — Live Evolution of Transformer Architectures\n"
+            "Tune the search, watch the population converge, and track **accuracy** in real time (PIQA/HellaSwag)."
         )
 
     with gr.Row():
@@ -435,9 +425,9 @@ with gr.Blocks(theme=gr.themes.Soft(), css=CUSTOM_CSS) as demo:
             with gr.Group():
                 dataset = gr.Dropdown(
                     label="Dataset",
-                    choices=["Demo (Surrogate)", "PIQA (Phase 2)", "HellaSwag (Phase 2)", "WikiText Perplexity (Phase 2)"],
+                    choices=["Demo (Surrogate)", "PIQA (Phase 2)", "HellaSwag (Phase 2)"],
                     value="Demo (Surrogate)",
-                    info="Demo is instant. PIQA/HellaSwag run a tiny CPU MLP proxy for real dataset fitness."
+                    info="PIQA/HellaSwag compute real proxy accuracy; Demo uses a fast surrogate."
                 )
                 pop = gr.Slider(8, 80, value=24, step=2, label="Population size")
                 gens = gr.Slider(5, 200, value=60, step=1, label="Max generations")
@@ -447,12 +437,13 @@ with gr.Blocks(theme=gr.themes.Soft(), css=CUSTOM_CSS) as demo:
                     exploit = gr.Slider(0.0, 1.0, value=0.65, step=0.05, label="Exploitation")
                 seed = gr.Number(value=42, label="Seed", precision=0)
                 pace = gr.Slider(0, 1000, value=120, step=10, label="Pace (ms between gens)")
+                metric_choice = gr.Radio(choices=["Accuracy", "Fitness"], value="Accuracy", label="History Metric")
                 with gr.Row():
                     start = gr.Button("▶ Start Evolution", variant="primary")
                     stop = gr.Button("⏹ Stop", variant="secondary")
 
             with gr.Group(elem_id="right-card"):
-                stats_md = gr.Markdown("Waiting…")
+                stats_md = gr.Markdown("Waiting…", elem_id="stats-md")
                 export_btn = gr.Button("Export Snapshot (JSON)")
                 export_file = gr.File(label="Download snapshot", visible=False)
 
@@ -461,16 +452,16 @@ with gr.Blocks(theme=gr.themes.Soft(), css=CUSTOM_CSS) as demo:
             with gr.Group(elem_id="viz-card"):
                 sphere_plot = gr.Plot(label="Evolution Sphere")
             with gr.Group(elem_id="viz-card"):
-                hist_plot = gr.Plot(label="Best Fitness History")
+                hist_plot = gr.Plot(label="Progress")
             with gr.Group(elem_id="table-card"):
                 top_df = gr.Dataframe(label="Top Genomes (live)", wrap=True, interactive=False)
 
     # Wiring
-    start.click(start_evo, [dataset, pop, gens, mut, explore, exploit, seed, pace], [start, stop])
+    start.click(start_evo, [dataset, pop, gens, mut, explore, exploit, seed, pace, metric_choice], [start, stop])
     stop.click(stop_evo, [], [start, stop])
     export_btn.click(export_snapshot, [], [export_file])
 
-    # Initial paint once when app loads
+    # Initial paint
     demo.load(poll_state, None, [sphere_plot, hist_plot, stats_md, top_df])
 
     # Continuous polling (every 0.7s)