update model script

vit_gpt2/modeling_flax_vit_gpt2_lm.py

@@ -1,4 +1,5 @@
-from typing import Callable, Optional, Tuple
+import os
+from typing import Callable, Optional, Tuple, Union
 
 import flax.linen as nn
 import jax
@@ -32,8 +33,8 @@ class FlaxViTGPT2LMModule(nn.Module):
 
     def setup(self):
 
-        self.encoder = FlaxViTModule(self.config.vit_config, dtype=self.dtype)
-        self.decoder = FlaxGPT2LMHeadModule(self.config.gpt2_config, dtype=self.dtype)
+        self.encoder = FlaxViTModule(self.config.vision_config, dtype=self.dtype)
+        self.decoder = FlaxGPT2LMHeadModule(self.config.text_config, dtype=self.dtype)
 
     def _get_encoder_module(self):
         return self.encoder
@@ -147,7 +148,7 @@ class FlaxViTGPT2LMPreTrainedModel(FlaxPreTrainedModel):
     ):
         if input_shape is None:
             input_shape = (
-                (1, config.vit_config.image_size, config.vit_config.image_size, 3),
+                (1, config.vision_config.image_size, config.vision_config.image_size, 3),
                 (1, 1),
             )
 
@@ -164,7 +165,7 @@ class FlaxViTGPT2LMPreTrainedModel(FlaxPreTrainedModel):
         attention_mask = None
         decoder_input_ids = jnp.zeros(decoder_input_shape, dtype="i4")
         # make sure initialization pass will work for FlaxBartForSequenceClassificationModule
-        decoder_input_ids = jax.ops.index_update(decoder_input_ids, (..., -1), self.config.eos_token_id)
+        decoder_input_ids = jax.ops.index_update(decoder_input_ids, (..., -1), self.config.text_config.eos_token_id)
         decoder_attention_mask = jnp.ones_like(decoder_input_ids)
 
         batch_size, sequence_length = decoder_input_ids.shape
@@ -221,11 +222,11 @@ class FlaxViTGPT2LMPreTrainedModel(FlaxPreTrainedModel):
         params: dict = None,
         dropout_rng: PRNGKey = None,
     ):
-        output_attentions = (output_attentions if output_attentions is not None else self.config.vit_config.output_attentions)
+        output_attentions = (output_attentions if output_attentions is not None else self.config.vision_config.output_attentions)
         output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.vit_config.output_hidden_states
+            output_hidden_states if output_hidden_states is not None else self.config.vision_config.output_hidden_states
         )
-        return_dict = return_dict if return_dict is not None else self.config.vit_config.return_dict
+        return_dict = return_dict if return_dict is not None else self.config.vision_config.return_dict
 
         # (`transpose` is done in `FlaxViTPreTrainedModel.__call__()`, so we do the same here.)
         pixel_values = jnp.transpose(pixel_values, (0, 2, 3, 1))
@@ -267,12 +268,12 @@ class FlaxViTGPT2LMPreTrainedModel(FlaxPreTrainedModel):
     ):
 
         output_attentions = (
-            output_attentions if output_attentions is not None else self.config.gpt2_config.output_attentions
+            output_attentions if output_attentions is not None else self.config.text_config.output_attentions
         )
         output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.gpt2_config.output_hidden_states
+            output_hidden_states if output_hidden_states is not None else self.config.text_config.output_hidden_states
         )
-        return_dict = return_dict if return_dict is not None else self.config.gpt2_config.return_dict
+        return_dict = return_dict if return_dict is not None else self.config.text_config.return_dict
 
         encoder_hidden_states = encoder_outputs[0]
         if encoder_attention_mask is None:
@@ -486,71 +487,78 @@ class FlaxViTGPT2LMForConditionalGeneration(FlaxViTGPT2LMPreTrainedModel):
         return model_kwargs
 
     @classmethod
