Update app.py

app.py

@@ -1,4 +1,4 @@
-# app.py
+# app.py — Minimal, pro UI with big transparent sphere and clean hover
 import math, json, random, time, threading
 from dataclasses import dataclass, asdict
 from typing import List, Tuple, Dict, Any, Optional
@@ -6,33 +6,32 @@ from functools import lru_cache
 
 import numpy as np
 import plotly.graph_objs as go
+import plotly.io as pio
 import gradio as gr
 import pandas as pd
 
-# Proxy fitness deps
 import torch
 import torch.nn as nn
 import torch.optim as optim
 
 from data_utils import load_piqa, load_hellaswag, hash_vectorize
 
-# =========================
-# UX THEME & STYLES (cleaner, pro)
-# =========================
+# ---------- Minimal style ----------
 CUSTOM_CSS = """
-: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;}
+:root { --radius: 14px; --fg:#0f172a; --muted:#64748b; --line:#e5e7eb; }
+* { font-family: Inter, ui-sans-serif, system-ui, -apple-system, Segoe UI, Roboto, Helvetica Neue, Arial, "Noto Sans", "Apple Color Emoji", "Segoe UI Emoji"; }
+.gradio-container { max-width: 1180px !important; }
+#header { border-radius: var(--radius); padding: 8px 6px; }
+h1, h2, h3, .gr-markdown { color: var(--fg); }
+.gr-button { border-radius: 10px; }
+.controls .gr-group, .panel { border: 1px solid var(--line); border-radius: var(--radius); }
+.panel { padding: 10px; }
+#stats { font-weight: 300; color: var(--fg); }
+#stats strong { font-weight: 500; }
+.small { font-size: 13px; color: var(--muted); }
 """
 
-# =========================
-# GENOME & EVOLUTION CORE
-# =========================
+# ---------- Genome ----------
 @dataclass
 class Genome:
     d_model: int
@@ -43,7 +42,7 @@ class Genome:
     dropout: float
     species: int = 0
     fitness: float = float("inf")
-    acc: Optional[float] = None  # accuracy when dataset is PIQA/HS
+    acc: Optional[float] = None
 
     def vector(self) -> np.ndarray:
         return np.array([
@@ -91,9 +90,7 @@ def crossover(a: Genome, b: Genome, rng: random.Random) -> Genome:
         acc = None
     )
 
-# =========================
-# PROXY FITNESS (Phase 2a)
-# =========================
+# ---------- Proxy fitness ----------
 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))
@@ -116,14 +113,11 @@ def _cached_dataset(name: str):
     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") -> Tuple[float, Optional[float]]:
-    """Returns (fitness, accuracy or None)."""
+def _train_eval_proxy(genome: Genome, dataset_name: str, explore: float, device: str = "cpu"):
     data = _cached_dataset(dataset_name)
     if data is None:
-        # Demo path handled elsewhere
         v = genome.vector() * 2 - 1
         return float(rastrigin(v)), None
-
     Xtr_txt, ytr, Xva_txt, yva = data
     nfeat = 4096
     Xtr = hash_vectorize(Xtr_txt, n_features=nfeat, seed=1234)
