Update app.py

app.py

@@ -1,253 +1,310 @@
 import gradio as gr
 import torch
 import yaml
-import plotly.graph_objects as go
-import plotly.express as px
+import os
 from core.graph_mamba import GraphMamba
 from data.loader import GraphDataLoader
 from utils.metrics import GraphMetrics
-import networkx as nx
-import numpy as np
+from utils.visualization import GraphVisualizer
+import warnings
+warnings.filterwarnings('ignore')
+
+# Force CPU for HuggingFace Spaces
+if os.getenv('SPACE_ID') or os.getenv('GRADIO_SERVER_NAME'):
+    device = torch.device('cpu')
+    print("Running on HuggingFace Spaces - using CPU")
+else:
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"Running locally - using {device}")
 
 # Load configuration
-with open('config.yaml', 'r') as f:
-    config = yaml.safe_load(f)
+config = {
+    'model': {
+        'd_model': 128,  # Smaller for demo
+        'd_state': 8,
+        'd_conv': 4,
+        'expand': 2,
+        'n_layers': 3,  # Fewer layers for speed
+        'dropout': 0.1
+    },
+    'data': {
+        'batch_size': 16,
+        'test_split': 0.2
+    },
+    'ordering': {
+        'strategy': 'bfs',
+        'preserve_locality': True
+    }
+}
 
-# Initialize model (will be loaded dynamically based on dataset)
+# Global model holder
 model = None
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+current_dataset = None
 
 def load_and_evaluate(dataset_name, ordering_strategy, num_layers):
-    """Load dataset, train/evaluate model, return results"""
-    global model, config
+    """Load dataset, configure model, return results"""
+    global model, config, current_dataset
     
     try:
         # Update config
         config['ordering']['strategy'] = ordering_strategy
         config['model']['n_layers'] = num_layers
         
+        print(f"Loading dataset: {dataset_name}")
+        
         # Load data
         data_loader = GraphDataLoader()
         
-        if dataset_name in ['Cora', 'CiteSeer', 'PubMed', 'Reddit', 'Flickr']:
+        if dataset_name in ['Cora', 'CiteSeer', 'PubMed']:
             dataset = data_loader.load_node_classification_data(dataset_name)
             data = dataset[0].to(device)
             task_type = 'node_classification'
+            current_dataset = data
+            print(f"Loaded {dataset_name}: {data.num_nodes} nodes, {data.num_edges} edges")
         else:
             dataset = data_loader.load_graph_classification_data(dataset_name)
-            train_loader, val_loader, test_loader = data_loader.create_dataloaders(
-                dataset, 'graph_classification'
-            )
             task_type = 'graph_classification'
+            print(f"Loaded {dataset_name}: {len(dataset)} graphs")
         
         # Get dataset info
         dataset_info = data_loader.get_dataset_info(dataset)
         
+        print(f"Dataset info: {dataset_info}")
+        
         # Initialize model
+        print("Initializing GraphMamba model...")
         model = GraphMamba(config).to(device)
         
-        # Quick evaluation (in production, you'd load pre-trained weights)
+        # Initialize classifier for evaluation
+        num_classes = dataset_info['num_classes']
+        model._init_classifier(num_classes, device)
+        
+        total_params = sum(p.numel() for p in model.parameters())
+        print(f"Model parameters: {total_params:,}")
+        
+        # Quick evaluation (random weights for demo)
+        print("Running evaluation...")
         if task_type == 'node_classification':
             # Use test mask for evaluation
