Create gnn_predictor.py

gnn_predictor.py

@@ -0,0 +1,26 @@
+# gnn_predictor.py
+import torch
+import torch.nn.functional as F
+try:
+    from torch_geometric.nn import GCNConv
+    TORCH_GEOMETRIC_AVAILABLE = True
+except ImportError:
+    TORCH_GEOMETRIC_AVAILABLE = False
+
+class FailureGNN(torch.nn.Module):
+    def __init__(self, num_features=5, hidden=16):
+        super().__init__()
+        if TORCH_GEOMETRIC_AVAILABLE:
+            self.conv1 = GCNConv(num_features, hidden)
+            self.conv2 = GCNConv(hidden, 2)
+        else:
+            self.dummy = torch.nn.Linear(num_features, 2)
+
+    def forward(self, x, edge_index=None):
+        if TORCH_GEOMETRIC_AVAILABLE and edge_index is not None:
+            x = self.conv1(x, edge_index)
+            x = F.relu(x)
+            x = self.conv2(x, edge_index)
+        else:
+            x = self.dummy(x)
+        return F.log_softmax(x, dim=1)