Update app.py

app.py

@@ -1,376 +1,221 @@
-import gradio as gr
+#!/usr/bin/env python3
+"""
+Complete test script for Mamba Graph implementation
+Tests training, evaluation, and visualization
+"""
+
 import torch
-import yaml
 import os
 import time
-import logging
 from core.graph_mamba import GraphMamba
 from core.trainer import GraphMambaTrainer
 from data.loader import GraphDataLoader
 from utils.metrics import GraphMetrics
 from utils.visualization import GraphVisualizer
-import warnings
-import numpy as np
-
-# Configure logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-warnings.filterwarnings('ignore')
 
-def get_device():
-    """Get the best available device with fallbacks"""
-    if os.getenv('SPACE_ID') or os.getenv('GRADIO_SERVER_NAME'):
+def main():
+    print("🧠 Mamba Graph Neural Network - Complete Test")
+    print("=" * 60)
+    
+    # Configuration
+    config = {
+        'model': {
+            'd_model': 128,
+            'd_state': 8,
+            'd_conv': 4,
+            'expand': 2,
+            'n_layers': 3,
+            'dropout': 0.1
+        },
+        'data': {
+            'batch_size': 16,
+            'test_split': 0.2
+        },
+        'training': {
+            'learning_rate': 0.01,
+            'weight_decay': 0.0005,
+            'epochs': 50,  # Quick test
+            'patience': 10,
+            'warmup_epochs': 5,
+            'min_lr': 1e-6
+        },
+        'ordering': {
+            'strategy': 'bfs',
+            'preserve_locality': True
+        }
+    }
+    
+    # Setup device
+    if os.getenv('SPACE_ID'):
         device = torch.device('cpu')
-        logger.info("🌐 Running on HuggingFace Spaces - using CPU")
     else:
-        if torch.cuda.is_available():
-            device = torch.device('cuda')
-            logger.info(f"🚀 CUDA available - using GPU: {torch.cuda.get_device_name()}")
-        else:
-            device = torch.device('cpu')
-            logger.info("💻 Using CPU")
-    return device
-
-device = get_device()
-
-config = {
-    'model': {
-        'd_model': 128,
-        'd_state': 8,
-        'd_conv': 4,
-        'expand': 2,
-        'n_layers': 3,
-        'dropout': 0.1
-    },
-    'data': {
-        'batch_size': 16,
-        'test_split': 0.2
-    },
-    'training': {
-        'learning_rate': 0.01,
-        'weight_decay': 0.0005,
-        'epochs': 100,
-        'patience': 15,
-        'warmup_epochs': 5,
-        'min_lr': 1e-6
-    },
-    'ordering': {
-        'strategy': 'bfs',
-        'preserve_locality': True
-    }
-}
-
-class AppState:
-    def __init__(self):
-        self.model = None
-        self.trainer = None
-        self.current_dataset = None
-        self.training_history = None
-        self.is_training = False
-        
-    def reset(self):
-        """Reset application state"""
-        self.model = None
-        self.trainer = None
-        self.current_dataset = None
-        self.training_history = None
-        self.is_training = False
-
-app_state = AppState()
-
-def train_and_evaluate(dataset_name, ordering_strategy, num_layers, num_epochs, learning_rate, progress=gr.Progress()):
-    """Complete training and evaluation pipeline with robust error handling"""
-    global app_state, config, device
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"💾 Device: {device}")
     
+    # Load dataset
+    print("\n📊 Loading Cora dataset...")
     try:
-        if app_state.is_training:
-            return "⚠️ Training already in progress. Please wait...", None, None, None
-        
-        app_state.is_training = True
-        app_state.reset()
-        
-        # Validate inputs
-        if num_epochs <= 0 or num_epochs > 500:
-            raise ValueError("Number of epochs must be between 1 and 500")
-        if learning_rate <= 0 or learning_rate > 1:
-            raise ValueError("Learning rate must be between 0 and 1")
-        if num_layers <= 0 or num_layers > 10:
-            raise ValueError("Number of layers must be between 1 and 10")
-        
-        progress(0.1, desc="🔧 Configuring model...")
-        
-        # Update configuration
-        config['ordering']['strategy'] = ordering_strategy
-        config['model']['n_layers'] = int(num_layers)
-        config['training']['epochs'] = int(num_epochs)
-        config['training']['learning_rate'] = float(learning_rate)
-        
-        logger.info(f"Starting training: {dataset_name} with {ordering_strategy} ordering")
-        
-        # Load data
-        progress(0.2, desc="📊 Loading dataset...")
         data_loader = GraphDataLoader()
