Update for ZeroGPU with @spaces.GPU decorator

README.md

@@ -7,12 +7,12 @@ sdk: gradio
 sdk_version: 4.44.0
 app_file: app.py
 pinned: false
-hardware: h200
+hardware: zero-a10g
 ---
 
 # MANIFOLD - CS2 Cheat Detection Training
 
-Train the MANIFOLD (Motor-Aware Neural Inference for Faithfulness Of Latent Dynamics) model for CS2 cheat detection.
+Train the MANIFOLD model for CS2 cheat detection using ZeroGPU (H200).
 
 ## Features
 - Synthetic data generation with realistic player behavior

app.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-"""MANIFOLD Training Interface for Hugging Face Spaces."""
+"""MANIFOLD Training Interface for Hugging Face Spaces with ZeroGPU."""
 
 import gradio as gr
 import torch
@@ -7,10 +7,8 @@ import numpy as np
 import json
 import time
 from pathlib import Path
-from threading import Thread
-import queue
+import spaces
 
-# Add src to path for manifold imports
 import sys
 sys.path.insert(0, str(Path(__file__).parent / "src"))
 
@@ -18,82 +16,47 @@ from manifold import MANIFOLDLite
 from manifold.config import ModelConfig, TrainingConfig
 from manifold.data.generator import SyntheticDataGenerator
 from manifold.data.dataset import MANIFOLDDataset, create_dataloader
-from manifold.training.trainer import MANIFOLDTrainer
-from manifold.training.callbacks import Callback, ProgressCallback
+from manifold.training.trainer import train_epoch, validate
 from manifold.training.curriculum import CurriculumScheduler
+from manifold.training.losses import compute_total_loss
 
-# Global state
-training_log = queue.Queue()
-is_training = False
 current_model = None
 
 
-class GradioCallback(Callback):
-    """Callback that logs to Gradio interface."""
-    
-    def on_epoch_end(self, trainer, epoch_info):
-        epoch = epoch_info["epoch"]
-        stage = epoch_info.get("stage", {}).get("stage_name", "")
-        train_loss = epoch_info.get("train", {}).get("loss", 0)
-        val_loss = epoch_info.get("val", {}).get("loss", 0)
-        val_acc = epoch_info.get("val", {}).get("accuracy", 0)
-        lr = epoch_info.get("lr", 0)
-        
-        log_entry = {
-            "epoch": epoch + 1,
-            "stage": stage,
-            "train_loss": f"{train_loss:.4f}",
-            "val_loss": f"{val_loss:.4f}",
-            "val_acc": f"{val_acc:.4f}",
-            "lr": f"{lr:.2e}",
-        }
-        training_log.put(log_entry)
-
-
 def get_device_info():
-    """Get GPU information."""
     if torch.cuda.is_available():
-        gpu_name = torch.cuda.get_device_name(0)
-        gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1e9
-        return f"GPU: {gpu_name} ({gpu_mem:.1f} GB)"
-    return "CPU only (no GPU detected)"
+        return f"GPU: {torch.cuda.get_device_name(0)} ({torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB)"
+    return "CPU (GPU will be allocated when training starts)"
 
 
 def generate_data(num_legit, num_cheaters, seed, progress=gr.Progress()):
-    """Generate synthetic training data."""
     progress(0, desc="Initializing generator...")
-    
-    generator = SyntheticDataGenerator(seed=seed, engagements_per_session=200)
+    generator = SyntheticDataGenerator(seed=int(seed), engagements_per_session=200)
     
     all_features = []
     all_labels = []
     total = num_legit + num_cheaters
     
-    # Generate legit players
     for i in progress.tqdm(range(num_legit), desc="Generating legit players"):
         session = generator.generate_player(is_cheater=False)
         all_features.append(session.to_tensor())
         all_labels.append(0)
     
-    # Generate cheaters
     for i in progress.tqdm(range(num_cheaters), desc="Generating cheaters"):
         session = generator.generate_player(is_cheater=True)
         all_features.append(session.to_tensor())
         all_labels.append(2)
     
