Added files

__init__.py

@@ -0,0 +1 @@
+from .modeling_sagvit import SAGViTClassifier

hubconf.py

@@ -1,6 +1,6 @@
 dependencies = ['torch']
 
-from sag_vit_model import SAGViTClassifier  
+from modeling_sagvit import SAGViTClassifier  
 import torch
 
 def SAGViT(pretrained=False, **kwargs):

sag_vit_model.py → modeling_sagvit.py

@@ -1,11 +1,16 @@
 import torch
 from torch import nn
+from transformers import AutoConfig, PretrainedConfig, AutoModel, PreTrainedModel
+from transformers.models.auto import AutoConfig, CONFIG_MAPPING, MODEL_MAPPING
+from transformers.utils import logging
+from transformers.modeling_utils import ModelOutput
 from huggingface_hub import PyTorchModelHubMixin
 
 from torch_geometric.data import Batch
 from model_components import EfficientNetV2FeatureExtractor, GATGNN, TransformerEncoder, MLPBlock
 from graph_construction import build_graph_from_patches, build_graph_data_from_patches
 
+
 ###############################################################################
 # SAG-ViT Model:
 # This class combines:
@@ -15,7 +20,26 @@ from graph_construction import build_graph_from_patches, build_graph_data_from_p
 # 4) A final MLP classifier.
 ###############################################################################
 
-class SAGViTClassifier(nn.Module, PyTorchModelHubMixin):
+
+# Custom model registration
+class SAGViTConfig(PretrainedConfig):
+    model_type = "sagvit"
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.d_model = kwargs.get("d_model", 64)
+        self.dim_feedforward = kwargs.get("dim_feedforward", 64)
+        self.gcn_hidden = kwargs.get("gcn_hidden", 128)
+        self.gcn_out = kwargs.get("gcn_out", 64)
+        self.hidden_mlp_features = kwargs.get("hidden_mlp_features", 64)
+        self.in_channels = kwargs.get("in_channels", 2560)
+        self.nhead = kwargs.get("nhead", 4)
+        self.num_classes = kwargs.get("num_classes", 10)
+        self.num_layers = kwargs.get("num_layers", 2)
+        self.patch_size = kwargs.get("patch_size", (4, 4))
+
+
+class SAGViTClassifier(PreTrainedModel):
     """
     SAG-ViT: Scale-Aware Graph Attention Vision Transformer
 
@@ -32,41 +56,41 @@ class SAGViTClassifier(nn.Module, PyTorchModelHubMixin):
     Outputs:
     - out (Tensor): Classification logits (B, num_classes)
     """
-    def __init__(
-        self,
-        patch_size=(4,4),
-        num_classes=10,
-        d_model=64,
-        nhead=4,
-        num_layers=2,
-        dim_feedforward=64,
-        hidden_mlp_features=64,
-        in_channels=2560,  # Derived from patch dimensions and CNN output channels
-        gcn_hidden=128,
-        gcn_out=64
-    ):
-        super(SAGViTClassifier, self).__init__()
+
+    config_class = SAGViTConfig
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.patch_size = config.patch_size
+        self.num_classes = config.num_classes
 
         # CNN feature extractor (frozen pre-trained EfficientNetv2)
         self.cnn = EfficientNetV2FeatureExtractor()
 
         # Graph Attention Network to process patch embeddings
-        self.gcn = GATGNN(in_channels=in_channels, hidden_channels=gcn_hidden, out_channels=gcn_out)
+        self.gcn = GATGNN(
+            in_channels=config.in_channels,
+            hidden_channels=config.gcn_hidden,
+            out_channels=config.gcn_out,
+        )
 
         # Learnable positional embedding for Transformer input
-        self.positional_embedding = nn.Parameter(torch.randn(1, 1, d_model))
+        self.positional_embedding = nn.Parameter(torch.randn(1, 1, config.d_model))
         # Extra embedding token (similar to class token) to summarize global info
-        self.extra_embedding = nn.Parameter(torch.randn(1, d_model))
+        self.extra_embedding = nn.Parameter(torch.randn(1, config.d_model))
 
         # Transformer encoder to capture long-range global dependencies
