Create app.py

app.py

@@ -0,0 +1,380 @@
+# app.py
+import math, json, random, time, threading, io, os
+from dataclasses import dataclass, asdict
+from typing import List, Tuple, Dict, Any
+import numpy as np
+import plotly.graph_objs as go
+import gradio as gr
+
+# =========================
+# UX THEME & STYLES
+# =========================
+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}
+"""
+
+# =========================
+# GENOME & EVOLUTION CORE
+# =========================
+@dataclass
+class Genome:
+    d_model: int
+    n_layers: int
+    n_heads: int
+    ffn_mult: float
+    memory_tokens: int
+    dropout: float
+    species: int = 0
+    fitness: float = float("inf")
+
+    def vector(self) -> np.ndarray:
+        # Normalized structural vector (0..1)
+        return np.array([
+            self.d_model / 1024.0,
+            self.n_layers / 24.0,
+            self.n_heads / 32.0,
+            self.ffn_mult / 8.0,
+            self.memory_tokens / 64.0,
+            self.dropout / 0.5
+        ], dtype=np.float32)
+
+def random_genome(rng: random.Random) -> Genome:
+    return Genome(
+        d_model=rng.choice([256, 384, 512, 640]),
+        n_layers=rng.choice([4, 6, 8, 10, 12]),
+        n_heads=rng.choice([4, 6, 8, 10, 12]),
+        ffn_mult=rng.choice([2.0, 3.0, 4.0, 6.0]),
+        memory_tokens=rng.choice([0, 4, 8, 16]),
+        dropout=rng.choice([0.0, 0.05, 0.1, 0.15]),
+        species=rng.randrange(5)
+    )
+
+def mutate(g: Genome, rng: random.Random, rate: float) -> Genome:
+    g = Genome(**asdict(g))
+    if rng.random() < rate: g.d_model = rng.choice([256, 384, 512, 640])
+    if rng.random() < rate: g.n_layers = rng.choice([4, 6, 8, 10, 12])
+    if rng.random() < rate: g.n_heads = rng.choice([4, 6, 8, 10, 12])
+    if rng.random() < rate: g.ffn_mult = rng.choice([2.0, 3.0, 4.0, 6.0])
+    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")
+    return g
+
+def crossover(a: Genome, b: Genome, rng: random.Random) -> Genome:
+    return Genome(
+        d_model = a.d_model if rng.random()<0.5 else b.d_model,
+        n_layers = a.n_layers if rng.random()<0.5 else b.n_layers,
+        n_heads = a.n_heads if rng.random()<0.5 else b.n_heads,
+        ffn_mult = a.ffn_mult if rng.random()<0.5 else b.ffn_mult,
+        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 HOOK (Phase 1: fast surrogate)
+# Swap this later for real PIQA/HellaSwag evaluation
+# =========================
+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))
+
+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)
+
+# =========================
+# PROJECTION & VIZ
+# =========================
+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
+
+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
+    ]
+
+    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'
+    )
+
+    # 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)
+
+    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
+    )
+    return go.Figure(data=[sphere, scatter], layout=layout)
+
+def make_history_figure(history: List[Tuple[int,float]]) -> go.Figure:
+    xs = [h[0] for h in history]
+    ys = [h[1] for h in history]
+    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))
+    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
+    # tiny bump for memory tokens pathways (illustrative only)
+    total += 1000 * g.memory_tokens
+    return int(total)
+
+# =========================
+# ORCHESTRATOR
+# =========================
+class EvoRunner:
+    def __init__(self):
+        self.lock = threading.Lock()
+        self.running = False
+        self.stop_flag = False
+        self.state: Dict[str, Any] = {}
+
+    def run(self, dataset, pop_size, generations, mutation_rate, explore, exploit, seed, pace_ms):
+        rng = random.Random(int(seed))
+        self.stop_flag = False
+        self.running = True
+
+        pop: List[Genome] = [random_genome(rng) for _ in range(pop_size)]
+        # initial eval
+        for g in pop:
+            g.fitness = fitness_hook(g, dataset, explore)
+
+        history: List[Tuple[int,float]] = []
+        best_overall: Genome | None = None
+
+        for gen in range(1, generations+1):
+            if self.stop_flag: break
+
+            # Selection: tournament size depends on exploitation
+            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]
+                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
+            for c in children:
+                c.fitness = fitness_hook(c, dataset, explore)
+
+            # 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
+
+            history.append((gen, best.fitness))
+
+            # 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)
+            top = sorted(pop, key=lambda x: x.fitness)[: min(12, len(pop))]
+            top_table = [
+                {
+                    "gen": gen,
+                    "fitness": round(t.fitness, 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 {}
+
+            with self.lock:
+                self.state = {
+                    "sphere": sphere_fig,
+                    "history": hist_fig,
+                    "top": top_table,
+                    "best": best_card,
+                    "gen": gen,
+                    "dataset": dataset
+                }
+
+            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()
+
+# =========================
+# GRADIO 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))
+    return (gr.update(interactive=False), gr.update(interactive=True))
+
+def stop_evo():
+    runner.stop()
+    return (gr.update(interactive=True), gr.update(interactive=False))
+
+def poll_state():
+    with runner.lock:
+        s = runner.state.copy()
+    # Defaults before first run
+    sphere = s.get("sphere", go.Figure())
+    history = s.get("history", go.Figure())
+    best = s.get("best", {})
+    gen = s.get("gen", 0)
+    dataset = s.get("dataset", "Demo (Surrogate)")
+    top = s.get("top", [])
+    # Build stats Markdown
+    if best:
+        stats_md = (
+            f"**Dataset:** {dataset}  \n"
+            f"**Generation:** {gen}  \n"
+            f"**Best fitness:** {best.get('fitness','–')}  \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"
+            f"**~Params (rough):** {best.get('params_approx'):,}"
+        )
+    else:
+        stats_md = "Waiting… click **Start Evolution**."
+
+    # Dataframe rows
+    import pandas as pd
+    df = pd.DataFrame(top)
+    return sphere, history, stats_md, df
+
+def export_snapshot():
+    with runner.lock:
+        payload = json.dumps(runner.state, default=lambda o: o, indent=2)
+    # Write to a temp file so user can download
+    path = "evo_snapshot.json"
+    with open(path, "w", encoding="utf-8") as f:
+        f.write(payload)
+    return path
+
+# =========================
+# BUILD UI
+# =========================
+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)."
+        )
+
+    with gr.Row():
+        # LEFT: Controls
+        with gr.Column(scale=1):
+            with gr.Group():
+                dataset = gr.Dropdown(
+                    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."
+                )
+                pop = gr.Slider(8, 80, value=24, step=2, label="Population size")
+                gens = gr.Slider(5, 200, value=60, step=1, label="Max generations")
+                mut = gr.Slider(0.05, 0.9, value=0.25, step=0.01, label="Mutation rate")
+                with gr.Row():
+                    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)")
+                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…")
+
+                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="Best Fitness History")
+            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])
+    stop.click(stop_evo, [], [start, stop])
+    export_btn.click(export_snapshot, [], [export_file])
+
+    # Live polling
+    demo.load(poll_state, None, [sphere_plot, hist_plot, stats_md, top_df], every=0.7)
+
+if __name__ == "__main__":
+    demo.launch()