-    # Convert and shuffle
     features = np.array(all_features)
     labels = np.array(all_labels)
     
-    rng = np.random.default_rng(seed)
+    rng = np.random.default_rng(int(seed))
     indices = rng.permutation(total)
     features = features[indices]
     labels = labels[indices]
     
-    # Split 90/10
     split_idx = int(total * 0.9)
     
-    # Save to temp location
     data_dir = Path("/tmp/manifold_data")
     data_dir.mkdir(exist_ok=True)
     
@@ -102,195 +65,191 @@ def generate_data(num_legit, num_cheaters, seed, progress=gr.Progress()):
     np.save(data_dir / "val_features.npy", features[split_idx:])
     np.save(data_dir / "val_labels.npy", labels[split_idx:])
     
-    return f"Generated {total} samples:\n- Train: {split_idx}\n- Val: {total - split_idx}\n- Features shape: {features.shape}"
+    return f"✅ Generated {total} samples:\n- Train: {split_idx}\n- Val: {total - split_idx}\n- Shape: {features.shape}"
 
 
-def train_model(batch_size, learning_rate, max_epochs, progress=gr.Progress()):
-    """Train the MANIFOLD model."""
-    global is_training, current_model
-    
-    if is_training:
-        return "Training already in progress!"
+@spaces.GPU(duration=300)
+def train_model(batch_size, learning_rate, num_epochs):
+    global current_model
     
-    is_training = True
     device = "cuda" if torch.cuda.is_available() else "cpu"
+    gpu_info = f"Using: {torch.cuda.get_device_name(0)}" if torch.cuda.is_available() else "CPU only"
     
-    try:
-        # Load data
-        data_dir = Path("/tmp/manifold_data")
-        if not (data_dir / "train_features.npy").exists():
-            is_training = False
-            return "No data found! Generate data first."
-        
-        progress(0.1, desc="Loading data...")
-        train_features = np.load(data_dir / "train_features.npy")
-        train_labels = np.load(data_dir / "train_labels.npy")
-        val_features = np.load(data_dir / "val_features.npy")
-        val_labels = np.load(data_dir / "val_labels.npy")
-        
-        train_dataset = MANIFOLDDataset(data=train_features, labels=train_labels)
-        val_dataset = MANIFOLDDataset(data=val_features, labels=val_labels)
-        
-        train_loader = create_dataloader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=0)
-        val_loader = create_dataloader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=0)
-        
-        progress(0.2, desc="Creating model...")
-        model_config = ModelConfig()
-        model = MANIFOLDLite.from_config(model_config)
+    data_dir = Path("/tmp/manifold_data")
+    if not (data_dir / "train_features.npy").exists():
+        return "❌ No data found! Generate data first.", ""
+    
+    train_features = np.load(data_dir / "train_features.npy")
+    train_labels = np.load(data_dir / "train_labels.npy")
+    val_features = np.load(data_dir / "val_features.npy")
+    val_labels = np.load(data_dir / "val_labels.npy")
+    
+    train_dataset = MANIFOLDDataset(data=train_features, labels=train_labels)
+    val_dataset = MANIFOLDDataset(data=val_features, labels=val_labels)
+    
+    train_loader = create_dataloader(train_dataset, batch_size=int(batch_size), shuffle=True, num_workers=0)
+    val_loader = create_dataloader(val_dataset, batch_size=int(batch_size), shuffle=False, num_workers=0)
+    
+    model = MANIFOLDLite.from_config(ModelConfig())
+    model = model.to(device)
+    
+    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=0.01)
+    scaler = torch.amp.GradScaler(enabled=True)
+    
+    scheduler = CurriculumScheduler()
+    logs = []
+    
+    logs.append(f"🚀 {gpu_info}")
+    logs.append(f"📊 Train: {len(train_dataset)}, Val: {len(val_dataset)}")
+    logs.append(f"🔧 Params: {model.get_num_params():,}")
+    logs.append("-" * 40)
+    
+    global_step = 0
+    
+    for epoch in range(int(num_epochs)):
+        stage_config = scheduler.get_stage_config()
         