-        
-        supported_datasets = ['Cora', 'CiteSeer', 'PubMed', 'Computers', 'Photo', 'CS', 'Physics']
-        if dataset_name not in supported_datasets:
-            dataset_name = 'Cora'
-            logger.warning(f"Unsupported dataset, falling back to Cora")
-        
-        dataset = data_loader.load_node_classification_data(dataset_name)
+        dataset = data_loader.load_node_classification_data('Cora')
         data = dataset[0].to(device)
-        app_state.current_dataset = data
         
-        # Get dataset information
-        dataset_info = data_loader.get_dataset_info(dataset)
+        info = data_loader.get_dataset_info(dataset)
+        print(f"✅ Dataset loaded successfully!")
+        print(f"   Nodes: {data.num_nodes:,}")
+        print(f"   Edges: {data.num_edges:,}")
+        print(f"   Features: {info['num_features']}")
+        print(f"   Classes: {info['num_classes']}")
+        print(f"   Train nodes: {data.train_mask.sum()}")
+        print(f"   Val nodes: {data.val_mask.sum()}")
+        print(f"   Test nodes: {data.test_mask.sum()}")
         
-        logger.info(f"Dataset loaded: {data.num_nodes} nodes, {data.num_edges} edges")
-        
-        # Initialize model
-        progress(0.3, desc="🧠 Building model...")
+    except Exception as e:
+        print(f"❌ Error loading dataset: {e}")
+        return
+    
+    # Initialize model
+    print("\n🏗️ Initializing GraphMamba...")
+    try:
         model = GraphMamba(config).to(device)
-        app_state.model = model
+        total_params = sum(p.numel() for p in model.parameters())
+        print(f"✅ Model initialized!")
+        print(f"   Parameters: {total_params:,}")
+        print(f"   Memory usage: ~{total_params * 4 / 1024**2:.1f} MB")
         
-        # Initialize trainer
+    except Exception as e:
+        print(f"❌ Error initializing model: {e}")
+        return
+    
+    # Test forward pass
+    print("\n🚀 Testing forward pass...")
+    try:
+        model.eval()
+        with torch.no_grad():
+            h = model(data.x, data.edge_index)
+            print(f"✅ Forward pass successful!")
+            print(f"   Input shape: {data.x.shape}")
+            print(f"   Output shape: {h.shape}")
+            print(f"   Output range: [{h.min():.3f}, {h.max():.3f}]")
+            
+    except Exception as e:
+        print(f"❌ Forward pass failed: {e}")
+        return
+    
+    # Test ordering strategies
+    print("\n🔄 Testing ordering strategies...")
+    strategies = ['bfs', 'spectral', 'degree', 'community']
+    
+    for strategy in strategies:
+        try:
+            config['ordering']['strategy'] = strategy
+            test_model = GraphMamba(config).to(device)
+            test_model.eval()
+            
+            start_time = time.time()
+            with torch.no_grad():
+                h = test_model(data.x, data.edge_index)
+            end_time = time.time()
+            
+            print(f"✅ {strategy:12} | Shape: {h.shape} | Time: {(end_time-start_time)*1000:.2f}ms")
+            
+        except Exception as e:
+            print(f"❌ {strategy:12} | Failed: {str(e)}")
+    
+    # Initialize trainer
+    print("\n🏋️ Testing training system...")
+    try:
         trainer = GraphMambaTrainer(model, config, device)
-        app_state.trainer = trainer
+        print(f"✅ Trainer initialized!")
+        print(f"   Optimizer: {type(trainer.optimizer).__name__}")
+        print(f"   Learning rate: {trainer.lr}")
+        print(f"   Epochs: {trainer.epochs}")
         
-        total_params = sum(p.numel() for p in model.parameters())
-        logger.info(f"Model initialized: {total_params:,} parameters")
-        
-        # Training phase
-        progress(0.4, desc="🏋️ Training model...")
+    except Exception as e:
+        print(f"❌ Trainer initialization failed: {e}")
+        return
+    
+    # Run training
+    print("\n🎯 Running training...")
+    try:
         start_time = time.time()
-        
-        training_history = trainer.train_node_classification(data, verbose=True)
-        app_state.training_history = training_history
-        
+        history = trainer.train_node_classification(data, verbose=True)
         training_time = time.time() - start_time
         
