Modified config to transformers standard

src/config.py

@@ -1,6 +1,7 @@
 import pathlib
 
 import pydantic
+from transformers import PretrainedConfig
 
 MAX_DOWNLOAD_TIME = 0.2
 
@@ -16,6 +17,74 @@ class DataConfig(pydantic.BaseModel):
     dataset: str = small_dataset
 
 
+class TinyCLIPTextConfig(PretrainedConfig):
+    model_type = "text"
+
+    def __init__(
+        self,
+        text_model: str = "microsoft/xtremedistil-l6-h256-uncased",
+        projection_layers: int = 3,
+        embed_dims: int = 512,
+        max_len: int = 128,
+        cls_type: bool = True,
+        **kwargs,
+    ):
+        self.text_model = text_model
+        self.projection_layers = projection_layers
+        self.embed_dims = embed_dims
+        self.max_len = max_len
+        self.cls_type = cls_type
+        super().__init__(**kwargs)
+
+
+class TinyCLIPVisionConfig(PretrainedConfig):
+    model_type = "vision"
+
+    def __init__(
+        self,
+        vision_model: str = "edgenext_small",
+        projection_layers: int = 3,
+        embed_dims: int = 512,
+        **kwargs,
+    ):
+        self.vision_model = vision_model
+        self.projection_layers = projection_layers
+        self.embed_dims = embed_dims
+        super().__init__(**kwargs)
+
+
+class TinyCLIPConfig(PretrainedConfig):
+    model_type = "clip"
+
+    def __init__(
+        self,
+        text_model: str = "microsoft/xtremedistil-l6-h256-uncased",
+        vision_model: str = "edgenext_small",
+        projection_layers: int = 3,
+        embed_dim: int = 512,
+        max_len: int = 128,
+        cls_type: bool = True,
+        freeze_vision_base: bool = False,
+        freeze_text_base: bool = False,
+        loss_type: str = "cyclip",
+        **kwargs,
+    ):
+        self.text_config = TinyCLIPTextConfig(
+            text_model=text_model,
+            projection_layers=projection_layers,
+            embed_dims=embed_dim,
+            max_len=max_len,
+            cls_type=cls_type,
+        )
+        self.vision_config = TinyCLIPVisionConfig(
+            vision_model=vision_model, projection_layers=projection_layers, embed_dims=embed_dim
+        )
+        self.freeze_vision_base = freeze_vision_base
+        self.freeze_text_base = freeze_text_base
+        self.loss_type = loss_type
+        super().__init__(**kwargs)
+
+
 class ModelConfig(pydantic.BaseModel):
     text_model: str = "microsoft/xtremedistil-l6-h256-uncased"  # 51 mb
     vision_model: str = "edgenext_small"  # 20 mb