-        train_config = TrainingConfig(
-            batch_size=batch_size,
-            learning_rate=learning_rate,
-            max_epochs=max_epochs,
-        )
+        for pg in optimizer.param_groups:
+            pg["lr"] = stage_config["learning_rate"]
         
-        callbacks = [GradioCallback()]
+        model.train()
+        train_loss = 0
+        for batch in train_loader:
+            batch = {k: v.to(device) for k, v in batch.items()}
+            
+            with torch.amp.autocast(device_type='cuda', dtype=torch.float16):
+                outputs = model(batch["features"], mask=batch.get("mask"), active_components=stage_config.get("components"))
+                loss, _ = compute_total_loss(outputs, {"labels": batch["labels"]}, stage_config["losses"], global_step)
+            
+            scaler.scale(loss).backward()
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            scaler.step(optimizer)
+            scaler.update()
+            optimizer.zero_grad(set_to_none=True)
+            
+            train_loss += loss.item()
+            global_step += 1
         
-        trainer = MANIFOLDTrainer(
-            model=model,
-            config=train_config,
-            train_dataloader=train_loader,
-            val_dataloader=val_loader,
-            callbacks=callbacks,
-        )
+        train_loss /= len(train_loader)
         
-        progress(0.3, desc="Training...")
+        model.eval()
+        val_loss = 0
+        correct = 0
+        total = 0
+        with torch.no_grad():
+            for batch in val_loader:
+                batch = {k: v.to(device) for k, v in batch.items()}
+                outputs = model(batch["features"], mask=batch.get("mask"), active_components=stage_config.get("components"))
+                loss, _ = compute_total_loss(outputs, {"labels": batch["labels"]}, stage_config["losses"])
+                val_loss += loss.item()
+                if "predicted_class" in outputs:
+                    correct += (outputs["predicted_class"] == batch["labels"]).sum().item()
+                    total += batch["labels"].size(0)
         
-        # Run training
-        history = trainer.train()
+        val_loss /= len(val_loader)
+        val_acc = correct / total if total > 0 else 0
         
-        # Save model
-        save_path = Path("/tmp/manifold_model.pt")
-        trainer.save_checkpoint(save_path)
-        current_model = model
+        step_info = scheduler.step_epoch()
+        stage_name = step_info["stage_name"].split(":")[0] if ":" in step_info["stage_name"] else step_info["stage_name"]
         
-        is_training = False
-        return f"Training complete!\nFinal val accuracy: {history['val'][-1].get('accuracy', 'N/A')}\nModel saved to {save_path}"
+        logs.append(f"Epoch {epoch+1:2d} | {stage_name:8s} | Loss: {train_loss:.4f} / {val_loss:.4f} | Acc: {val_acc:.4f}")
         
-    except Exception as e:
-        is_training = False
-        return f"Training failed: {str(e)}"
-
-
-def get_training_logs():
-    """Get accumulated training logs."""
-    logs = []
-    while not training_log.empty():
-        try:
-            logs.append(training_log.get_nowait())
-        except:
-            break
-    
-    if not logs:
-        return "No logs yet. Start training to see progress."
-    
-    # Format as table
-    header = "| Epoch | Stage | Train Loss | Val Loss | Val Acc | LR |\n|-------|-------|------------|----------|---------|----|\n"
-    rows = "\n".join([
-        f"| {l['epoch']} | {l['stage'][:20]} | {l['train_loss']} | {l['val_loss']} | {l['val_acc']} | {l['lr']} |"
-        for l in logs
-    ])
-    return header + rows
+        if step_info.get("stage_changed"):
+            logs.append(f"  → Advanced to {scheduler.current_stage.name}")
+    
+    save_path = Path("/tmp/manifold_model.pt")
+    torch.save({"model_state_dict": model.state_dict(), "config": ModelConfig()}, save_path)
+    current_model = model.cpu()
+    
+    logs.append("-" * 40)
+    logs.append(f"✅ Training complete! Final val accuracy: {val_acc:.4f}")
+    
+    return "✅ Training complete!", "\n".join(logs)
 
 
+@spaces.GPU(duration=60)
 def test_inference(num_samples):
