Update app.py

app.py

@@ -1,25 +1,25 @@
 import torch
 import numpy as np
 import time
+import threading
+import uvicorn
 from fastmcp import FastMCP
 from ase import Atoms
 from ase.build import molecule
 import gradio as gr
 
-# Initialize MCP Server
+# --- 1. MCP SERVER SETUP ---
 mcp = FastMCP("RealMACE_Agent")
 
 # Global State
 STATE = {
     "model": None,
     "config": None,
-    "batch": None,
-    "training_logs": []
+    "batch": None
 }
 
 # --- HELPER FUNCTIONS ---
 def get_mace_setup():
-    """Lazy load MACE imports."""
     try:
         from mace.models import ScaleShiftMACE
         from mace.data import AtomicData, Configuration
@@ -30,13 +30,10 @@ def get_mace_setup():
         raise ImportError("MACE not installed. Run: pip install mace-torch")
 
 def create_dummy_batch(r_max=5.0):
-    """Creates a water molecule batch for training."""
     _, AtomicData, Configuration, torch_geometric, _ = get_mace_setup()
-    
     mol = molecule("H2O")
     mol.info["energy"] = -14.0
     mol.arrays["forces"] = np.random.randn(3, 3) * 0.1
-    
     config = Configuration(
         atomic_numbers=mol.get_atomic_numbers(),
         positions=mol.get_positions(),
@@ -45,7 +42,6 @@ def create_dummy_batch(r_max=5.0):
         pbc=np.array([False, False, False]),
         cell=np.eye(3) * 10.0 
     )