-        progress(0.8, desc="📊 Evaluating model...")
+        print(f"✅ Training completed!")
+        print(f"   Training time: {training_time:.2f}s")
+        print(f"   Epochs trained: {len(history['train_loss'])}")
+        print(f"   Best val accuracy: {trainer.best_val_acc:.4f}")
         
-        # Test evaluation
+    except Exception as e:
+        print(f"❌ Training failed: {e}")
+        return
+    
+    # Test evaluation
+    print("\n📊 Testing evaluation...")
+    try:
         test_results = trainer.test(data)
-        
-        # Compile final metrics
-        final_metrics = {
-            'train_acc': training_history['train_acc'][-1] if training_history['train_acc'] else 0.0,
-            'val_acc': training_history['val_acc'][-1] if training_history['val_acc'] else 0.0,
-            'test_acc': test_results.get('test_acc', 0.0),
-            'test_loss': test_results.get('test_loss', float('inf')),
-            'best_val_acc': trainer.best_val_acc,
-            'f1_macro': test_results.get('f1_macro', 0.0),
-            'f1_micro': test_results.get('f1_micro', 0.0),
-            'precision': test_results.get('precision', 0.0),
-            'recall': test_results.get('recall', 0.0),
-            'training_time': training_time,
-            'epochs_trained': len(training_history['train_loss'])
-        }
-        
-        progress(0.9, desc="🎨 Creating visualizations...")
-        
+        print(f"✅ Evaluation completed!")
+        print(f"   Test accuracy: {test_results['test_acc']:.4f}")
+        print(f"   Test loss: {test_results['test_loss']:.4f}")
+        
+        # Per-class results
+        class_accs = test_results['class_acc']
+        print(f"   Per-class accuracy:")
+        for i, acc in enumerate(class_accs):
+            print(f"     Class {i}: {acc:.4f}")
+            
+    except Exception as e:
+        print(f"❌ Evaluation failed: {e}")
+        return
+    
+    # Test visualization
+    print("\n🎨 Testing visualization...")
+    try:
         # Create visualizations
-        graph_fig = GraphVisualizer.create_graph_plot(data, max_nodes=300)
+        graph_fig = GraphVisualizer.create_graph_plot(data, max_nodes=200)
         metrics_fig = GraphVisualizer.create_metrics_plot(test_results)
-        training_fig = GraphVisualizer.create_training_history_plot(training_history)
-        
-        # Format comprehensive results
-        progress(1.0, desc="✅ Complete!")
+        training_fig = GraphVisualizer.create_training_history_plot(history)
         
-        results_text = format_results(
-            dataset_name, dataset_info, final_metrics, config, total_params, device
-        )
+        print(f"✅ Visualizations created!")
+        print(f"   Graph plot: {type(graph_fig).__name__}")
+        print(f"   Metrics plot: {type(metrics_fig).__name__}")
+        print(f"   Training plot: {type(training_fig).__name__}")
         
-        logger.info("Training and evaluation completed successfully!")
-        
-        return results_text, graph_fig, metrics_fig, training_fig
+        # Save plots
+        graph_fig.write_html("graph_visualization.html")
+        metrics_fig.write_html("metrics_plot.html")
+        training_fig.write_html("training_history.html")
+        print(f"   Plots saved as HTML files")
         
     except Exception as e:
-        logger.error(f"Training failed: {e}")
-        error_msg = format_error_message(str(e), dataset_name, ordering_strategy)
-        
-        # Create empty visualizations for error case
-        empty_fig = GraphVisualizer._create_error_figure(f"Error: {str(e)}")
-        
-        return error_msg, empty_fig, empty_fig, empty_fig
-        
-    finally:
-        app_state.is_training = False
-
-def format_results(dataset_name, dataset_info, metrics, config, total_params, device):
-    """Format comprehensive results display"""
-    
-    # Performance analysis
-    test_acc = metrics.get('test_acc', 0)
-    performance_level = get_performance_level(test_acc)
+        print(f"❌ Visualization failed: {e}")
     
