Update app.py

app.py

@@ -1,4 +1,4 @@
-# app.py — Minimal, pro UI with big transparent sphere, accuracy, and robust rendering
+# app.py — Minimal dark UI, default idle sphere, Clear button, inline Plotly
 import math, json, random, time, threading
 from dataclasses import dataclass, asdict
 from typing import List, Tuple, Dict, Any, Optional
@@ -16,22 +16,28 @@ import torch.optim as optim
 
 from data_utils import load_piqa, load_hellaswag, hash_vectorize
 
-# ---------- Minimal style ----------
+# =========================
+# STYLE — calm, dark, thin
+# =========================
 CUSTOM_CSS = """
-: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; }
+:root { --radius: 14px; --fg:#E5E7EB; --muted:#94A3B8; --line:#111827; --bg:#0F1A24; }
+* { font-family: Inter, ui-sans-serif, system-ui, -apple-system, Segoe UI, Roboto, Helvetica Neue, Arial; font-weight: 300; }
+.gradio-container { max-width: 1140px !important; background: var(--bg); }
+#header { border-radius: var(--radius); padding: 6px 2px; }
 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); }
+.gr-button { border-radius: 10px; }
+.controls .gr-group, .panel { border: 1px solid #1f2b36; border-radius: var(--radius); background: #0c161f; }
 .panel { padding: 10px; }
-#stats { font-weight: 300; color: var(--fg); }
+#stats { color: var(--fg); }
 #stats strong { font-weight: 500; }
-.small { font-size: 13px; color: var(--muted); }
+.small { font-size: 12px; color: var(--muted); }
+label, .gradio-container * { color: var(--fg); }
+input, textarea, select { color: var(--fg) !important; }
 """
 
-# ---------- Genome ----------
+# =========================
+# GENOME
+# =========================
 @dataclass
 class Genome:
     d_model: int
@@ -86,11 +92,12 @@ 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"),
-        acc = None
+        fitness = float("inf"), acc=None
     )
 
-# ---------- Proxy fitness ----------
+# =========================
+# 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))
@@ -107,9 +114,9 @@ class TinyMLP(nn.Module):
         )
     def forward(self, x): return self.net(x).squeeze(-1)
 
+from functools import lru_cache
 @lru_cache(maxsize=4)
 def _cached_dataset(name: str):
-    # Defensive: if loading fails (e.g., datasets version / no internet), return None
     try:
         if name.startswith("PIQA"): return load_piqa(subset=800, seed=42)
         if name.startswith("HellaSwag"): return load_hellaswag(subset=800, seed=42)
@@ -117,10 +124,9 @@ def _cached_dataset(name: str):
         return None
     return None
 
-def _train_eval_proxy(genome: Genome, dataset_name: str, explore: float, device: str = "cpu"):
+def _train_eval_proxy(genome: Genome, dataset_name: str, explore: float, device: str="cpu"):
     data = _cached_dataset(dataset_name)
     if data is None:
-        # Fallback to surrogate to keep the UI alive
         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)
@@ -132,17 +138,14 @@ def _train_eval_proxy(genome: Genome, dataset_name: str, explore: float, device:
     Xtr = hash_vectorize(Xtr_txt, n_features=nfeat, seed=1234)
     Xva = hash_vectorize(Xva_txt, n_features=nfeat, seed=5678)
 
-    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))
+    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()
 
-    model.train()
-    steps, bs, N = 120, 256, Xtr_t.size(0)
+    model.train(); 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)
@@ -157,12 +160,12 @@ def _train_eval_proxy(genome: Genome, dataset_name: str, explore: float, device:
         probs = torch.sigmoid(logits).cpu().numpy()
 
     if dataset_name.startswith("PIQA"):
-        probs = probs.reshape(-1, 2); yva2 = yva.reshape(-1, 2)
+        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:
-        probs = probs.reshape(-1, 4); yva2 = yva.reshape(-1, 4)
+        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())
 
@@ -178,39 +181,53 @@ def evaluate_genome(genome: Genome, dataset: str, explore: float):
         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), None
-    if dataset.startswith("PIQA"):
-        return _train_eval_proxy(genome, "PIQA", explore)
-    if dataset.startswith("HellaSwag"):
-        return _train_eval_proxy(genome, "HellaSwag", explore)
-    # Fallback
+    if dataset.startswith("PIQA"): return _train_eval_proxy(genome, "PIQA", explore)
+    if dataset.startswith("HellaSwag"): return _train_eval_proxy(genome, "HellaSwag", explore)
     v = genome.vector() * 2 - 1
     return float(rastrigin(v)), None
 
-# ---------- Viz helpers (bigger, transparent sphere) ----------
-PALETTE = ["#111827", "#334155", "#475569", "#64748b", "#94a3b8"]  # muted grayscale/blue
-BG = "white"
+# =========================
+# VIZ — big transparent sphere
+# =========================
+BG = "#0F1A24"
+DOT = "#93C5FD"   # soft blue dot
+SPHERE = "#cbd5e1" # subtle sphere tint
 
 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.2
+    return (Y / norms) * 1.22
 
