Update app.py

app.py

@@ -1,48 +1,50 @@
-# app.py
 import gradio as gr
 from evo_transformer import EvoTransformer
-import matplotlib.pyplot as plt
 
-def run_evolution(generations):
-    evo = EvoTransformer()
+# Initialize global EvoTransformer object
+evo = EvoTransformer()
+
+# Define Gradio functions
+def evolve_and_display(generations):
     evo.evolve(generations)
+    latest_config = evo.config
     history = evo.get_history()
-    results = evo.evaluate()
-
-    # Format history for output
-    trait_logs = ""
-    for i, config in enumerate(history):
-        trait_logs += f"Generation {i}: {config}\n"
+    evaluation = evo.evaluate()
 
-    # Plot number of attention heads over generations
-    heads = [conf["attention_heads"] for conf in history]
-    layers = [conf["layers"] for conf in history]
-    ffn = [conf["ffn_dim"] for conf in history]
+    history_table = "\n".join(
+        [f"Gen {i}: {cfg}" for i, cfg in enumerate(history[-generations:])]
+    )
 
-    fig, ax = plt.subplots()
-    ax.plot(range(len(heads)), heads, label="Attention Heads")
-    ax.plot(range(len(layers)), layers, label="Layers")
-    ax.plot(range(len(ffn)), ffn, label="FFN Dim")
-    ax.set_title("Evolution of Traits")
-    ax.set_xlabel("Generation")
-    ax.set_ylabel("Value")
-    ax.legend()
+    return (
+        f"### Final Evolved Configuration\n"
+        f"**Layers**: {latest_config['layers']}  \n"
+        f"**Attention Heads**: {latest_config['attention_heads']}  \n"
+        f"**FFN Dimension**: {latest_config['ffn_dim']}  \n"
+        f"**Dropout**: {latest_config['dropout']}  \n"
+        f"**Memory**: {latest_config['memory']}  \n\n"
+        f"### Evaluation\n"
+        f"**Simulated Accuracy**: {evaluation['accuracy']}  \n"
+        f"**Estimated Parameters**: {evaluation['params']}M  \n\n"
+        f"### Trait History (last {generations} generations)\n"
+        f"```\n{history_table}\n```"
+    )
 
-    return trait_logs, results["accuracy"], round(results["params"], 2), fig
+# Build Gradio Interface
+with gr.Blocks(title="EvoTransformer Demo") as demo:
+    gr.Markdown("# 🧬 EvoTransformer")
+    gr.Markdown("Simulate in-training architectural evolution of a Transformer model.")
+    
+    with gr.Row():
+        generations_slider = gr.Slider(1, 10, value=3, label="Generations to Evolve")
+        evolve_btn = gr.Button("Evolve")
+    
+    output_box = gr.Markdown()
 
-# Gradio UI
-demo = gr.Interface(
-    fn=run_evolution,
-    inputs=gr.Slider(minimum=1, maximum=10, step=1, label="Generations"),
-    outputs=[
-        gr.Textbox(label="Evolution History"),
-        gr.Number(label="Simulated Accuracy"),
-        gr.Number(label="Estimated Parameters (M)"),
-        gr.Plot(label="Trait Evolution Plot"),
-    ],
-    title="🧬 EvoTransformer Demo",
-    description="An evolving Transformer that mutates architecture traits during training. Watch the architecture change in real time!"
-)
+    evolve_btn.click(
+        evolve_and_display,
+        inputs=[generations_slider],
+        outputs=[output_box],
+    )
 
-if __name__ == "__main__":
-    demo.launch()
+# Launch the app
+demo.launch()