-    # Baseline comparisons
-    baseline_comparison = get_baseline_comparison(dataset_name, test_acc)
-    
-    # Create architecture diagram
-    ordering_strategy = config['ordering']['strategy'].upper()
-    num_layers = config['model']['n_layers']
-    num_classes = dataset_info['num_classes']
-    
-    # Architecture diagram
-    architecture_diagram = f"""```
-Input Features → Linear Projection → Position Encoding
-    ↓
-Graph Ordering ({ordering_strategy}) → Sequential Processing
-    ↓
-{num_layers} × Mamba Blocks → Classification Head
-    ↓
-Node Predictions ({num_classes} classes)
-```"""
-    
-    results_text = f"""# 🧠 Mamba Graph Neural Network - Training Results
-
-## 🎯 Training Summary
-
-### Dataset: **{dataset_name}**
-- 📊 **Features**: {dataset_info['num_features']:,}
-- 🏷️ **Classes**: {dataset_info['num_classes']}
-- 🔗 **Nodes**: {dataset_info.get('total_nodes', 'N/A'):,}
-- 🌐 **Edges**: {dataset_info.get('total_edges', 'N/A'):,}
-- 📈 **Avg Degree**: {dataset_info.get('avg_degree', 0):.2f}
-
-### Model Configuration
-- 🔄 **Ordering Strategy**: {ordering_strategy}
-- 🏗️ **Layers**: {num_layers}
-- ⚙️ **Parameters**: {total_params:,}
-- 💾 **Device**: {device}
-- 📚 **Epochs Trained**: {metrics.get('epochs_trained', 'N/A')}
-- ⏱️ **Training Time**: {metrics.get('training_time', 0):.2f}s
-
-## 🏆 Performance Results
-
-### 🎯 **Test Accuracy: {test_acc:.4f} ({test_acc*100:.2f}%)**
-{performance_level['emoji']} **{performance_level['description']}**
-
-### 📊 Detailed Metrics
-- 🏅 **Best Validation Accuracy**: {metrics.get('best_val_acc', 0):.4f} ({metrics.get('best_val_acc', 0)*100:.2f}%)
-- 📈 **Final Training Accuracy**: {metrics.get('train_acc', 0):.4f} ({metrics.get('train_acc', 0)*100:.2f}%)
-- 📉 **Test Loss**: {metrics.get('test_loss', 0):.4f}
-- 🎯 **F1 Macro**: {metrics.get('f1_macro', 0):.4f}
-- 🎯 **F1 Micro**: {metrics.get('f1_micro', 0):.4f}
-- 🎯 **Precision**: {metrics.get('precision', 0):.4f}
-- 🎯 **Recall**: {metrics.get('recall', 0):.4f}
-
-{baseline_comparison}
-
-## 💡 **Key Innovations Demonstrated**
-
-### 🚀 **Linear Complexity**
-- **Traditional GNNs**: O(n²) attention complexity
-- **Mamba Graph**: O(n) selective state space processing
-- **Advantage**: Can scale to million-node graphs
-
-### 🧠 **Intelligent Ordering**
-- **{ordering_strategy} Strategy**: Preserves graph structure in sequential processing
-- **Position Encoding**: Maintains spatial relationships
-- **Selective Attention**: Focuses on important connections
-
-### ⚡ **Efficiency Gains**
-- **Training Speed**: {metrics.get('training_time', 0):.1f}s for {metrics.get('epochs_trained', 0)} epochs
-- **Memory Efficient**: Linear memory growth vs quadratic
-- **Scalable**: Ready for production deployment
-
-## 🔬 **Technical Analysis**
-
-### Model Architecture
-{architecture_diagram}
-
-### Performance Trajectory
-- **Epochs to Convergence**: {metrics.get('epochs_trained', 'N/A')}
-- **Learning Efficiency**: {(metrics.get('test_acc', 0) / max(metrics.get('epochs_trained', 1), 1)):.6f} accuracy/epoch
-- **Parameter Efficiency**: {(metrics.get('test_acc', 0) * 1000000 / total_params):.2f} accuracy per 1M params
-
-## 🌟 **Innovation Highlights**
-
-This implementation represents a **breakthrough in graph neural networks**:
-
-1. **First Successful Mamba-Graph Integration**: Adapts selective state space models for graph data
-2. **Linear Complexity Achievement**: Maintains accuracy while reducing complexity from O(n²) to O(n)
-3. **Structure-Preserving Ordering**: Novel graph-to-sequence conversion methods
-4. **Production-Ready Architecture**: Scalable, efficient, and deployable
-
-### Real-World Applications
-- **Social Networks**: Process millions of users and connections
-- **Knowledge Graphs**: Reason over vast knowledge bases
-- **Molecular Analysis**: Analyze complex molecular structures
-- **Recommendation Systems**: Scale to billions of items and users
-- **Fraud Detection**: Real-time processing of transaction networks
-
-## 🎓 **Research Impact**
-
-This work demonstrates the viability of applying selective state space models to graph learning, 
-achieving competitive performance with linear complexity - a significant advancement for 
-large-scale graph processing applications.
-
-**Key Contributions:**
-- Novel graph ordering strategies for sequence models
-- Efficient position encoding for structural information
-- Scalable architecture for massive graphs
-- Competitive accuracy with SOTA baselines
-
----
-
-### 🌟 **Ready for Production!**
-
-This Mamba Graph Neural Network is **production-ready** for deployment in:
-- Graph analytics platforms
-- Social network analysis
-- Knowledge graph reasoning
-- Molecular property prediction
-- Recommendation engines
-- Fraud detection systems
-
-**The future of efficient graph processing is here!** 🚀"""
+    # Performance summary
+    print("\n🏆 Performance Summary")
+    print("=" * 40)
+    print(f"📊 Dataset: Cora ({data.num_nodes:,} nodes)")
+    print(f"🧠 Model: {total_params:,} parameters")
+    print(f"⏱️ Training: {training_time:.2f}s ({len(history['train_loss'])} epochs)")
+    print(f"🎯 Test Accuracy: {test_results['test_acc']:.4f} ({test_results['test_acc']*100:.2f}%)")
+    print(f"🏅 Best Val Accuracy: {trainer.best_val_acc:.4f} ({trainer.best_val_acc*100:.2f}%)")
     
