Upload model.py with huggingface_hub

model.py

@@ -0,0 +1,37 @@
+from transformers import PreTrainedModel, PretrainedConfig
+import torch.nn as nn
+from torchvision.models import efficientnet_v2_s, EfficientNet_V2_S_Weights
+
+class CheckboxConfig(PretrainedConfig):
+    model_type = "checkbox-classifier"
+    
+    def __init__(self, num_labels=2, dropout_rate=0.3, **kwargs):
+        super().__init__(num_labels=num_labels, **kwargs)
+        self.dropout_rate = dropout_rate
+
+class CheckboxClassifier(PreTrainedModel):
+    config_class = CheckboxConfig
+    
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        
+        self.backbone = efficientnet_v2_s(weights=EfficientNet_V2_S_Weights.IMAGENET1K_V1)
+        num_features = self.backbone.classifier[1].in_features
+        
+        self.backbone.classifier = nn.Sequential(
+            nn.Dropout(config.dropout_rate),
+            nn.Linear(num_features, 512),
+            nn.SiLU(inplace=True),
+            nn.BatchNorm1d(512),
+            nn.Dropout(config.dropout_rate),
+            nn.Linear(512, 256),
+            nn.SiLU(inplace=True),
+            nn.BatchNorm1d(256),
+            nn.Dropout(config.dropout_rate/2),
+            nn.Linear(256, config.num_labels)
+        )
+    
+    def forward(self, pixel_values):
+        outputs = self.backbone(pixel_values)
+        return {"logits": outputs}