Create demo.py

demo.py

@@ -0,0 +1,135 @@
+#!/usr/bin/env python3
+"""
+Quick demo script to test Mamba Graph implementation
+Device-safe version
+"""
+
+import torch
+import os
+from core.graph_mamba import GraphMamba
+from data.loader import GraphDataLoader
+from utils.metrics import GraphMetrics
+
+def main():
+    print("🧠 Testing Mamba Graph Neural Network")
+    print("=" * 50)
+    
+    # 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
+        },
+        'ordering': {
+            'strategy': 'bfs',
+            'preserve_locality': True
+        }
+    }
+    
+    # Setup device
+    if os.getenv('SPACE_ID'):
+        device = torch.device('cpu')
+    else:
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"Device: {device}")
+    
+    # Load dataset
+    print("\n📊 Loading Cora dataset...")
+    try:
+        data_loader = GraphDataLoader()
+        dataset = data_loader.load_node_classification_data('Cora')
+        data = dataset[0].to(device)
+        
+        # Dataset info
+        info = data_loader.get_dataset_info(dataset)
+        print(f"✅ Success!")
+        print(f"Nodes: {data.num_nodes}")
+        print(f"Edges: {data.num_edges}")
+        print(f"Features: {info['num_features']}")
+        print(f"Classes: {info['num_classes']}")
+        
+    except Exception as e:
+        print(f"❌ Error loading dataset: {e}")
+        return
+    
+    # Initialize model
+    print("\n🏗️ Initializing GraphMamba...")
+    try:
+        model = GraphMamba(config).to(device)
+        total_params = sum(p.numel() for p in model.parameters())
+        print(f"✅ Model initialized!")
+        print(f"Parameters: {total_params:,}")
+        
+    except Exception as e:
+        print(f"❌ Error initializing model: {e}")
+        return
+    
+    # Forward pass test
+    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 different ordering strategies
+    print("\n🔄 Testing ordering strategies...")
+    
+    strategies = ['bfs', 'spectral', 'degree', 'community']
+    
+    for strategy in strategies:
+        try:
+            config['ordering']['strategy'] = strategy
+            model_test = GraphMamba(config).to(device)
+            model_test.eval()
+            
+            with torch.no_grad():
+                h = model_test(data.x, data.edge_index)
+            print(f"✅ {strategy}: Success - Shape {h.shape}")
+            
+        except Exception as e:
+            print(f"❌ {strategy}: Failed - {str(e)}")
+    
+    # Test evaluation
+    print("\n📈 Testing evaluation...")
+    try:
+        # Initialize classifier
+        num_classes = info['num_classes']
+        model._init_classifier(num_classes, device)
+        
+        # Create test mask if not available
+        if hasattr(data, 'test_mask'):
+            mask = data.test_mask
+        else:
+            mask = torch.zeros(data.num_nodes, dtype=torch.bool, device=device)
+            mask[data.num_nodes//2:] = True
+        
+        metrics = GraphMetrics.evaluate_node_classification(model, data, mask, device)
+        print("✅ Evaluation successful!")
+        for metric, value in metrics.items():
+            if isinstance(value, float):
+                print(f"  {metric}: {value:.4f}")
+            
+    except Exception as e:
+        print(f"❌ Evaluation failed: {e}")
+    
+    print("\n✨ Demo completed!")
+    print("🚀 Ready for production deployment!")
+
+if __name__ == "__main__":
+    main()