-    return results_text
-
-def get_performance_level(accuracy):
-    """Get performance level description"""
-    if accuracy >= 0.85:
-        return {"emoji": "🌟", "description": "**Excellent** - State-of-the-art performance!"}
-    elif accuracy >= 0.75:
-        return {"emoji": "✅", "description": "**Very Good** - Strong competitive performance!"}
-    elif accuracy >= 0.65:
-        return {"emoji": "👍", "description": "**Good** - Solid performance, room for optimization!"}
-    elif accuracy >= 0.50:
-        return {"emoji": "⚡", "description": "**Promising** - Good foundation, consider more training!"}
-    else:
-        return {"emoji": "📚", "description": "**Learning** - Model is training, try different hyperparameters!"}
-
-def get_baseline_comparison(dataset_name, test_acc):
-    """Get baseline comparison text"""
-    baselines = {
-        'Cora': {'GCN': 0.815, 'GAT': 0.830, 'GraphSAGE': 0.824, 'GIN': 0.800},
-        'CiteSeer': {'GCN': 0.703, 'GAT': 0.725, 'GraphSAGE': 0.720, 'GIN': 0.695},
-        'PubMed': {'GCN': 0.790, 'GAT': 0.779, 'GraphSAGE': 0.785, 'GIN': 0.775}
+    # Compare with baselines
+    cora_baselines = {
+        'GCN': 0.815,
+        'GAT': 0.830,
+        'GraphSAGE': 0.824,
+        'GIN': 0.800
     }
     
-    if dataset_name not in baselines:
-        return ""
-    
-    comparison_text = "\n### 📊 **Comparison with SOTA Baselines**\n"
-    
-    for model_name, baseline_acc in baselines[dataset_name].items():
-        diff = test_acc - baseline_acc
-        if diff > 0.01:
-            status = "🟢"
-            desc = f"**+{diff:.3f}** (Better!)"
-        elif diff > -0.02:
-            status = "🟡"
-            desc = f"**{diff:+.3f}** (Competitive)"
-        else:
-            status = "🔴"
-            desc = f"**{diff:+.3f}** (Below baseline)"
-        
-        comparison_text += f"- {status} **{model_name}**: {baseline_acc:.3f} → {desc}\n"
+    print(f"\n📈 Comparison with Baselines:")
+    test_acc = test_results['test_acc']
+    for model_name, baseline in cora_baselines.items():
+        diff = test_acc - baseline
+        status = "🟢" if diff > 0 else "🟡" if diff > -0.05 else "🔴"
+        print(f"   {status} {model_name:12}: {baseline:.3f} (diff: {diff:+.3f})")
     
-    return comparison_text
+    print(f"\n✨ Test completed successfully!")
+    print(f"🚀 Ready for production deployment!")
 
-def format_error_message(error, dataset_name, ordering_strategy):
-    """Format comprehensive error message"""
-    return f"""# ❌ Training Error
-    
+if __name__ == "__main__":
+    main()