Upload models_arch.py with huggingface_hub

models_arch.py

@@ -0,0 +1,34 @@
+"""
+GNN architectures for AntioxFP: AttentiveFP (primary model).
+Copied from the research codebase and trimmed to inference-only.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch_geometric.nn import AttentiveFP
+
+
+NODE_DIM = 40  # atom feature dimension
+EDGE_DIM = 6   # bond feature dimension
+
+
+class AttentiveFPModel(nn.Module):
+    def __init__(self, hidden=200, num_layers=2, num_timesteps=2, dropout=0.2):
+        super().__init__()
+        self.gnn = AttentiveFP(
+            in_channels=NODE_DIM,
+            hidden_channels=hidden,
+            out_channels=1,
+            edge_dim=EDGE_DIM,
+            num_layers=num_layers,
+            num_timesteps=num_timesteps,
+            dropout=dropout,
+        )
+
+    def forward(self, x, edge_index=None, edge_attr=None, batch=None, **kwargs):
+        if hasattr(x, 'edge_index'):
+            data = x
+            x, edge_index, edge_attr, batch = (
+                data.x, data.edge_index, data.edge_attr, data.batch)
+        return self.gnn(x, edge_index, edge_attr, batch).squeeze(-1)