update 11

run_image_captioning_flax_reduced.py

@@ -507,12 +507,6 @@ def main():
     decoder_config.is_decoder = True
     decoder_config.add_cross_attention = True
 
-    # GPT2 only has bos/eos token but not decoder_start/pad token
-    if decoder_config.decoder_start_token_id is None:
-        decoder_config.decoder_start_token_id = decoder_config.bos_token_id
-    if decoder_config.pad_token_id is None:
-        decoder_config.pad_token_id = decoder_config.eos_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,
@@ -523,10 +517,19 @@ def main():
         encoder_dtype=getattr(jnp, model_args.dtype),
         decoder_dtype=getattr(jnp, model_args.dtype),
     )
-    # necessary to make Flax's generate() work
+
+    # GPT2 only has bos/eos tokens but not decoder_start/pad tokens
+    decoder_start_token_id = decoder_config.decoder_start_token_id
+    pad_token_id = decoder_config.pad_token_id
+    if decoder_start_token_id is None:
+        decoder_config.pad_token_id = decoder_config.bos_token_id
+    if pad_token_id is None:
+        pad_token_id = decoder_config.pad_token_id
+            
+    # This is necessary to make Flax's generate() work
     model.config.eos_token_id = decoder_config.eos_token_id
-    model.config.decoder_start_token_id = decoder_config.decoder_start_token_id
-    model.config.pad_token_id = decoder_config.pad_token_id
+    model.config.decoder_start_token_id = decoder_start_token_id
+    model.config.pad_token_id = pad_token_id
 
     if model_args.feature_extractor_name:
         feature_extractor = AutoFeatureExtractor.from_pretrained(
@@ -556,7 +559,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(model.config.decoder.pad_token_id)
+    tokenizer.pad_token = tokenizer.convert_ids_to_tokens(model.config.pad_token_id)
 
     # Preprocessing the datasets.
     # We need to tokenize inputs and targets.
@@ -631,7 +634,7 @@ def main():
 
         model_inputs["labels"] = labels["input_ids"]
         decoder_input_ids = shift_tokens_right_fn(
-            labels["input_ids"], model.config.decoder.pad_token_id, model.config.decoder_start_token_id
+            labels["input_ids"], model.config.pad_token_id, model.config.decoder_start_token_id
         )
         model_inputs["decoder_input_ids"] = np.asarray(decoder_input_ids)