-    """Test model inference."""
     global current_model
     
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
     if current_model is None:
-        # Try to load saved model
         model_path = Path("/tmp/manifold_model.pt")
         if model_path.exists():
             current_model = MANIFOLDLite.from_config(ModelConfig())
             ckpt = torch.load(model_path, map_location="cpu")
             current_model.load_state_dict(ckpt["model_state_dict"])
         else:
-            return "No model available! Train a model first."
+            return "❌ No model! Train first."
     
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    current_model.to(device)
-    current_model.eval()
-    
-    # Generate test samples
-    generator = SyntheticDataGenerator(seed=12345)
+    model = current_model.to(device)
+    model.eval()
     
+    generator = SyntheticDataGenerator(seed=99999)
     results = []
-    for i in range(num_samples):
-        is_cheater = i % 2 == 1  # Alternate
+    
+    for i in range(int(num_samples)):
+        is_cheater = i % 2 == 1
         session = generator.generate_player(is_cheater=is_cheater)
         features = torch.tensor(session.to_tensor(), dtype=torch.float32).unsqueeze(0).to(device)
         
         with torch.no_grad():
-            outputs = current_model(features)
+            outputs = model(features)
         
-        pred_class = outputs["predicted_class"].item()
-        uncertainty = outputs["uncertainty"].item()
-        probs = outputs["verdict_probs"][0].cpu().numpy()
+        pred = outputs["predicted_class"].item()
+        conf = outputs["verdict_probs"][0].max().item()
+        unc = outputs["uncertainty"].item()
         
