separate tokenization and feature extraction

run_image_captioning_flax.py

@@ -680,6 +680,54 @@ def main():
 
         return bools
 
+    def tokenization_fn(examples, max_target_length):
+
+        captions = []
+        for caption in examples[caption_column]:
+                captions.append(caption.lower() + ' ' + tokenizer.eos_token)
+
+        targets = captions
+
+        model_inputs = {}
+
+        # Setup the tokenizer for targets
+        with tokenizer.as_target_tokenizer():
+            labels = tokenizer(
+                targets, max_length=max_target_length, padding="max_length", truncation=True, return_tensors="np"
+            )
+
+        model_inputs["labels"] = labels["input_ids"]
+        decoder_input_ids = shift_tokens_right_fn(
+            labels["input_ids"], config.pad_token_id, config.decoder_start_token_id
+        )
+        model_inputs["decoder_input_ids"] = np.asarray(decoder_input_ids)
+
+        # We need decoder_attention_mask so we can ignore pad tokens from loss
+        model_inputs["decoder_attention_mask"] = labels["attention_mask"]
+
+        model_inputs[image_column] = examples[image_column]
+
+        return model_inputs
+
+    def feature_extraction_fn(examples):
+
+        pixel_values = []
+
+        for image_file in examples[image_column]:
+            with Image.open(image_file) as image:
+                try:
+                    encoder_inputs = feature_extractor(images=image, return_tensors="np")
+                except:
+                    continue
+                pixel_values.append(encoder_inputs.pixel_values)
+
+        pixel_values = np.concatenate(pixel_values)
+
+        model_inputs = examples
+        model_inputs['pixel_values'] = pixel_values
+
+        return model_inputs
+
     # Setting padding="max_length" as we need fixed length inputs for jitted functions
     def preprocess_function(examples, max_target_length):
 
@@ -741,6 +789,16 @@ def main():
         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))
+        train_dataset = train_dataset.map(
+            tokenization_fn,
+            batched=True,
+            num_proc=data_args.preprocessing_num_workers,
+            # kept image paths
+            remove_columns=column_names.remove(image_column),
+            load_from_cache_file=not data_args.overwrite_cache,
+            desc=f"Running tokenizer on train dataset",
+            fn_kwargs={"max_target_length": data_args.max_target_length},
+        )
 
     if training_args.do_eval:
         if "validation" not in dataset:
@@ -750,6 +808,16 @@ def main():
         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))
+        eval_dataset = eval_dataset.map(
+            tokenization_fn,
+            batched=True,
+            num_proc=data_args.preprocessing_num_workers,
+            # kept image paths
+            remove_columns=column_names.remove(image_column),
+            load_from_cache_file=not data_args.overwrite_cache,
+            desc=f"Running tokenizer on validation dataset",
+            fn_kwargs={"max_target_length": data_args.val_max_target_length},
+        )
 
     if training_args.do_predict:
         if "test" not in dataset:
@@ -759,6 +827,16 @@ def main():
         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))
+        predict_dataset = predict_dataset.map(
+            tokenization_fn,
+            batched=True,
+            num_proc=data_args.preprocessing_num_workers,
+            # kept image paths
+            remove_columns=column_names.remove(image_column),
+            load_from_cache_file=not data_args.overwrite_cache,
+            desc=f"Running tokenizer on prediction dataset",
+            fn_kwargs={"max_target_length": data_args.val_max_target_length},
+        )
 
     # Split the dataset into several chunks - each chunk is processed (.map) without cache to create a
     # data loader separately (in a sequential order).
@@ -804,7 +882,6 @@ def main():
         else:
             indices = jnp.arange(len(ds))
 
-        # Temporarily set max_target_length for training or evaluation/prediction.
         max_target_length = data_args.max_target_length
         if split in ["valid", "test"]:
             max_target_length = data_args.val_max_target_length
@@ -825,14 +902,17 @@ def main():
             }
 
             _ds =_ds.map(
-                preprocess_function,
+                # preprocess_function,
+                feature_extraction_fn,
                 batched=True,
                 num_proc=data_args.preprocessing_num_workers,
-                remove_columns=column_names,
+                # remove_columns=column_names,
+                remove_columns=[image_column],
                 load_from_cache_file=not data_args.overwrite_cache,
                 features=features,
-                desc=f"Running tokenizer on {names[split]} dataset".replace("  ", " "),
-                fn_kwargs={"max_target_length": max_target_length},
+                # desc=f"Running tokenizer on {names[split]} dataset".replace("  ", " "),
+                desc=f"Running feature extraction on {names[split]} dataset".replace("  ", " "),
+                # fn_kwargs={"max_target_length": max_target_length},
             )
             _ds = _ds.with_format("numpy")