@@ -139,8 +133,7 @@ def _train_eval_proxy(genome: Genome, dataset_name: str, explore: float, device:
     lossf = nn.BCEWithLogitsLoss()
 
     model.train()
-    steps, bs = 120, 256
-    N = Xtr_t.size(0)
+    steps, bs, N = 120, 256, 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)
@@ -169,7 +162,7 @@ def _train_eval_proxy(genome: Genome, dataset_name: str, explore: float, device:
     fitness = (1.0 - acc) + parsimony + noise
     return float(max(0.0, min(1.5, fitness))), float(acc)
 
-def evaluate_genome(genome: Genome, dataset: str, explore: float) -> Tuple[float, Optional[float]]:
+def evaluate_genome(genome: Genome, dataset: str, explore: float):
     if dataset == "Demo (Surrogate)":
         v = genome.vector() * 2 - 1
         base = rastrigin(v)
@@ -180,23 +173,30 @@ def evaluate_genome(genome: Genome, dataset: str, explore: float) -> Tuple[float
         return _train_eval_proxy(genome, "PIQA", explore)
     if dataset.startswith("HellaSwag"):
         return _train_eval_proxy(genome, "HellaSwag", explore)
-    # fallback
     v = genome.vector() * 2 - 1
     return float(rastrigin(v)), None
 
-# =========================
-# PROJECTION & VIZ (bigger, transparent sphere, rich hover)
-# =========================
+# ---------- Viz helpers (bigger, transparent sphere) ----------
+PALETTE = ["#111827", "#334155", "#475569", "#64748b", "#94a3b8"]  # muted grayscale/blue
+BG = "white"
+
 def sphere_project(points: np.ndarray) -> np.ndarray:
     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) * 1.15  # slightly larger radius
+    return (Y / norms) * 1.2
+
+def _species_colors(species: np.ndarray) -> list:
+    colors = []
+    for s in species:
+        c = PALETTE[int(s) % len(PALETTE)]
+        colors.append(c)
+    return colors
 
 def make_sphere_figure(points3d: np.ndarray, genomes: List[Genome], gen_idx: int) -> go.Figure:
     species = np.array([g.species for g in genomes])
-    # Prepare hover with all fields
+    colors = _species_colors(species)
     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)
@@ -204,11 +204,11 @@ def make_sphere_figure(points3d: np.ndarray, genomes: List[Genome], gen_idx: int
     scatter = go.Scatter3d(
         x=points3d[:,0], y=points3d[:,1], z=points3d[:,2],
         mode='markers',
-        marker=dict(size=7, color=species, opacity=0.95),
+        marker=dict(size=6.5, color=colors, opacity=0.92),
         customdata=custom,
         hovertemplate=(
-            "d_model=%{customdata[0]:.0f}<br>"
-            "layers=%{customdata[1]:.0f} · heads=%{customdata[2]:.0f}<br>"
+            "<b>Genome</b><br>"
+            "d_model=%{customdata[0]:.0f} · 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>"
@@ -216,46 +216,63 @@ def make_sphere_figure(points3d: np.ndarray, genomes: List[Genome], gen_idx: int
         )
     )
 
-    # Faint sphere
-    u = np.linspace(0, 2*np.pi, 64)
-    v = np.linspace(0, np.pi, 32)
-    r = 1.15
+    # Subtle, large sphere
+    u = np.linspace(0, 2*np.pi, 72)
+    v = np.linspace(0, np.pi, 36)
+    r = 1.2
     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)
+    sphere = go.Surface(
+        x=xs, y=ys, z=zs,
+        opacity=0.08,
+        showscale=False,
+        colorscale=[[0, "#cbd5e1"], [1, "#cbd5e1"]],
+        hoverinfo="skip"
+    )
 
     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),
+        paper_bgcolor=BG, plot_bgcolor=BG,
+        title=f"Evo Architecture Sphere — Gen {gen_idx}",
+        scene=dict(
+            xaxis=dict(visible=False), yaxis=dict(visible=False), zaxis=dict(visible=False),
+            bgcolor=BG
+        ),
+        margin=dict(l=0, r=0, t=36, b=0),
         showlegend=False,
-        height=680
+        height=720,
+        font=dict(family="Inter, Arial, sans-serif", size=14)
     )
     return go.Figure(data=[sphere, scatter], layout=layout)
 
 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]
     if metric == "Accuracy":
-        ys = [h[2] if (h[2] == h[2]) else None for h in history]  # keep None for Demo
+        ys = [h[2] if (h[2] == h[2]) else None for h in history]
         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=title, xaxis_title="Generation", yaxis_title=ylab,
