Upload model.py with huggingface_hub

model.py

@@ -0,0 +1,96 @@
+"""
+GNN graph classifiers for Android ransomware detection.
+
+Architectures: GIN, GCN, GAT — all using BatchNorm + global mean pooling.
+Uses `PyTorchModelHubMixin` for native `from_pretrained()` / `save_pretrained()`.
+
+Usage:
+    from model import GINClassifier
+
+    # Load from a local directory saved by save_pretrained()
+    model = GINClassifier.from_pretrained("./GIN/internal_only")
+
+    # Or load directly from the Hugging Face Hub
+    model = GINClassifier.from_pretrained(
+        "USER/android-ransomware-gnn-baseline",
+        subfolder="GIN/internal_only",
+    )
+    model.eval()
+"""
+
+import torch
+import torch.nn as nn
+from torch_geometric.nn import GINConv, GCNConv, GATConv, global_mean_pool
+from huggingface_hub import PyTorchModelHubMixin
+
+
+class GINClassifier(nn.Module, PyTorchModelHubMixin):
+    """Graph Isomorphism Network for graph-level binary classification."""
+
+    def __init__(self, in_dim=5, hidden=128, num_layers=3, num_classes=2, dropout=0.5):
+        super().__init__()
+        self.convs = nn.ModuleList()
+        self.bns = nn.ModuleList()
+        for i in range(num_layers):
+            dim_in = in_dim if i == 0 else hidden
+            mlp = nn.Sequential(
+                nn.Linear(dim_in, hidden), nn.ReLU(), nn.Linear(hidden, hidden),
+            )
+            self.convs.append(GINConv(mlp))
+            self.bns.append(nn.BatchNorm1d(hidden))
+        self.classifier = nn.Sequential(
+            nn.Linear(hidden, hidden), nn.ReLU(), nn.Dropout(dropout),
+            nn.Linear(hidden, num_classes),
+        )
+
+    def forward(self, x, edge_index, batch):
+        for conv, bn in zip(self.convs, self.bns):
+            x = torch.relu(bn(conv(x, edge_index)))
+        return self.classifier(global_mean_pool(x, batch))
+
+
+class GCNClassifier(nn.Module, PyTorchModelHubMixin):
+    """GCN baseline for graph-level classification."""
+
+    def __init__(self, in_dim=5, hidden=128, num_layers=3, num_classes=2, dropout=0.5):
+        super().__init__()
+        self.convs = nn.ModuleList()
+        self.bns = nn.ModuleList()
+        for i in range(num_layers):
+            dim_in = in_dim if i == 0 else hidden
+            self.convs.append(GCNConv(dim_in, hidden))
+            self.bns.append(nn.BatchNorm1d(hidden))
+        self.classifier = nn.Sequential(
+            nn.Linear(hidden, hidden), nn.ReLU(), nn.Dropout(dropout),
+            nn.Linear(hidden, num_classes),
+        )
+
+    def forward(self, x, edge_index, batch):
+        for conv, bn in zip(self.convs, self.bns):
+            x = torch.relu(bn(conv(x, edge_index)))
+        return self.classifier(global_mean_pool(x, batch))
+
+
+class GATClassifier(nn.Module, PyTorchModelHubMixin):
+    """GAT baseline for graph-level classification."""
+
+    def __init__(self, in_dim=5, hidden=128, num_layers=3, num_classes=2, dropout=0.5, heads=4):
+        super().__init__()
+        self.convs = nn.ModuleList()
+        self.bns = nn.ModuleList()
+        for i in range(num_layers):
+            dim_in = in_dim if i == 0 else hidden
+            self.convs.append(GATConv(dim_in, hidden // heads, heads=heads, concat=True))
+            self.bns.append(nn.BatchNorm1d(hidden))
+        self.classifier = nn.Sequential(
+            nn.Linear(hidden, hidden), nn.ReLU(), nn.Dropout(dropout),
+            nn.Linear(hidden, num_classes),
+        )
+
+    def forward(self, x, edge_index, batch):
+        for conv, bn in zip(self.convs, self.bns):
+            x = torch.relu(bn(conv(x, edge_index)))
+        return self.classifier(global_mean_pool(x, batch))
+
+
+MODEL_REGISTRY = {"GIN": GINClassifier, "GCN": GCNClassifier, "GAT": GATClassifier}