+            if hasattr(data, 'test_mask'):
+                mask = data.test_mask
+            else:
+                # Create a test mask if not available
+                num_nodes = data.num_nodes
+                mask = torch.zeros(num_nodes, dtype=torch.bool)
+                mask[num_nodes//2:] = True
+            
             metrics = GraphMetrics.evaluate_node_classification(
-                model, data, data.test_mask, device
+                model, data, mask, device
             )
             
             # Create visualization
-            fig = create_graph_visualization(data)
+            print("Creating visualization...")
+            fig = GraphVisualizer.create_graph_plot(data)
             
         else:
             # Graph classification
+            train_loader, val_loader, test_loader = data_loader.create_dataloaders(
+                dataset, 'graph_classification'
+            )
             metrics = GraphMetrics.evaluate_graph_classification(
                 model, test_loader, device
             )
-            fig = create_dataset_stats_plot(dataset_info)
+            fig = GraphVisualizer.create_metrics_plot(metrics)
         
         # Format results
         results_text = f"""
-        ## Dataset: {dataset_name}
-        
-        **Dataset Statistics:**
-        - Features: {dataset_info['num_features']}
-        - Classes: {dataset_info['num_classes']}
-        - Graphs: {dataset_info['num_graphs']}
-        - Avg Nodes: {dataset_info['avg_nodes']:.1f}
-        - Avg Edges: {dataset_info['avg_edges']:.1f}
-        
-        **Model Configuration:**
-        - Ordering Strategy: {ordering_strategy}
-        - Layers: {num_layers}
-        - Model Parameters: {sum(p.numel() for p in model.parameters()):,}
-        
-        **Performance Metrics:**
+## 🧠 Mamba Graph Neural Network Results
+
+### Dataset: {dataset_name}
+
+**Dataset Statistics:**
+- 📊 Features: {dataset_info['num_features']}
+- 🏷️ Classes: {dataset_info['num_classes']}
+- 📈 Graphs: {dataset_info['num_graphs']}
+- 🔗 Avg Nodes: {dataset_info['avg_nodes']:.1f}
+- 🌐 Avg Edges: {dataset_info['avg_edges']:.1f}
+
+**Model Configuration:**
+- 🔄 Ordering Strategy: {ordering_strategy}
+- 🏗️ Layers: {num_layers}
+- ⚙️ Parameters: {sum(p.numel() for p in model.parameters()):,}
+- 💾 Device: {device}
+
+**Performance Metrics:**
         """
         
         for metric, value in metrics.items():
-            if isinstance(value, float):
-                results_text += f"- {metric.replace('_', ' ').title()}: {value:.4f}\n"
+            if isinstance(value, float) and metric != 'error':
+                results_text += f"- 📈 {metric.replace('_', ' ').title()}: {value:.4f}\n"
+            elif metric == 'error':
+                results_text += f"- ⚠️ Error: {value}\n"
+        
+        results_text += f"""
+
+**Status:** ✅ Model successfully loaded and evaluated!
+
+*Note: This is a demo with random weights. In production, the model would be trained on the dataset.*
+        """
         
+        print("Evaluation completed successfully!")
         return results_text, fig
         
     except Exception as e:
-        return f"Error: {str(e)}", None
+        error_msg = f"""
+## ❌ Error Loading Model
 
-def create_graph_visualization(data):
-    """Create interactive graph visualization"""
-    try:
-        # Convert to NetworkX
-        G = nx.Graph()
-        edge_list = data.edge_index.t().cpu().numpy()
-        G.add_edges_from(edge_list)
-        
-        # Limit to first 1000 nodes for visualization
-        if len(G.nodes()) > 1000:
-            nodes_to_keep = list(G.nodes())[:1000]
-            G = G.subgraph(nodes_to_keep)
-        
-        # Layout
-        pos = nx.spring_layout(G, k=1, iterations=50)
-        
-        # Node colors based on labels if available
-        node_colors = []
-        if hasattr(data, 'y') and data.y is not None:
-            labels = data.y.cpu().numpy()
-            for node in G.nodes():
-                if node < len(labels):
-                    node_colors.append(labels[node])
-                else:
-                    node_colors.append(0)
-        else:
-            node_colors = [0] * len(G.nodes())
-        
-        # Create traces
-        edge_x, edge_y = [], []
-        for edge in G.edges():
-            x0, y0 = pos[edge[0]]
-            x1, y1 = pos[edge[1]]
-            edge_x.extend([x0, x1, None])
-            edge_y.extend([y0, y1, None])
-        
-        node_x = [pos[node][0] for node in G.nodes()]
-        node_y = [pos[node][1] for node in G.nodes()]
-        
-        fig = go.Figure()
-        
-        # Add edges
-        fig.add_trace(go.Scatter(
-            x=edge_x, y=edge_y,
-            line=dict(width=0.5, color='#888'),
-            hoverinfo='none',
-            mode='lines'
-        ))
-        
-        # Add nodes
-        fig.add_trace(go.Scatter(
-            x=node_x, y=node_y,
-            mode='markers',
-            hoverinfo='text',
-            text=[f'Node {i}' for i in G.nodes()],
-            marker=dict(
-                size=8,
-                color=node_colors,
-                colorscale='Viridis',
-                line=dict(width=2)
-            )
-        ))
-        
-        fig.update_layout(
-            title='Graph Visualization',
-            showlegend=False,
-            hovermode='closest',
-            margin=dict(b=20,l=5,r=5,t=40),
-            annotations=[
-                dict(
-                    text="Graph structure visualization",
-                    showarrow=False,
-                    xref="paper", yref="paper",
-                    x=0.005, y=-0.002,
-                    xanchor='left', yanchor='bottom',
-                    font=dict(color="black", size=12)
-                )
-            ],
-            xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
-            yaxis=dict(showgrid=False, zeroline=False, showticklabels=False)
-        )
+**Error:** {str(e)}
+
+**Troubleshooting:**
+- Check dataset availability
+- Verify device compatibility  
+- Try different ordering strategy
+
+**Debug Info:**
+- Device: {device}
+- Dataset: {dataset_name}
+- Strategy: {ordering_strategy}
+        """
         
-        return fig
+        print(f"Error: {e}")
         
-    except Exception as e:
         # Return empty plot on error
+        import plotly.graph_objects as go
         fig = go.Figure()
-        fig.add_annotation(text=f"Visualization error: {str(e)}", x=0.5, y=0.5)
-        return fig
-
-def create_dataset_stats_plot(dataset_info):
-    """Create dataset statistics visualization"""
-    stats = [
-        ['Features', dataset_info['num_features']],
-        ['Classes', dataset_info['num_classes']],
-        ['Avg Nodes', dataset_info['avg_nodes']],
-        ['Avg Edges', dataset_info['avg_edges']]
-    ]
-    
-    fig = go.Figure(data=[
-        go.Bar(
-            x=[stat[0] for stat in stats],
-            y=[stat[1] for stat in stats],
-            marker_color='lightblue'
+        fig.add_annotation(
+            text=f"Error: {str(e)}",
+            x=0.5, y=0.5,
+            xref="paper", yref="paper",
+            showarrow=False
         )
-    ])
-    
-    fig.update_layout(
-        title='Dataset Statistics',
-        xaxis_title='Metric',
-        yaxis_title='Value'
-    )
-    
-    return fig
+        
+        return error_msg, fig
 
 # Gradio interface
-with gr.Blocks(title="Mamba Graph Neural Network") as demo:
+with gr.Blocks(
+    title="🧠 Mamba Graph Neural Network",
+    theme=gr.themes.Soft(),
+    css="""
+    .gradio-container {
+        max-width: 1200px !important;
+    }
+    """
+) as demo:
+    
     gr.Markdown("""
     # 🧠 Mamba Graph Neural Network
     
-    Real-time evaluation of Graph-Mamba on standard benchmarks.
-    This uses actual datasets and trained models - no synthetic data.
+    **Real-time evaluation of Graph-Mamba on standard benchmarks.**
+    
+    This demonstrates the revolutionary combination of Mamba's linear complexity with graph neural networks.
+    Uses actual datasets and real model architectures - no synthetic data.
+    
+    🚀 **Features:**
+    - Linear O(n) complexity for massive graphs
+    - Multiple graph ordering strategies  
+    - Real benchmark datasets (Cora, CiteSeer, etc.)
+    - Interactive visualizations
     """)
     
     with gr.Row():
-        with gr.Column():
+        with gr.Column(scale=1):
+            gr.Markdown("### 🎮 Model Configuration")
+            
             dataset_choice = gr.Dropdown(
-                choices=['Cora', 'CiteSeer', 'PubMed', 'MUTAG', 'ENZYMES', 'PROTEINS'],
+                choices=['Cora', 'CiteSeer', 'PubMed', 'MUTAG', 'ENZYMES'],
                 value='Cora',
-                label="Dataset"
+                label="📊 Dataset",
+                info="Choose a graph dataset for evaluation"
             )
             
             ordering_choice = gr.Dropdown(
                 choices=['bfs', 'spectral', 'degree', 'community'],
                 value='bfs',
-                label="Graph Ordering Strategy"
+                label="🔄 Graph Ordering Strategy",
+                info="How to convert graph to sequence"
             )
             
             layers_slider = gr.Slider(
-                minimum=2, maximum=8, value=4, step=1,
-                label="Number of Mamba Layers"
+                minimum=2, maximum=6, value=3, step=1,
+                label="🏗️ Number of Mamba Layers",
+                info="More layers = more capacity"
             )
             
-            evaluate_btn = gr.Button("Evaluate Model", variant="primary")
+            evaluate_btn = gr.Button(
+                "🚀 Evaluate Model", 
+                variant="primary",
+                size="lg"
+            )
+            
+            gr.Markdown("""
+            ### 📖 Ordering Strategies:
+            - **BFS**: Breadth-first traversal
+            - **Spectral**: Eigenvalue-based ordering  
+            - **Degree**: High-degree nodes first
+            - **Community**: Cluster-aware ordering
+            """)
         
-        with gr.Column():
-            results_text = gr.Markdown("Select parameters and click 'Evaluate Model'")
+        with gr.Column(scale=2):
+            results_text = gr.Markdown("""
+            ### 👋 Welcome!
             
+            Select your parameters and click **'🚀 Evaluate Model'** to see Mamba Graph in action.
+            
+            The model will:
+            1. 📥 Load the selected dataset
+            2. 🔄 Apply graph ordering strategy  
+            3. 🧠 Process through Mamba layers
+            4. 📊 Evaluate performance
+            5. 📈 Show results and visualization
+            """)
+    
     with gr.Row():
-        visualization = gr.Plot(label="Graph Visualization")
+        with gr.Column():
+            visualization = gr.Plot(
+                label="📈 Graph Visualization",
+                container=True
+            )
     
     # Event handlers
     evaluate_btn.click(
         fn=load_and_evaluate,
         inputs=[dataset_choice, ordering_choice, layers_slider],
-        outputs=[results_text, visualization]
+        outputs=[results_text, visualization],
+        show_progress=True
     )
+    
+    # Example section
+    gr.Markdown("""
+    ---
+    ### 🎯 What Makes This Special?
+    
+    **Traditional GNNs:** O(n²) complexity limits them to small graphs
+    
+    **Mamba Graph:** O(n) complexity enables processing of massive graphs
+    
+    **Key Innovation:** Smart graph-to-sequence conversion preserves structural information while enabling linear-time processing.
+    
+    ### 🔬 Technical Details:
+    - **Selective State Space Models** for sequence processing
+    - **Structure-preserving ordering** algorithms
+    - **Position encoding** to maintain graph relationships
+    - **Multi-scale processing** for different graph properties
+    
+    ### 📚 References:
+    - Mamba: Linear-Time Sequence Modeling (Gu & Dao, 2023)
+    - Graph Neural Networks (Kipf & Welling, 2017)
+    - Spectral Graph Theory applications
+    """)
 
 if __name__ == "__main__":
-    demo.launch(share=True)
+    print("🧠 Starting Mamba Graph Demo...")
+    print(f"Device: {device}")
+    print("Loading Gradio interface...")
+    
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_error=True,
+        share=False  # Set to False for HuggingFace Spaces
+    )