clean up config script

vit_gpt2/configuration_vit_gpt2.py

@@ -1,9 +1,10 @@
 import copy
 
-from transformers import GPT2Config, ViTConfig
+from transformers import ViTConfig, GPT2Config
 from transformers.configuration_utils import PretrainedConfig
 from transformers.utils import logging
 
+
 logger = logging.get_logger(__name__)
 
 
@@ -12,34 +13,28 @@ class ViTGPT2Config(PretrainedConfig):
     model_type = "vit-gpt2"
     is_composition = True
 
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-        if "vit_config" not in kwargs:
-            raise ValueError("`vit_config` can not be `None`.")
+    def __init__(self, vision_config_dict=None, text_config_dict=None, **kwargs):
+        super().__init__(text_config_dict=text_config_dict, vision_config_dict=vision_config_dict, **kwargs)
 
-        if "gpt2_config" not in kwargs:
-            raise ValueError("`gpt2_config` can not be `None`.")
+        if vision_config_dict is None:
+            vision_config_dict = {}
+            logger.info("vision_config_dict is None. initializing the ViTConfig with default values.")
 
-        vit_config = kwargs.pop("vit_config")
-        gpt2_config = kwargs.pop("gpt2_config")
+        if text_config_dict is None:
+            text_config_dict = {}
+            logger.info("text_config_dict is None. Initializing the GPT2Config with default values.")
 
-        self.vit_config = ViTConfig(**vit_config)
-        self.gpt2_config = GPT2Config(**gpt2_config)
+        self.vision_config = ViTConfig(**vision_config_dict)
+        self.text_config = GPT2Config(**text_config_dict)
 
     @classmethod
-    def from_vit_gpt2_configs(
-        cls, vit_config: PretrainedConfig, gpt2_config: PretrainedConfig, **kwargs
-    ):
-        return cls(
-            vit_config=vit_config.to_dict(),
-            gpt2_config=gpt2_config.to_dict(),
-            **kwargs
-        )
+    def from_vit_gpt2_configs(cls, vision_config: ViTConfig, text_config: GPT2Config, **kwargs):
+
+        return cls(vision_config_dict=vision_config.to_dict(), text_config_dict=text_config.to_dict(), **kwargs)
 
     def to_dict(self):
         output = copy.deepcopy(self.__dict__)
-        output["vit_config"] = self.vit_config.to_dict()
-        output["gpt2_config"] = self.gpt2_config.to_dict()
+        output["vision_config"] = self.vision_config.to_dict()
+        output["text_config"] = self.text_config.to_dict()
         output["model_type"] = self.__class__.model_type
-        return output
+        return output