fix

run_image_captioning_flax.py

@@ -835,22 +835,23 @@ def main():
         }
     )
 
-    # If `block_size` is `0`, tokenization & image feature extraction is done before training
-    run_feat_ext_before_training = training_args.block_size == 0
+    # If `block_size` is `0`, tokenization & image feature extraction is done at the beginning
+    run_feat_ext_at_beginning = training_args.block_size == 0
     # Used in .map() below
-    function_kwarg = preprocess_fn if run_feat_ext_before_training else tokenization_fn
+    function_kwarg = preprocess_fn if run_feat_ext_at_beginning else tokenization_fn
     # `features` is used only for the final preprocessed dataset (for the performance purpose).
-    features_kwarg = features if run_feat_ext_before_training else None
+    features_kwarg = features if run_feat_ext_at_beginning else None
     # Keep `image_column` if the feature extraction is done during training
-    remove_columns_kwarg = [x for x in column_names if x != image_column or run_feat_ext_before_training]
-    processor_names = "tokenizer and feature extractor" if run_feat_ext_before_training else "tokenizer"
+    remove_columns_kwarg = [x for x in column_names if x != image_column or run_feat_ext_at_beginning]
+    processor_names = "tokenizer and feature extractor" if run_feat_ext_at_beginning else "tokenizer"
 
     if training_args.do_train:
         if "train" not in dataset:
             raise ValueError("--do_train requires a train dataset")
         train_dataset = dataset["train"]
-        # remove problematic examples
-        if not run_feat_ext_before_training:
+        # remove problematic examples 
+        # (if feature extraction is performed at the beginning, the filtering is done during preprocessing not here)
+        if not run_feat_ext_at_beginning:
             train_dataset = train_dataset.filter(filter_fn, batched=True, num_proc=data_args.preprocessing_num_workers)
         if data_args.max_train_samples is not None:
             train_dataset = train_dataset.select(range(data_args.max_train_samples))
@@ -865,15 +866,17 @@ def main():
             fn_kwargs={"max_target_length": data_args.max_target_length},
             features=features_kwarg,
         )
-        if run_feat_ext_before_training:
+        if run_feat_ext_at_beginning:
+            # set format (for performance) since the dataset is ready to be used
             train_dataset = train_dataset.with_format("numpy")
 
     if training_args.do_eval:
         if "validation" not in dataset:
             raise ValueError("--do_eval requires a validation dataset")
         eval_dataset = dataset["validation"]
-        # remove problematic examples
-        if not run_feat_ext_before_training:
+        # remove problematic examples 
+        # (if feature extraction is performed at the beginning, the filtering is done during preprocessing not here)
+        if not run_feat_ext_at_beginning:
             eval_dataset = eval_dataset.filter(filter_fn, batched=True, num_proc=data_args.preprocessing_num_workers)
         if data_args.max_eval_samples is not None:
             eval_dataset = eval_dataset.select(range(data_args.max_eval_samples))
@@ -888,15 +891,17 @@ def main():
             fn_kwargs={"max_target_length": data_args.val_max_target_length},
             features=features_kwarg,
         )
-        if run_feat_ext_before_training:
+        if run_feat_ext_at_beginning:
+            # set format (for performance) since the dataset is ready to be used
             eval_dataset = eval_dataset.with_format("numpy")
 
     if training_args.do_predict:
         if "test" not in dataset:
             raise ValueError("--do_predict requires a test dataset")
         predict_dataset = dataset["test"]
-        # remove problematic examples
-        if not run_feat_ext_before_training:
+        # remove problematic examples 
+        # (if feature extraction is performed at the beginning, the filtering is done during preprocessing not here)
+        if not run_feat_ext_at_beginning:
             predict_dataset = predict_dataset.filter(filter_fn, batched=True, num_proc=data_args.preprocessing_num_workers)
         if data_args.max_predict_samples is not None:
             predict_dataset = predict_dataset.select(range(data_args.max_predict_samples))
@@ -911,7 +916,8 @@ def main():
             fn_kwargs={"max_target_length": data_args.val_max_target_length},
             features=features_kwarg,
         )
-        if run_feat_ext_before_training:
+        if run_feat_ext_at_beginning:
+            # set format (for performance) since the dataset is ready to be used
             predict_dataset = predict_dataset.with_format("numpy")
 
     # Store some constant