-def _species_colors(species: np.ndarray) -> list:
-    return [PALETTE[int(s) % len(PALETTE)] for s in species]
+def make_idle_sphere() -> go.Figure:
+    # empty scatter, only sphere
+    u = np.linspace(0, 2*np.pi, 72)
+    v = np.linspace(0, np.pi, 36)
+    r = 1.22
+    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,
+                        colorscale=[[0, SPHERE],[1, SPHERE]], hoverinfo="skip")
+    layout = go.Layout(
+        paper_bgcolor=BG, plot_bgcolor=BG,
+        title="Architecture Sphere (idle)", titlefont=dict(color="#E5E7EB"),
+        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=720,
+        font=dict(family="Inter, Arial, sans-serif", size=14, color="#E5E7EB")
+    )
+    return go.Figure(data=[sphere], layout=layout)
 
 def make_sphere_figure(points3d: np.ndarray, genomes: List[Genome], gen_idx: int) -> go.Figure:
-    species = np.array([g.species for g in genomes])
-    colors = _species_colors(species)
+    # single-color dots for a sober look
     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.5, color=colors, opacity=0.92),
+        marker=dict(size=7.2, color=DOT, opacity=0.92),
         customdata=custom,
         hovertemplate=(
             "<b>Genome</b><br>"
@@ -221,35 +238,10 @@ def make_sphere_figure(points3d: np.ndarray, genomes: List[Genome], gen_idx: int
             "accuracy=%{customdata[8]:.3f}<extra></extra>"
         )
     )
-
-    # 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.08,
-        showscale=False,
-        colorscale=[[0, "#cbd5e1"], [1, "#cbd5e1"]],
-        hoverinfo="skip"
-    )
-
-    layout = go.Layout(
-        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=720,
-        font=dict(family="Inter, Arial, sans-serif", size=14)
-    )
-    return go.Figure(data=[sphere, scatter], layout=layout)
+    idle = make_idle_sphere()
+    layout = idle.layout.update(title=f"Evo Architecture Sphere — Gen {gen_idx}")
+    fig = go.Figure(data=idle.data + (scatter,), layout=layout)
+    return fig
 
 def make_history_figure(history: List[Tuple[int,float,float]], metric: str) -> go.Figure:
     xs = [h[0] for h in history]
@@ -259,37 +251,36 @@ def make_history_figure(history: List[Tuple[int,float,float]], metric: str) -> g
     else:
         ys = [h[1] for h in history]
         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 = go.Figure(data=[go.Scatter(x=xs, y=ys, mode="lines+markers", line=dict(width=2), marker=dict(color=DOT))])
     fig.update_layout(
-        paper_bgcolor=BG, plot_bgcolor=BG,
+        paper_bgcolor=BG, plot_bgcolor=BG, font=dict(color="#E5E7EB"),
         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)
+        margin=dict(l=30, r=10, t=36, b=30), height=340
     )
+    fig.update_xaxes(gridcolor="#1f2b36"); fig.update_yaxes(gridcolor="#1f2b36")
     return fig
 
 def fig_to_html(fig: go.Figure) -> str:
-    # Inline Plotly JS to avoid any CDN/network dependency in Spaces
-    return pio.to_html(
-        fig,
-        include_plotlyjs=True,   # IMPORTANT: inline JS so the sphere always renders
-        full_html=False,
-        config=dict(displaylogo=False)
-    )
+    return pio.to_html(fig, include_plotlyjs=True, 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 ----------
+# =========================
+# RUNNER
+# =========================
 class EvoRunner:
     def __init__(self):
         self.lock = threading.Lock()
         self.running = False
         self.stop_flag = False
         self.state: Dict[str, Any] = {}
+        # seed the idle sphere immediately
+        idle = fig_to_html(make_idle_sphere())
+        self.state = {"sphere_html": idle, "history_html": fig_to_html(make_history_figure([], "Accuracy")),
+                      "top": [], "best": {}, "gen": 0, "dataset": "Demo (Surrogate)", "metric": "Accuracy"}
 
     def run(self, dataset, pop_size, generations, mutation_rate, explore, exploit, seed, pace_ms, metric_choice):
         rng = random.Random(int(seed))
@@ -302,16 +293,12 @@ class EvoRunner:
             g.fitness, g.acc = fit, acc
 
         history: List[Tuple[int,float,float]] = []
-        best_overall: Optional[Genome] = None
 
         for gen in range(1, generations+1):
             if self.stop_flag: break
 
             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))
+            parents = [min(rng.sample(pop, k=k), key=lambda x: x.fitness) for _ in range(pop_size)]
 
             children = []
             for i in range(0, pop_size, 2):
@@ -331,8 +318,6 @@ class EvoRunner:
             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
-
             history.append((gen, best.fitness, (best.acc if best.acc is not None else float("nan"))))
 
             P = np.stack([g.vector() for g in pop], axis=0)
@@ -341,22 +326,14 @@ class EvoRunner:
             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,
