Merge pull request #5323 from naem1023/feat/add-dataset-map-batch-size-argument

Add batch size of map function in the preprocessed dataset

src/llamafactory/data/loader.py

@@ -179,6 +179,9 @@ def _get_preprocessed_dataset(
             load_from_cache_file=(not data_args.overwrite_cache) or (training_args.local_process_index != 0),
             desc="Running tokenizer on dataset",
         )
+        if data_args.dataset_map_batch_size:
+            # Set the batch size conditionally without considering the default variable of the batch size in the map function
+            kwargs.update(batch_size=data_args.dataset_map_batch_size)
 
     dataset = dataset.map(
         preprocess_func,

src/llamafactory/hparams/data_args.py

@@ -113,6 +113,10 @@ class DataArguments:
         default=None,
         metadata={"help": "Path to save or load the tokenized datasets."},
     )
+    dataset_map_batch_size: Optional[int] = field(
+        default=None,
+        metadata={"help": "Batch size for dataset mapping."},
+    )
 
     def __post_init__(self):
         def split_arg(arg):