more changes

.gitignore

@@ -7,4 +7,5 @@ model_weights.pth
 # Python metadata
 venv
 __pycache__
-symmetric_test.egg-info
+symmetric_test.egg-info
+.ruff_cache

config.json

@@ -1,9 +1,14 @@
 {
     "architectures": ["DigitClassifier"],
-    "model_type": "pytorch",
+    "model_type": "custom",
     "num_labels": 10,
     "id2label": {
-      "0": "0", "1": "1", "2": "2", "3": "3", "4": "4", "5": "5", "6": "6", "7": "7", "8": "8", "9": "9"
+      "0": "0", "1": "1", "2": "2", "3": "3", "4": "4",
+      "5": "5", "6": "6", "7": "7", "8": "8", "9": "9"
     },
-    "preprocessor": "symmetric_test.preprocessor.get_preprocessor"
+    "preprocessor": "symmetric_test.preprocessor.get_transform",
+    "auto_map": {
+      "AutoModel": "symmetric_test.model.DigitClassifier",
+      "AutoConfig": "symmetric_test.model.DigitClassifierConfig"
+    }
   }

model_weights.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:7dc1c5d0edd67dcbd3746a7d79567c4825cf261944f7be55ed55d1386d3b7339
-size 903624

symmetric_test/model.py

@@ -1,9 +1,20 @@
+import torch
 import torch.nn as nn
+from transformers import PretrainedConfig
+
+
+class DigitClassifierConfig(PretrainedConfig):
+    model_type = "custom"
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.num_labels = 10
 
 
 class DigitClassifier(nn.Module):
-    def __init__(self):
+    def __init__(self, config):
         super().__init__()
+        self.config = config
         self.conv_block = nn.Sequential(
             nn.Conv2d(1, 32, 3),
             nn.ReLU(),
@@ -20,3 +31,12 @@ class DigitClassifier(nn.Module):
         x = self.conv_block(x)
         x = x.view(x.size(0), -1)
         return self.classifier(x)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        config = DigitClassifierConfig()
+        model = cls(config)
+        model.load_state_dict(
+            torch.load(f"{pretrained_model_name_or_path}/pytorch_model.bin")
+        )
+        return model

symmetric_test/preprocessor.py

@@ -1,12 +1,9 @@
 from torchvision import transforms
 
-
-def get_preprocessor():
-    return transforms.Compose(
-        [
-            transforms.Resize((28, 28)),
-            transforms.Grayscale(num_output_channels=1),
-            transforms.ToTensor(),
-            transforms.Normalize((0.1307,), (0.3081,)),
-        ]
-    )
+def get_transform():
+    return transforms.Compose([
+        transforms.Resize((28, 28)),
+        transforms.Grayscale(num_output_channels=1),
+        transforms.ToTensor(),
+        transforms.Normalize((0.1307,), (0.3081,))
+    ])

symmetric_test/train.py

@@ -2,7 +2,7 @@ import torch
 import torch.nn as nn
 from torch import optim
 from torchvision import datasets, transforms
-from huggingface_hub import HfApi, Repository
+from huggingface_hub import HfApi
 from .model import DigitClassifier
 
 # Config (better to put in separate config.yaml)