Upload model.py with huggingface_hub

model.py

@@ -1,10 +1,6 @@
-from transformers import PreTrainedModel, PretrainedConfig, ImageProcessingMixin
-import torch
+from transformers import PreTrainedModel, PretrainedConfig
 import torch.nn as nn
 from torchvision.models import efficientnet_v2_s, EfficientNet_V2_S_Weights
-from torchvision import transforms
-from PIL import Image
-import numpy as np
 
 class CheckboxConfig(PretrainedConfig):
     model_type = "checkbox-classifier"
@@ -13,44 +9,6 @@ class CheckboxConfig(PretrainedConfig):
         super().__init__(num_labels=num_labels, **kwargs)
         self.dropout_rate = dropout_rate
 
-class CheckboxImageProcessor(ImageProcessingMixin):
-    """Simple image processor for checkbox classifier"""
-    
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-        self.size = {"height": 128, "width": 128}
-        self.image_mean = [0.485, 0.456, 0.406]
-        self.image_std = [0.229, 0.224, 0.225]
-        
-        self.transform = transforms.Compose([
-            transforms.Resize((self.size["height"], self.size["width"])),
-            transforms.ToTensor(),
-            transforms.Normalize(mean=self.image_mean, std=self.image_std)
-        ])
-    
-    def preprocess(self, images, **kwargs):
-        """Preprocess images for model input"""
-        if not isinstance(images, list):
-            images = [images]
-        
-        processed = []
-        for image in images:
-            if isinstance(image, str):
-                image = Image.open(image).convert('RGB')
-            elif isinstance(image, np.ndarray):
-                image = Image.fromarray(image).convert('RGB')
-            elif not isinstance(image, Image.Image):
-                raise ValueError(f"Unsupported image type: {type(image)}")
-            
-            processed.append(self.transform(image))
-        
-        # Stack into batch
-        pixel_values = torch.stack(processed)
-        return {"pixel_values": pixel_values}
-    
-    def __call__(self, images, **kwargs):
-        return self.preprocess(images, **kwargs)
-
 class CheckboxClassifier(PreTrainedModel):
     config_class = CheckboxConfig
     
@@ -58,7 +16,7 @@ class CheckboxClassifier(PreTrainedModel):
         super().__init__(config)
         self.num_labels = config.num_labels
         
-        self.backbone = efficientnet_v2_s(weights=None)  # Don't load pretrained weights here
+        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(
@@ -74,15 +32,6 @@ class CheckboxClassifier(PreTrainedModel):
             nn.Linear(256, config.num_labels)
         )
     
-    def forward(self, pixel_values, labels=None):
+    def forward(self, pixel_values):
         outputs = self.backbone(pixel_values)
-        
-        loss = None
-        if labels is not None:
-            loss_fct = nn.CrossEntropyLoss()
-            loss = loss_fct(outputs, labels)
-            
-        return {
-            "loss": loss,
-            "logits": outputs,
-        }
+        return {"logits": outputs}