-        class_names = ["Clean", "Suspicious", "Cheating"]
-        results.append({
-            "Sample": i + 1,
-            "Actual": "Cheater" if is_cheater else "Legit",
-            "Predicted": class_names[pred_class],
-            "Confidence": f"{probs.max():.2%}",
-            "Uncertainty": f"{uncertainty:.4f}",
-            "Correct": "✓" if (pred_class > 0) == is_cheater else "✗",
-        })
-    
-    # Format as markdown table
-    header = "| # | Actual | Predicted | Confidence | Uncertainty | Correct |\n|---|--------|-----------|------------|-------------|---------|"
-    rows = "\n".join([
-        f"| {r['Sample']} | {r['Actual']} | {r['Predicted']} | {r['Confidence']} | {r['Uncertainty']} | {r['Correct']} |"
-        for r in results
-    ])
-    
-    correct = sum(1 for r in results if r["Correct"] == "✓")
-    summary = f"\n\n**Accuracy: {correct}/{num_samples} ({100*correct/num_samples:.1f}%)**"
-    
-    return header + "\n" + rows + summary
+        classes = ["Clean", "Suspicious", "Cheating"]
+        actual = "Cheater" if is_cheater else "Legit"
+        correct = "✓" if (pred > 0) == is_cheater else "✗"
+        
+        results.append(f"| {i+1} | {actual} | {classes[pred]} | {conf:.1%} | {unc:.3f} | {correct} |")
+    
+    current_model = model.cpu()
+    
+    header = "| # | Actual | Predicted | Conf | Uncert | ✓/✗ |\n|---|--------|-----------|------|--------|-----|"
+    correct_count = sum(1 for r in results if "✓" in r)
+    footer = f"\n\n**Accuracy: {correct_count}/{num_samples} ({100*correct_count/num_samples:.1f}%)**"
+    
+    return header + "\n" + "\n".join(results) + footer
 
 
-# Build Gradio interface
-with gr.Blocks(title="MANIFOLD Training") as demo:
-    gr.Markdown("# 🎯 MANIFOLD - CS2 Cheat Detection Training")
-    gr.Markdown(f"**Device:** {get_device_info()}")
+with gr.Blocks(title="MANIFOLD Training", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🎯 MANIFOLD - CS2 Cheat Detection")
+    gr.Markdown(f"**{get_device_info()}** | ZeroGPU will allocate H200 on demand")
     
     with gr.Tabs():
-        with gr.TabItem("1. Generate Data"):
-            gr.Markdown("Generate synthetic CS2 player behavior data for training.")
+        with gr.TabItem("1️⃣ Generate Data"):
+            gr.Markdown("Generate synthetic CS2 player data")
             with gr.Row():
-                num_legit = gr.Slider(100, 10000, value=1000, step=100, label="Legit Players")
-                num_cheaters = gr.Slider(100, 5000, value=500, step=100, label="Cheaters")
-                seed = gr.Number(value=42, label="Random Seed")
-            gen_btn = gr.Button("Generate Data", variant="primary")
-            gen_output = gr.Textbox(label="Generation Status", lines=5)
+                num_legit = gr.Slider(100, 10000, value=2000, step=100, label="Legit Players")
+                num_cheaters = gr.Slider(100, 5000, value=1000, step=100, label="Cheaters")
+                seed = gr.Number(value=42, label="Seed")
+            gen_btn = gr.Button("🎲 Generate Data", variant="primary")
+            gen_output = gr.Textbox(label="Status", lines=3)
             gen_btn.click(generate_data, [num_legit, num_cheaters, seed], gen_output)
         
-        with gr.TabItem("2. Train Model"):
-            gr.Markdown("Train the MANIFOLD model with curriculum learning.")
+        with gr.TabItem("2️⃣ Train Model"):
+            gr.Markdown("Train with 4-stage curriculum learning (ZeroGPU: 5 min limit)")
             with gr.Row():
-                batch_size = gr.Slider(8, 128, value=32, step=8, label="Batch Size")
-                learning_rate = gr.Number(value=3e-4, label="Learning Rate")
-                max_epochs = gr.Slider(5, 100, value=20, step=5, label="Max Epochs")
-            train_btn = gr.Button("Start Training", variant="primary")
-            train_output = gr.Textbox(label="Training Status", lines=5)
-            train_btn.click(train_model, [batch_size, learning_rate, max_epochs], train_output)
-            
-            gr.Markdown("### Training Logs")
-            logs_output = gr.Markdown("No logs yet.")
-            refresh_btn = gr.Button("Refresh Logs")
-            refresh_btn.click(get_training_logs, [], logs_output)
-        
-        with gr.TabItem("3. Test Model"):
-            gr.Markdown("Test the trained model on synthetic samples.")
-            num_test = gr.Slider(5, 50, value=10, step=5, label="Number of Test Samples")
-            test_btn = gr.Button("Run Inference", variant="primary")
-            test_output = gr.Markdown("Click 'Run Inference' to test the model.")
+                batch_size = gr.Slider(16, 128, value=64, step=16, label="Batch Size")
+                lr = gr.Number(value=3e-4, label="Learning Rate")
+                epochs = gr.Slider(5, 50, value=15, step=5, label="Epochs")
+            train_btn = gr.Button("🚀 Start Training", variant="primary")
+            train_status = gr.Textbox(label="Status", lines=2)
+            train_logs = gr.Textbox(label="Training Logs", lines=15)
+            train_btn.click(train_model, [batch_size, lr, epochs], [train_status, train_logs])
+        
+        with gr.TabItem("3️⃣ Test Model"):
+            gr.Markdown("Test on synthetic samples")
+            num_test = gr.Slider(5, 30, value=10, step=5, label="Test Samples")
+            test_btn = gr.Button("🔍 Run Inference", variant="primary")
+            test_output = gr.Markdown()
             test_btn.click(test_inference, [num_test], test_output)
     
-    gr.Markdown("---")
-    gr.Markdown("**MANIFOLD** - Motor-Aware Neural Inference for Faithfulness Of Latent Dynamics")
-
+    gr.Markdown("---\n*MANIFOLD: Motor-Aware Neural Inference for Faithfulness Of Latent Dynamics*")
 
 if __name__ == "__main__":
     demo.launch()

requirements.txt

@@ -1,9 +1,9 @@
 torch>=2.1.0
 numpy>=1.24.0
 scipy>=1.11.0
-pandas>=2.0.0
 pydantic>=2.0.0
 tqdm>=4.66.0
 einops>=0.7.0
 scikit-learn>=1.3.0
 gradio>=4.0.0
+spaces