-                    "ffn_mult": t.ffn_mult,
-                    "mem": t.memory_tokens,
-                    "dropout": t.dropout,
-                    "species": t.species,
-                    "params_approx": approx_params(t)
-                } for t in top
-            ]
-            best_card = top_table[0] if len(top_table) else {}
+            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,
+                "ffn_mult": t.ffn_mult, "mem": t.memory_tokens, "dropout": t.dropout,
+                "params_approx": approx_params(t)
+            } for t in top]
+            best_card = top_table[0] if top_table else {}
 
             with self.lock:
                 self.state = {
@@ -376,24 +353,42 @@ class EvoRunner:
         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
 
+    def clear(self):
+        # stop and reset to idle sphere
+        self.stop_flag = True
+        idle = fig_to_html(make_idle_sphere())
+        with self.lock:
+            self.running = False
+            self.state = {"sphere_html": idle, "history_html": fig_to_html(make_history_figure([], "Accuracy")),
+                          "top": [], "best": {}, "gen": 0, "dataset": "Demo (Surrogate)", "metric": "Accuracy"}
+
 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))
+    return (gr.update(interactive=False), gr.update(interactive=True), gr.update(interactive=False))
 
 def stop_evo():
     runner.stop()
-    return (gr.update(interactive=True), gr.update(interactive=False))
+    return (gr.update(interactive=True), gr.update(interactive=False), gr.update(interactive=True))
+
+def clear_evo():
+    runner.clear()
+    # return updated visuals + reset buttons
+    sphere_html, history_html, stats_md, df = poll_state()
+    return sphere_html, history_html, stats_md, df, gr.update(interactive=True), gr.update(interactive=False), gr.update(interactive=True)
 
 def poll_state():
     with runner.lock:
         s = runner.state.copy()
-    sphere_html = s.get("sphere_html", "")
-    history_html = s.get("history_html", "")
+    sphere_html = s.get("sphere_html", fig_to_html(make_idle_sphere()))
+    history_html = s.get("history_html", fig_to_html(make_history_figure([], "Accuracy")))
     best = s.get("best", {})
     gen = s.get("gen", 0)
     dataset = s.get("dataset", "Demo (Surrogate)")
@@ -411,7 +406,7 @@ def poll_state():
             f"**~Params (rough):** {best.get('params_approx'):,}"
         )
     else:
-        stats_md = "Waiting… click **Start Evolution**."
+        stats_md = "Ready. Press **Start** to evolve, or **Clear** anytime."
     df = pd.DataFrame(top)
     return sphere_html, history_html, stats_md, df
 
@@ -424,10 +419,12 @@ def export_snapshot():
         f.write(payload)
     return path
 
-# ---------- Build UI (minimal layout) ----------
-with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Soft()) as demo:
+# =========================
+# BUILD UI
+# =========================
+with gr.Blocks(css=CUSTOM_CSS) as demo:
     with gr.Column(elem_id="header"):
-        gr.Markdown("## Evo Playground — Minimal Live Evolution (PIQA / HellaSwag accuracy)")
+        gr.Markdown("### Evo Playground — Live Evolution (clean dark)")
 
     with gr.Row():
         with gr.Column(scale=1, elem_classes=["controls"]):
@@ -436,7 +433,7 @@ with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Soft()) as demo:
                     label="Dataset",
                     choices=["Demo (Surrogate)", "PIQA (Phase 2)", "HellaSwag (Phase 2)"],
                     value="Demo (Surrogate)",
-                    info="PIQA/HellaSwag compute real proxy accuracy; Demo is a fast surrogate."
+                    info="PIQA/HellaSwag compute 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")
@@ -445,14 +442,16 @@ with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Soft()) as demo:
                     explore = gr.Slider(0.0, 1.0, value=0.35, step=0.05, label="Exploration")
                     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)")
+                pace = gr.Slider(0, 1000, value=120, step=10, label="Pace (ms)")
                 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")
+                    start = gr.Button("▶ Start", variant="primary")
+                    stop = gr.Button("⏹ Stop", variant="secondary", interactive=False)
+                    clear = gr.Button("↺ Clear", variant="secondary")
 
             with gr.Group(elem_classes=["panel"]):
-                stats_md = gr.Markdown("Waiting…", elem_id="stats")
+                stats_md = gr.Markdown("Ready. Press **Start** to evolve, or **Clear** anytime.", elem_id="stats")
 
             with gr.Group(elem_classes=["panel"]):
                 export_btn = gr.Button("Export Snapshot (JSON)")
@@ -466,12 +465,13 @@ with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Soft()) as demo:
             with gr.Group(elem_classes=["panel"]):
                 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, metric_choice], [start, stop])
-    stop.click(stop_evo, [], [start, stop])
+    # wiring
+    start.click(start_evo, [dataset, pop, gens, mut, explore, exploit, seed, pace, metric_choice], [start, stop, clear])
+    stop.click(stop_evo, [], [start, stop, clear])
+    clear.click(clear_evo, [], [sphere_html, hist_html, stats_md, top_df, start, stop, clear])
     export_btn.click(export_snapshot, [], [export_file])
 
-    # Initial paint + polling
+    # 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])