-    def from_vit_gpt2_pretrained(
+    def from_vision_text_pretrained(
         cls,
-        vit_model_name_or_path: str = None,
-        gpt2_model_name_or_path: str = None,
+        vision_pretrained_model_name_or_path: Optional[Union[str, os.PathLike]],
+        text_pretrained_model_name_or_path: Optional[Union[str, os.PathLike]],
         *model_args,
         **kwargs,
     ) -> FlaxViTGPT2LMPreTrainedModel:
 
-        kwargs_gpt2 = {
-            argument[len("gpt2_") :]: value
-            for argument, value in kwargs.items()
-            if argument.startswith("gpt2_")
+        vision_kwargs = {
+            kwarg[len("vision_"):]: value
+            for kwarg, value in kwargs.items()
+            if kwarg.startswith("vision_")
         }
 
-        kwargs_vit = {
-            argument[len("vit_") :]: value
-            for argument, value in kwargs.items()
-            if argument.startswith("vit_")
+        text_kwargs = {
+            kwarg[len("text_"):]: value
+            for kwarg, value in kwargs.items()
+            if kwarg.startswith("text_")
         }
 
-        # remove gpt2, vit kwargs from kwargs
-        for key in kwargs_gpt2.keys():
-            del kwargs["gpt2_" + key]
-        for key in kwargs_vit.keys():
-            del kwargs["vit_" + key]
+        # remove vit & gpt2 kwargs from kwargs
+        for key in vision_kwargs.keys():
+            del kwargs["vision_" + key]
+        for key in text_kwargs.keys():
+            del kwargs["text_" + key]
 
-        # Load and initialize the gpt2 and vit model
-        gpt2_model = kwargs_gpt2.pop("model", None)
-        if gpt2_model is None:
+        vision_model_args = vision_kwargs.pop('model_args', None)
+        text_model_args = text_kwargs.pop('model_args', None)
+
+        # Load and initialize the vit & gpt2 model
+        vision_model = vision_kwargs.pop("model", None)
+        text_model = text_kwargs.pop("model", None)
+
+        if vision_model is None:
             assert (
-                gpt2_model_name_or_path is not None
-            ), "If `model` is not defined as an argument, a `gpt2_model_name_or_path` has to be defined"
+                vision_pretrained_model_name_or_path is not None
+            ), "If `model` is not defined as an argument, a `vision_pretrained_model_name_or_path` has to be defined"
 
-            if "config" not in kwargs_gpt2:
-                gpt2_config = GPT2Config.from_pretrained(gpt2_model_name_or_path)
-                kwargs_gpt2["config"] = gpt2_config
+            if "config" not in vision_kwargs:
+                vision_config = ViTConfig.from_pretrained(vision_pretrained_model_name_or_path)
+                vision_kwargs["config"] = vision_config
 
-            kwargs_gpt2["config"].add_cross_attention = True
-            gpt2_model = FlaxGPT2LMHeadModel.from_pretrained(
-                gpt2_model_name_or_path, *model_args, **kwargs_gpt2
+            # TODO: How to deal with model_args?
+            vision_model = FlaxViTModel.from_pretrained(
+                vision_pretrained_model_name_or_path, *vision_model_args, **vision_kwargs
             )
 
-        vit_model = kwargs_vit.pop("model", None)
-        if vit_model is None:
+        if text_model is None:
             assert (
-                vit_model_name_or_path is not None
-            ), "If `model` is not defined as an argument, a `vit_model_name_or_path` has to be defined"
+                text_pretrained_model_name_or_path is not None
+            ), "If `model` is not defined as an argument, a `text_pretrained_model_name_or_path` has to be defined"
+
+            if "config" not in text_kwargs:
+                text_config = GPT2Config.from_pretrained(text_pretrained_model_name_or_path)
+                text_kwargs["config"] = text_config
 
-            if "config" not in kwargs_vit:
-                vit_config = ViTConfig.from_pretrained(vit_model_name_or_path)
-                kwargs_vit["config"] = vit_config
+            text_kwargs["config"].add_cross_attention = True
 
-            vit_model = FlaxViTModel.from_pretrained(
-                vit_model_name_or_path, *model_args, **kwargs_vit
+            # TODO: How to deal with model_args?
+            text_model = FlaxGPT2LMHeadModel.from_pretrained(
+                text_pretrained_model_name_or_path, *text_model_args, **text_kwargs
             )
 
         # instantiate config with corresponding kwargs
         dtype = kwargs.pop("dtype", jnp.float32)
-        config = ViTGPT2Config.from_vit_gpt2_configs(
-            vit_model.config, gpt2_model.config, **kwargs
+        config = ViTGPT2Config.from_vision_text_configs(
+            vision_model.config, text_model.config, **kwargs
         )
 
         # init model
         model = cls(config, *model_args, dtype=dtype, **kwargs)
-        model.params["model"]["encoder"] = vit_model.params
-        model.params["model"]["decoder"] = gpt2_model.params
+        model.params["model"]["encoder"] = vision_model.params
+        model.params["model"]["decoder"] = text_model.params
 
         return model