-                      margin=dict(l=30,r=10,t=40,b=30), height=360)
+        title, ylab = "Best Fitness per Generation", "Fitness (↓ better)"
+    fig = go.Figure(data=[go.Scatter(x=xs, y=ys, mode="lines+markers", line=dict(width=2))])
+    fig.update_layout(
+        paper_bgcolor=BG, plot_bgcolor=BG,
+        title=title, xaxis_title="Generation", yaxis_title=ylab,
+        margin=dict(l=30, r=10, t=36, b=30),
+        height=340,
+        font=dict(family="Inter, Arial, sans-serif", size=14)
+    )
     return fig
 
+def fig_to_html(fig: go.Figure) -> str:
+    # Robust Plotly rendering inside Gradio
+    return pio.to_html(fig, include_plotlyjs="cdn", full_html=False, config=dict(displaylogo=False))
+
 def approx_params(g: Genome) -> int:
     per_layer = (4.0 + 2.0 * float(g.ffn_mult)) * (g.d_model ** 2)
     total = per_layer * g.n_layers + 1000 * g.memory_tokens
     return int(total)
 
-# =========================
-# ORCHESTRATOR
-# =========================
+# ---------- Orchestrator ----------
 class EvoRunner:
     def __init__(self):
         self.lock = threading.Lock()
@@ -269,25 +286,22 @@ class EvoRunner:
         self.running = True
 
         pop: List[Genome] = [random_genome(rng) for _ in range(pop_size)]
-        # initial eval
         for g in pop:
             fit, acc = evaluate_genome(g, dataset, explore)
             g.fitness, g.acc = fit, acc
 
-        history: List[Tuple[int,float,float]] = []  # (gen, best_fitness, best_acc or NaN)
+        history: List[Tuple[int,float,float]] = []
         best_overall: Optional[Genome] = None
 
         for gen in range(1, generations+1):
             if self.stop_flag: break
 
-            # Selection (tournament)
             k = max(2, int(2 + exploit * 5))
             parents = []
             for _ in range(pop_size):
                 sample = rng.sample(pop, k=k)
                 parents.append(min(sample, key=lambda x: x.fitness))
 
-            # Reproduce
             children = []
             for i in range(0, pop_size, 2):
                 a = parents[i]; b = parents[(i+1) % pop_size]
@@ -296,30 +310,25 @@ class EvoRunner:
                 children.extend([child1, child2])
             children = children[:pop_size]
 
-            # Evaluate children
             for c in children:
                 fit, acc = evaluate_genome(c, dataset, explore)
                 c.fitness, c.acc = fit, acc
 
-            # Elitism
             elite_n = max(1, pop_size // 10)
             elites = sorted(pop, key=lambda x: x.fitness)[:elite_n]
-
-            # Next pop
             pop = sorted(children, key=lambda x: x.fitness)
             pop[-elite_n:] = elites
 
             best = min(pop, key=lambda x: x.fitness)
-            if best_overall is None or best.fitness < best_overall.fitness:
-                best_overall = best
+            if best_overall is None or best.fitness < best_overall.fitness: best_overall = best
 
             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, metric_choice)
+
             top = sorted(pop, key=lambda x: x.fitness)[: min(12, len(pop))]
             top_table = [
                 {
@@ -340,8 +349,8 @@ class EvoRunner:
 
             with self.lock:
                 self.state = {
-                    "sphere": sphere_fig,
-                    "history": hist_fig,
+                    "sphere_html": fig_to_html(sphere_fig),
+                    "history_html": fig_to_html(hist_fig),
                     "top": top_table,
                     "best": best_card,
                     "gen": gen,
@@ -350,21 +359,17 @@ class EvoRunner:
                 }
 
             time.sleep(max(0.0, pace_ms/1000.0))
-
         self.running = False
 
     def start(self, *args, **kwargs):
         if self.running: return
         t = threading.Thread(target=self.run, args=args, kwargs=kwargs, daemon=True)
         t.start()
-
     def stop(self): self.stop_flag = True
 
 runner = EvoRunner()
 
-# =========================
-# UI CALLBACKS
-# =========================
+# ---------- UI callbacks ----------
 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))
