improve doc

run_image_captioning_flax.py

@@ -365,7 +365,7 @@ def data_loader(rng: jax.random.PRNGKey, dataset: Dataset, batch_size: int, shuf
     Returns batches of size `batch_size` from truncated `dataset`, sharded over all local devices.
     Shuffle batches if `shuffle` is `True`.
     """
-    steps = len(dataset) // batch_size
+    steps = len(dataset) // batch_size  # Skip incomplete batch.
 
     if shuffle:
         batch_idx = jax.random.permutation(rng, len(dataset))
@@ -924,7 +924,7 @@ def main():
         num_test_examples = len(predict_dataset)
         test_steps = num_test_examples // eval_batch_size
 
-    def get_batch_iter(
+    def blockwise_data_loader(
             rng: jax.random.PRNGKey,
             ds: Dataset,
             block_size: int,
@@ -933,6 +933,13 @@ def main():
             keep_in_memory: bool = False,
             split: str = ""
     ):
+        """
+        Wrap the simple `data_loader` in a block-wise way if `block_size` > 0, else it's the same as `data_loader`.
+
+        If `block_size` > 0, it requires `ds` to have a column that gives image paths in order to perform image feature
+        extraction (with the column name being specified by `image_column`). The tokenization should be done before
+        training in this case.
+        """
 
         if shuffle:
             indices = jax.random.permutation(rng, len(ds))
@@ -1197,7 +1204,7 @@ def main():
         preds = []
         labels = []
 
-        batches = get_batch_iter(rng, dataset, block_size=training_args.block_size, batch_size=eval_batch_size, keep_in_memory=False, shuffle=False, split=split)
+        batches = blockwise_data_loader(rng, dataset, block_size=training_args.block_size, batch_size=eval_batch_size, keep_in_memory=False, shuffle=False, split=split)
         steps = len(dataset) // eval_batch_size
         for _ in tqdm(range(steps), desc=f"{'Predicting' if split == 'test' else 'Evaluating'}...", position=2, leave=False):
             # Model forward
@@ -1305,7 +1312,7 @@ def main():
 
             train_metrics = []
 
-            train_batches = get_batch_iter(input_rng, train_dataset, block_size=training_args.block_size, batch_size=train_batch_size, keep_in_memory=True, shuffle=True, split="train")
+            train_batches = blockwise_data_loader(input_rng, train_dataset, block_size=training_args.block_size, batch_size=train_batch_size, keep_in_memory=True, shuffle=True, split="train")
 
             # train
             for (batch_idx, _) in enumerate(tqdm(range(steps_per_epoch), desc="Training...", position=1, leave=False)):