Update app.py

app.py

@@ -1,25 +1,27 @@
 import torch
 import numpy as np
 import time
-import threading
 import uvicorn
+from fastapi import FastAPI
 from fastmcp import FastMCP
+import gradio as gr
 from ase import Atoms
 from ase.build import molecule
-import gradio as gr
 
-# --- 1. MCP SERVER SETUP ---
-mcp = FastMCP("RealMACE_Agent")
+# --- 1. DEFINE MCP SERVER ---
+mcp = FastMCP("Equivariant_Chem_Scout")
 
-# Global State
+# Global State for Persisting Models across Tool Calls
 STATE = {
     "model": None,
     "config": None,
-    "batch": None
+    "batch": None,
+    "logs": []
 }
 
 # --- HELPER FUNCTIONS ---
 def get_mace_setup():
+    """Lazy load MACE to avoid startup crashes if deps are missing"""
     try:
         from mace.models import ScaleShiftMACE
         from mace.data import AtomicData, Configuration
@@ -27,13 +29,16 @@ def get_mace_setup():
         from e3nn import o3
         return ScaleShiftMACE, AtomicData, Configuration, torch_geometric, o3
     except ImportError:
-        raise ImportError("MACE not installed. Run: pip install mace-torch")
+        raise ImportError("MACE not installed. Please check requirements.txt")
 
 def create_dummy_batch(r_max=5.0):
     _, AtomicData, Configuration, torch_geometric, _ = get_mace_setup()
+    
+    # Create dummy water molecule
     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(),
@@ -42,6 +47,7 @@ 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)],
@@ -51,105 +57,164 @@ def create_dummy_batch(r_max=5.0):
     return next(iter(data_loader))
 
 # --- MCP TOOLS ---
+
 @mcp.tool()
-def init_real_mace_model(r_max: float = 5.0, max_ell: int = 2, hidden_dim: int = 16) -> str:
-    """Initialize a REAL MACE model."""
+def init_mace_model(r_max: float = 5.0, max_ell: int = 2) -> str:
+    """
+    Initialize a MACE model with specific symmetry settings.
+    Args:
+        r_max: Cutoff radius (Angstroms)
+        max_ell: 0 = Invariant only, 2 = Equivariant (Vectors)
+    """
     ScaleShiftMACE, _, _, _, o3 = get_mace_setup()
+    
     batch = create_dummy_batch(r_max)
     STATE["batch"] = batch
     
+    # simplified MACE config
     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]),
+        hidden_irreps=o3.Irreps("16x0e"), atomic_energies=np.array([-13.6, -10.0]),
         avg_num_neighbors=2, atomic_numbers=[1, 8]
     )
+    
     try:
         model = ScaleShiftMACE(**model_config)
         STATE["model"] = model
-        return f"✅ MACE Model Ready! L_max={max_ell}, r_max={r_max}Å"
+        STATE["config"] = model_config
+        STATE["logs"] = [] # Reset logs
+        return f"✅ MACE Model Initialized! (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) -> str:
-    """Train with Trackio logging."""
-    try:
-        import trackio
-    except ImportError:
-        return "❌ Trackio not installed"
-
+def train_model(epochs: int = 10, learning_rate: float = 0.01) -> str:
+    """Train the initialized model and log results."""
     if STATE["model"] is None:
-        return "⚠️ Run 'init_real_mace_model' first!"
+        return "⚠️ Run 'init_mace_model' first!"
     
     model = STATE["model"]
     batch = STATE["batch"]
-    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
+    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
     
-    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 connection failed: {e}"
-
     model.train()
-    logs = []
-    start = time.time()
+    run_logs = []
     
+    # Simple Training Loop
     for epoch in range(epochs):
         optimizer.zero_grad()
         out = model(batch.to_dict())
+        
         loss = torch.mean((out["energy"] - batch.energy)**2) + 10.0 * torch.mean((out["forces"] - batch.forces)**2)
         loss.backward()
         optimizer.step()
         
