update 7

run_image_captioning_flax_reduced.py

@@ -510,20 +510,18 @@ def main():
     if decoder_config.pad_token_id is None:
         decoder_config.pad_token_id = decoder_config.eos_token_id
 
-    config = VisionEncoderDecoderConfig.from_encoder_decoder_configs(encoder_config, decoder_config)
-    # Necessary for Flax's generate()
-    config.decoder_start_token_id = config.decoder.decoder_start_token_id
-
     model = FlaxVisionEncoderDecoderModel.from_encoder_decoder_pretrained(
         encoder_pretrained_model_name_or_path=model_args.encoder_model_name_or_path,
         decoder_pretrained_model_name_or_path=model_args.decoder_model_name_or_path,
-        encoder_config=config.encoder,
-        decoder_config=config.decoder,
+        encoder_config=encoder_config,
+        decoder_config=decoder_config,
         encoder_seed=training_args.seed,
         decoder_seed=training_args.seed,
         encoder_dtype=getattr(jnp, model_args.dtype),
         decoder_dtype=getattr(jnp, model_args.dtype),
     )
+    # Necessary for Flax's generate()
+    model.config.decoder_start_token_id = decoder_config.decoder_start_token_id
 
     if model_args.feature_extractor_name:
         feature_extractor = AutoFeatureExtractor.from_pretrained(
@@ -553,7 +551,7 @@ def main():
             "You are instantiating a new tokenizer from scratch. This is not supported by this script."
             "You can do it from another script, save it, and load it from here, using --tokenizer_name."
         )
-    tokenizer.pad_token = tokenizer.convert_ids_to_tokens(config.decoder.pad_token_id)
+    tokenizer.pad_token = tokenizer.convert_ids_to_tokens(model.config.decoder.pad_token_id)
 
     # Preprocessing the datasets.
     # We need to tokenize inputs and targets.
@@ -628,7 +626,7 @@ def main():
 
         model_inputs["labels"] = labels["input_ids"]
         decoder_input_ids = shift_tokens_right_fn(
-            labels["input_ids"], config.decoder.pad_token_id, config.decoder.decoder_start_token_id
+            labels["input_ids"], model.config.decoder.pad_token_id, model.config.decoder_start_token_id
         )
         model_inputs["decoder_input_ids"] = np.asarray(decoder_input_ids)
 
@@ -687,9 +685,9 @@ def main():
         {
             "pixel_values": datasets.Array3D(
                 shape=(
-                    getattr(config.encoder, "num_channels", 3),
-                    config.encoder.image_size,
-                    config.encoder.image_size,
+                    getattr(model.config.encoder, "num_channels", 3),
+                    model.config.encoder.image_size,
+                    model.config.encoder.image_size,
                 ),
                 dtype="float32",
             ),