@@ -376,13 +381,12 @@ def stop_evo():
 def poll_state():
     with runner.lock:
         s = runner.state.copy()
-    sphere = s.get("sphere", go.Figure())
-    history = s.get("history", go.Figure())  # already built by runner
+    sphere_html = s.get("sphere_html", "")
+    history_html = s.get("history_html", "")
     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 = (
@@ -398,7 +402,7 @@ def poll_state():
     else:
         stats_md = "Waiting… click **Start Evolution**."
     df = pd.DataFrame(top)
-    return sphere, history, stats_md, df
+    return sphere_html, history_html, stats_md, df
 
 def export_snapshot():
     from json import dumps
@@ -409,25 +413,19 @@ def export_snapshot():
         f.write(payload)
     return path
 
-# =========================
-# 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 Evolution of Transformer Architectures\n"
-            "Tune the search, watch the population converge, and track **accuracy** in real time (PIQA/HellaSwag)."
-        )
+# ---------- Build UI (minimal layout) ----------
+with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Soft()) as demo:
+    with gr.Column(elem_id="header"):
+        gr.Markdown("## Evo Playground — Minimal Live Evolution (PIQA / HellaSwag accuracy)")
 
     with gr.Row():
-        # LEFT: Controls
-        with gr.Column(scale=1):
+        with gr.Column(scale=1, elem_classes=["controls"]):
             with gr.Group():
                 dataset = gr.Dropdown(
                     label="Dataset",
                     choices=["Demo (Surrogate)", "PIQA (Phase 2)", "HellaSwag (Phase 2)"],
                     value="Demo (Surrogate)",
-                    info="PIQA/HellaSwag compute real proxy accuracy; Demo uses a fast surrogate."
+                    info="PIQA/HellaSwag compute real proxy accuracy; Demo is 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")
@@ -442,18 +440,19 @@ with gr.Blocks(theme=gr.themes.Soft(), css=CUSTOM_CSS) as demo:
                     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…", elem_id="stats-md")
+            with gr.Group(elem_classes=["panel"]):
+                stats_md = gr.Markdown("Waiting…", elem_id="stats")
+
+            with gr.Group(elem_classes=["panel"]):
                 export_btn = gr.Button("Export Snapshot (JSON)")
                 export_file = gr.File(label="Download snapshot", visible=False)
 
-        # RIGHT: Viz + Table
         with gr.Column(scale=2):
-            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="Progress")
-            with gr.Group(elem_id="table-card"):
+            with gr.Group(elem_classes=["panel"]):
+                sphere_html = gr.HTML()
+            with gr.Group(elem_classes=["panel"]):
+                hist_html = gr.HTML()
+            with gr.Group(elem_classes=["panel"]):
                 top_df = gr.Dataframe(label="Top Genomes (live)", wrap=True, interactive=False)
 
     # Wiring
@@ -461,12 +460,9 @@ with gr.Blocks(theme=gr.themes.Soft(), css=CUSTOM_CSS) as demo:
     stop.click(stop_evo, [], [start, stop])
     export_btn.click(export_snapshot, [], [export_file])
 
-    # Initial paint
-    demo.load(poll_state, None, [sphere_plot, hist_plot, stats_md, top_df])
-
-    # Continuous polling (every 0.7s)
-    poller = gr.Timer(0.7)
-    poller.tick(poll_state, None, [sphere_plot, hist_plot, stats_md, top_df])
+    # Initial paint + polling
+    demo.load(poll_state, None, [sphere_html, hist_html, stats_md, top_df])
+    gr.Timer(0.7).tick(poll_state, None, [sphere_html, hist_html, stats_md, top_df])
 
 if __name__ == "__main__":
     demo.launch()