Refactoring models

src/models.py

@@ -70,9 +70,7 @@ class TinyCLIPVisionEncoder(PreTrainedModel):
 
     def __init__(self, config: TinyCLIPVisionConfig):
         super().__init__(config)
-        self.base = timm.create_model(config.vision_model, num_classes=0)
-        timm_config = data.resolve_data_config({}, model=self.base)
-        self.transform = data.transforms_factory.create_transform(**timm_config)
+
         self.projection = projection_layers(
             self.base.num_features, config.embed_dims, config.projection_layers
         )

src/trainer.py

@@ -0,0 +1,7 @@
+from src import data
+from src import config
+from src import vision_model
+
+
+def train(config: config.TrainerConfig):
+    train_dl, valid_dl = data.get_dataset()

src/vision_model.py

@@ -0,0 +1,11 @@
+import timm
+from timm import data
+
+from src import config
+
+
+def get_vision_base_and_transform(config: config.TrainerConfig):
+    base = timm.create_model(config._model_config.vision_model, num_classes=0)
+    timm_config = data.resolve_data_config({}, model=base)
+    transform = data.transforms_factory.create_transform(**timm_config)
+    return base, transform