app.py

import gradio as gr
from evo_transformer import 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()
    evaluation = evo.evaluate()

    history_table = "\n".join(
        [f"Gen {i}: {cfg}" for i, cfg in enumerate(history[-generations:])]
    )

    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```"
    )

# 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()

    evolve_btn.click(
        evolve_and_display,
        inputs=[generations_slider],
        outputs=[output_box],
    )

# Launch the app
demo.launch()