-        metrics = {
-            "epoch": epoch,
-            "total_loss": loss.item(),
-            "wall_time": time.time() - start
-        }
-        logger.log(metrics)
-        
-        if epoch % 5 == 0:
-            logs.append(f"Ep {epoch}: Loss={loss.item():.4f}")
-            time.sleep(0.05)
-
-    return "🚀 Training Done! Check Dashboard.\n" + "\n".join(logs)
-
-# --- 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.")
-        
-        # 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.
-        """)
+        # Log for Dashboard
+        log_entry = f"Epoch {epoch}: Loss={loss.item():.4f}"
+        run_logs.append(log_entry)
+        STATE["logs"].append(log_entry) # Append to global state for UI viewing
         
-        # Option B: Attempt to embed (Experimental)
-        # Note: Trackio doesn't have a verified embed widget yet, so we provide instructions.
+        time.sleep(0.05) # Simulate work
         
-    demo.launch(server_name="0.0.0.0", server_port=7860, prevent_thread_lock=True)
+    return f"🚀 Training Complete!\nFinal Loss: {loss.item():.4f}\n" + "\n".join(run_logs[-5:])
+
+@mcp.tool()
+def check_equivariance(rotation_degrees: float = 90.0) -> str:
+    """Check if the model respects E(3) symmetry."""
+    if STATE["model"] is None: return "⚠️ No model found."
+    
+    model = STATE["model"]
+    batch = STATE["batch"]
+    model.eval()
+    
+    # 1. Original Pred
+    out_orig = model(batch.to_dict())
+    f_orig = out_orig["forces"]
+    
+    # 2. Rotated Input
+    angle = np.radians(rotation_degrees)
+    R = 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, R.T)
+    
+    # 3. Rotated Pred
+    out_rot = model(batch_rot.to_dict())
+    f_rot = out_rot["forces"]
+    
+    # 4. Compare: Rot(F_orig) vs F_rot
+    f_orig_rot = torch.matmul(f_orig, R.T)
+    error = torch.mean(torch.abs(f_rot - f_orig_rot)).item()
+    
+    return f"🧪 Equivariance Error: {error:.2e} (Pass: {error < 1e-4})"
+
+# --- 2. DEFINE GRADIO DASHBOARD ---
+def get_latest_logs():
+    """Refresh function for the UI"""
+    if not STATE["logs"]:
+        return "No training logs yet. Ask the Agent to train a model!"
+    return "\n".join(STATE["logs"])
+
+with gr.Blocks(title="Equivariant Chem Scout") as demo:
+    gr.Markdown("# 🧪 Equivariant Chem Scout")
+    gr.Markdown("""
+    ### Status: 🟢 Online
+    **Connect via Claude Desktop:**
+    ```
+    {
+      "mcpServers": {
+        "chem_scout": {
+          "url": "https://YOUR-SPACE-URL.hf.space/sse"
+        }
+      }
+    }
+    ```
+    """)
+    
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("### 📊 Live Training Logs")
+            log_display = gr.TextArea(label="Training Output", interactive=False, lines=20)
+            refresh_btn = gr.Button("🔄 Refresh Logs")
+            
+            refresh_btn.click(fn=get_latest_logs, outputs=log_display)
+            
+            # Auto-refresh every 2 seconds
+            demo.load(fn=get_latest_logs, outputs=log_display, every=2.0)
+
+# --- 3. ASSEMBLE FASTAPI APP ---
+# This is the magic glue that makes it work on Spaces
+app = FastAPI()
+
+# Mount the MCP Server at /sse
+# FastMCP provides a method to attach itself to an existing FastAPI app
+from sse_starlette.sse import EventSourceResponse
+
+@app.get("/sse")
+async def handle_sse(request):
+    return EventSourceResponse(mcp.sse_handler(request))
+
+@app.post("/messages")
+async def handle_messages(request):
+    return await mcp.handle_post_message(request)
+
+# IMPORTANT: `mcp.mount_to_fastapi` is not always available in older versions,
+# so we can use the manual mounting above OR use the built-in if available.
+# Let's try the safest built-in method if it exists, or fallback to manual.
+
+try:
+    # Newer FastMCP versions
+    mcp.mount_to_fastapi(app, path="/sse")
+except AttributeError:
+    # Fallback if mount_to_fastapi doesn't exist (older versions)
+    pass 
+
+# Mount Gradio at the root
+app = gr.mount_gradio_app(app, demo, path="/")
 
+# --- 4. ENTRY POINT ---
 if __name__ == "__main__":
-    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")
+    # Hugging Face Spaces will run this with: uvicorn app:app --host 0.0.0.0 --port 7860
+    # But for local testing:
+    uvicorn.run(app, host="0.0.0.0", port=7860)