Update utils/metrics.py

utils/metrics.py

@@ -1,10 +1,10 @@
 import torch
 import torch.nn.functional as F
-from sklearn.metrics import accuracy_score, f1_score, roc_auc_score
+from sklearn.metrics import accuracy_score, f1_score, roc_auc_score, classification_report
 import numpy as np
 
 class GraphMetrics:
-    """Production-ready evaluation metrics - device safe"""
+    """Comprehensive evaluation metrics for graph learning"""
     
     @staticmethod
     def accuracy(pred, target):
@@ -44,7 +44,40 @@ class GraphMetrics:
             return 0.0
     
     @staticmethod
-    def evaluate_node_classification(model, data, mask, device):
+    def roc_auc(pred, target):
+        """ROC AUC score"""
+        try:
+            if pred.dim() > 1:
+                # Multi-class
+                pred_probs = F.softmax(pred, dim=1).cpu().numpy()
+                target_onehot = F.one_hot(target, num_classes=pred.size(1)).cpu().numpy()
+                return roc_auc_score(target_onehot, pred_probs, multi_class='ovr', average='macro')
+            else:
+                # Binary
+                pred_probs = torch.sigmoid(pred).cpu().numpy()
+                target_labels = target.cpu().numpy()
+                return roc_auc_score(target_labels, pred_probs)
+        except:
+            return 0.0
+    
+    @staticmethod
+    def classification_report_dict(pred, target):
+        """Detailed classification report"""
+        try:
+            if pred.dim() > 1:
+                pred_labels = pred.argmax(dim=1)
+            else:
+                pred_labels = pred
+            pred_labels = pred_labels.cpu().numpy()
+            target_labels = target.cpu().numpy()
+            
+            report = classification_report(target_labels, pred_labels, output_dict=True, zero_division=0)
+            return report
+        except:
+            return {}
+    
+    @staticmethod
+    def evaluate_node_classification(model, data, mask, device, detailed=False):
         """Comprehensive node classification evaluation"""
         model.eval()
         
@@ -74,59 +107,16 @@ class GraphMetrics:
                     'f1_micro': GraphMetrics.f1_score_micro(pred_masked, target_masked),
                 }
                 
-        except Exception as e:
-            print(f"Evaluation error: {e}")
-            metrics = {
-                'accuracy': 0.0,
-                'f1_macro': 0.0,
-                'f1_micro': 0.0,
-                'error': str(e)
-            }
-                
-        return metrics
-    
-    @staticmethod
-    def evaluate_graph_classification(model, dataloader, device):
-        """Comprehensive graph classification evaluation"""
-        model.eval()
-        
-        all_preds = []
-        all_targets = []
-        
-        try:
-            with torch.no_grad():
-                for batch in dataloader:
-                    batch = batch.to(device)
-                    h = model(batch.x, batch.edge_index, batch.batch)
-                    
-                    # Graph-level prediction
-                    graph_h = model.get_graph_embedding(h, batch.batch)
-                    
-                    if hasattr(model, 'classifier') and model.classifier is not None:
-                        pred = model.classifier(graph_h)
-                    else:
-                        # Initialize classifier
-                        num_classes = len(torch.unique(batch.y))
-                        model._init_classifier(num_classes, device)
-                        pred = model.classifier(graph_h)
-                    
-                    all_preds.append(pred.cpu())
-                    all_targets.append(batch.y.cpu())
-            
-            if all_preds:
-                all_preds = torch.cat(all_preds, dim=0)
-                all_targets = torch.cat(all_targets, dim=0)
+                # Add detailed metrics if requested
+                if detailed:
+                    metrics['roc_auc'] = GraphMetrics.roc_auc(pred_masked, target_masked)
+                    metrics['classification_report'] = GraphMetrics.classification_report_dict(pred_masked, target_masked)
                 
-                metrics = {
-                    'accuracy': GraphMetrics.accuracy(all_preds, all_targets),
-                    'f1_macro': GraphMetrics.f1_score_macro(all_preds, all_targets),
-                    'f1_micro': GraphMetrics.f1_score_micro(all_preds, all_targets),
-                }
-            else:
-                metrics = {'error': 'No predictions generated'}
+                # Add loss
+                criterion = torch.nn.CrossEntropyLoss()
+                metrics['loss'] = criterion(pred_masked, target_masked).item()
                 
         except Exception as e:
-            print(f"Graph classification evaluation error: {e}")
-            metrics = {'error': str(e)}
-        
-        return metrics
+            print(f"Evaluation error: {e}")
+            metrics = {
+                'accuracy': 0.0