-        self.transformer_encoder = TransformerEncoder(d_model, nhead, num_layers, dim_feedforward)
+        self.transformer_encoder = TransformerEncoder(
+            d_model=config.d_model, 
+            nhead=config.nhead, 
+            num_layers=config.num_layers, 
+            dim_feedforward=config.dim_feedforward,
+        )
 
         # MLP classification head
-        self.mlp = MLPBlock(d_model, hidden_mlp_features, num_classes)
+        self.mlp = MLPBlock(config.d_model, config.hidden_mlp_features, config.num_classes)
 
-        self.patch_size = patch_size
-
-    def forward(self, x):
+    def forward(self, x, **kwargs):
         # Step 1: High-fidelity feature extraction from CNN
         feature_map = self.cnn(x)
 
@@ -103,5 +127,9 @@ class SAGViTClassifier(nn.Module, PyTorchModelHubMixin):
         x_pooled = x_trans.mean(dim=1)  # (B, D)
 
         # Classification
-        out = self.mlp(x_pooled)
-        return out
+        logits = self.mlp(x_pooled)
+        return ModelOutput(logits=logits)
+
+# Register custom model and config
+CONFIG_MAPPING.register("sagvit", SAGViTConfig)
+MODEL_MAPPING.register(SAGViTConfig, SAGViTClassifier)

push_model_to_hfhub.py

@@ -0,0 +1,9 @@
+from transformers import AutoConfig, AutoModel
+from modeling_sagvit import SAGViTClassifier
+
+# Initialize config and model
+config = AutoConfig.from_pretrained("shravvvv/SAG-ViT")
+model = AutoModel.from_pretrained("shravvvv/SAG-ViT", config=config)
+
+# Push model to the Hub
+model.push_to_hub("shravvvv/SAG-ViT")

register_model.py

@@ -3,7 +3,7 @@ from transformers.models.auto.configuration_auto import CONFIG_MAPPING
 from transformers.models.auto.modeling_auto import MODEL_MAPPING
 
 from sagvit_config import SAGViTConfig
-from sag_vit_model import SAGViTClassifier
+from modeling_sagvit import SAGViTClassifier
 
 # Register the configuration
 CONFIG_MAPPING.register("sagvit", SAGViTConfig)

sagvit_config.py

@@ -1,28 +0,0 @@
-from transformers import PretrainedConfig
-
-class SAGViTConfig(PretrainedConfig):
-    model_type = "sagvit"
-
-    def __init__(self,
-                 d_model=64,
-                 dim_feedforward=64,
-                 gcn_hidden=128,
-                 gcn_out=64,
-                 hidden_mlp_features=64,
-                 in_channels=2560,
-                 nhead=4,
-                 num_classes=10,
-                 num_layers=2,
-                 patch_size=(4, 4),
-                 **kwargs):
-        super().__init__(**kwargs)
-        self.d_model = d_model
-        self.dim_feedforward = dim_feedforward
-        self.gcn_hidden = gcn_hidden
-        self.gcn_out = gcn_out
-        self.hidden_mlp_features = hidden_mlp_features
-        self.in_channels = in_channels
-        self.nhead = nhead
-        self.num_classes = num_classes
-        self.num_layers = num_layers
-        self.patch_size = patch_size

tests/test_sag_vit_model.py

@@ -1,6 +1,6 @@
 import unittest
 import torch
-from sag_vit_model import SAGViTClassifier
+from modeling_sagvit import SAGViTClassifier
 
 class TestSAGViTModel(unittest.TestCase):
     def test_forward_pass(self):

tests/test_train.py

@@ -3,7 +3,7 @@ from unittest.mock import MagicMock, patch
 import torch
 import torch.nn as nn
 from train import train_model
-from sag_vit_model import SAGViTClassifier
+from modeling_sagvit import SAGViTClassifier
 
 class TestTrain(unittest.TestCase):
     @patch("train.optim.Adam")

train.py

@@ -8,7 +8,7 @@ from sklearn.metrics import (precision_score, recall_score, f1_score,
                              roc_auc_score, cohen_kappa_score, matthews_corrcoef,
                              confusion_matrix)
 
-from sag_vit_model import SAGViTClassifier
+from modeling_sagvit import SAGViTClassifier
 from data_loader import get_dataloaders
 
 #####################################################################