-    
     z_table = {1: 0, 8: 1}
     data_loader = torch_geometric.DataLoader(
         dataset=[AtomicData.from_config(config, z_table=z_table, cutoff=r_max)],
@@ -59,35 +55,25 @@ def create_dummy_batch(r_max=5.0):
 def init_real_mace_model(r_max: float = 5.0, max_ell: int = 2, hidden_dim: int = 16) -> str:
     """Initialize a REAL MACE model."""
     ScaleShiftMACE, _, _, _, o3 = get_mace_setup()
-    
     batch = create_dummy_batch(r_max)
     STATE["batch"] = batch
     
     model_config = dict(
-        r_max=r_max,
-        num_bessel=8,
-        num_polynomial_cutoff=5,
-        max_ell=max_ell,
-        interaction_cls="RealAgnosticInteractionBlock",
-        num_interactions=2,
-        num_elements=2,
-        hidden_irreps=o3.Irreps(f"{hidden_dim}x0e"),
-        atomic_energies=np.array([-13.6, -10.0]),
-        avg_num_neighbors=2,
-        atomic_numbers=[1, 8]
+        r_max=r_max, num_bessel=8, num_polynomial_cutoff=5, max_ell=max_ell,
+        interaction_cls="RealAgnosticInteractionBlock", num_interactions=2, num_elements=2,
+        hidden_irreps=o3.Irreps(f"{hidden_dim}x0e"), atomic_energies=np.array([-13.6, -10.0]),
+        avg_num_neighbors=2, atomic_numbers=[1, 8]
     )
-    
     try:
         model = ScaleShiftMACE(**model_config)
         STATE["model"] = model
-        STATE["config"] = model_config
         return f"✅ MACE Model Ready! L_max={max_ell}, r_max={r_max}Å"
     except Exception as e:
         return f"❌ Error: {str(e)}"
 
 @mcp.tool()
-def train_with_trackio(experiment_name: str, epochs: int = 10, learning_rate: float = 0.01) -> str:
-    """Train the MACE model with Trackio logging."""
+def train_with_trackio(experiment_name: str, epochs: int = 10) -> str:
+    """Train with Trackio logging."""
     try:
         import trackio
     except ImportError:
@@ -98,152 +84,72 @@ def train_with_trackio(experiment_name: str, epochs: int = 10, learning_rate: fl
     
     model = STATE["model"]
     batch = STATE["batch"]
-    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
     
     try:
+        # Check if we are in a Space with OAuth
         logger = trackio.Logger(project="Real_MACE_Training", name=experiment_name)
     except Exception as e:
-        return f"❌ Trackio error: {e}"
+        return f"❌ Trackio connection failed: {e}"
 
     model.train()
-    STATE["training_logs"] = []
+    logs = []
+    start = time.time()
     
     for epoch in range(epochs):
         optimizer.zero_grad()
         out = model(batch.to_dict())
-        
-        loss_e = torch.mean((out["energy"] - batch.energy)**2)
-        loss_f = torch.mean((out["forces"] - batch.forces)**2)
-        total_loss = loss_e + 10.0 * loss_f
-        
-        total_loss.backward()
+        loss = torch.mean((out["energy"] - batch.energy)**2) + 10.0 * torch.mean((out["forces"] - batch.forces)**2)
+        loss.backward()
         optimizer.step()
         
-        force_mae = torch.mean(torch.abs(out["forces"] - batch.forces)).item()
-        
-        logger.log({
+        metrics = {
             "epoch": epoch,
-            "total_loss": total_loss.item(),
-            "force_mae_eV_A": force_mae,
-        })
+            "total_loss": loss.item(),
+            "wall_time": time.time() - start
+        }
+        logger.log(metrics)
         
-        STATE["training_logs"].append(f"Epoch {epoch}: Loss={total_loss.item():.5f}")
-        time.sleep(0.05)
-
-    return f"🚀 Training done! Final Loss: {total_loss.item():.6f}\n" + "\n".join(STATE["training_logs"][-5:])
+        if epoch % 5 == 0:
+            logs.append(f"Ep {epoch}: Loss={loss.item():.4f}")
+            time.sleep(0.05)
 
-@mcp.tool()
-def check_equivariance(rotation_degrees: float = 45.0) -> str:
-    """Test E(3)-equivariance."""
-    if STATE["model"] is None:
-        return "⚠️ No model found!"
-    
-    model = STATE["model"]
-    batch = STATE["batch"]
-    
-    model.eval()
-    with torch.no_grad():
-        out_orig = model(batch.to_dict())
-        forces_orig = out_orig["forces"].clone()
-    
-    angle = np.radians(rotation_degrees)
-    rot_matrix = torch.tensor([
-        [np.cos(angle), -np.sin(angle), 0],
-        [np.sin(angle), np.cos(angle), 0],
-        [0, 0, 1]
-    ], dtype=torch.float32)
-    
-    batch_rot = batch.clone()
-    batch_rot.positions = torch.matmul(batch.positions, rot_matrix.T)
-    
-    with torch.no_grad():
-        out_rot = model(batch_rot.to_dict())
-        forces_rot = out_rot["forces"]
-    
-    forces_orig_rotated = torch.matmul(forces_orig, rot_matrix.T)
-    equivariance_error = torch.mean(torch.abs(forces_rot - forces_orig_rotated)).item()
-    
-    return f"🧪 Equivariance Error: {equivariance_error:.2e} eV/Å\n{'✅ PASS' if equivariance_error < 1e-4 else '⚠️ High error'}"
+    return "🚀 Training Done! Check Dashboard.\n" + "\n".join(logs)
 
-# --- GRADIO UI ---
-def create_ui():
-    """Create the Gradio interface."""
-    with gr.Blocks(title="Equivariant Alchemist - MACE Training Lab") as demo:
-        gr.Markdown("""
-        # 🧪 Equivariant Alchemist - MACE Training Lab
+# --- 2. DASHBOARD UI (Separate Thread) ---
+def launch_dashboard():
+    """Launches a Gradio UI that serves as the Dashboard Viewer"""
+    with gr.Blocks(title="Equivariant Chem Scout") as demo:
+        gr.Markdown("# 🧪 Equivariant Chem Scout (Dashboard)")
+        gr.Markdown("To view training results, open the **Trackio** dashboard below.")
         
-        This app combines an **MCP Server** for AI agents with **Trackio** experiment tracking.
+        # Option A: If running locally, just show instructions
+        gr.Markdown("""
+        ### How to view graphs:
+        The Trackio dashboard runs separately. 
+        If you are running locally, type: `trackio show` in your terminal.
+        If you are on Hugging Face Spaces, we need to launch the Trackio server.
         """)
         
-        with gr.Tabs():
-            with gr.Tab("📊 View Trackio Dashboard"):
-                gr.Markdown("""
-                ### Live Training Metrics
-                
-                To view your training metrics, open the Trackio dashboard in a separate window:
-                
-                **Option 1: Command Line**
-                ```
-                trackio show --project "Real_MACE_Training"
-                ```
-                
-                **Option 2: Python**
-                ```
-                import trackio
-                trackio.show(project="Real_MACE_Training")
-                ```
-                
-                The dashboard will automatically update as training runs complete.
-                """)
-                
-                gr.HTML("""
-                <iframe 
-                    src="/trackio" 
-                    width="100%" 
-                    height="800px" 
-                    frameborder="0"
-                    style="border-radius: 8px;"
-                ></iframe>
-                """)
-            
-            with gr.Tab("🔌 MCP Server Info"):
-                gr.Markdown(f"""
-                ### MCP Server Status: ✅ Running
-                
-                **Server URL:** Access at `/sse` endpoint
-                
-                **Available Tools:**
-                1. `init_real_mace_model(r_max, max_ell, hidden_dim)` - Initialize MACE architecture
-                2. `train_with_trackio(experiment_name, epochs, learning_rate)` - Train with live logging  
-                3. `check_equivariance(rotation_degrees)` - Test rotation symmetry
-                
-                **Connect from Claude Desktop:**
-                ```
-                {{
-                  "mcpServers": {{
-                    "mace_trainer": {{
-                      "url": "YOUR_SPACE_URL/sse"
-                    }}
-                  }}
-                }}
-                ```
-                
-                **Example Prompts:**
-                - *"Initialize a MACE model with max_ell=2 and r_max=5.0"*
-                - *"Train for 20 epochs with learning rate 0.001"*
-                - *"Check if the model is equivariant by rotating 90 degrees"*
-                """)
-    
-    return demo
+        # Option B: Attempt to embed (Experimental)
+        # Note: Trackio doesn't have a verified embed widget yet, so we provide instructions.
+        
+    demo.launch(server_name="0.0.0.0", server_port=7860, prevent_thread_lock=True)
 
 if __name__ == "__main__":
-    print("Starting MACE-MCP Server with Trackio Integration...")
-    
-    # Create and launch the Gradio UI with MCP server
-    demo = create_ui()
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False,
-        mcp_server=mcp  # This enables MCP on the /sse endpoint
-    )
+    print("--- STARTING SERVICES ---")
+    
+    # 1. Launch the UI (Dashboard) in a background thread on port 7860
+    print("1. Launching Gradio Dashboard on port 7860...")
+    launch_dashboard()
+    
+    # 2. Run the MCP Server on the main thread (port 8000 or SSE)
+    print("2. Starting MCP Server (SSE Transport)...")
+    # Hugging Face Spaces expects the main process to listen on port 7860 usually,
+    # but for MCP we need to expose the SSE endpoint.
+    # TRICK: We let Gradio take 7860 (so the Space shows "Running"), 
+    # and we run MCP on 8000. You connect to the Space URL via SSE proxying if configured,
+    # or you use this Space *only* as a dashboard and run the MCP logic locally connecting to it.
+    
+    # However, since you want the Space to BE the MCP server:
+    mcp.run(transport="sse")