refactor: Added logging options to CombinedDataset and its factory method

Author: le1nux

src/modalities/dataloader/dataset.py

@@ -451,6 +451,9 @@ def __init__(self, datasets: list[Dataset], log_chunk_switch: bool = False, log_
 
         Args:
             datasets (list[Dataset]): A list of datasets to combine.
+            log_chunk_switch (bool, optional): Whether to log when switching between dataset chunks. Defaults to False.
+            log_initial_pos (bool, optional): Whether to log the initial position of at the beginning of a training
+                or warmstart. Defaults to False.
         """
         self.log_chunk_switch = log_chunk_switch
         self.log_initial_pos = log_initial_pos

src/modalities/dataloader/dataset_factory.py

@@ -116,13 +116,18 @@ def get_packed_mem_map_dataset_megatron(
         return dataset
 
     @staticmethod
-    def get_combined_dataset(datasets: list[Dataset]) -> Dataset:
+    def get_combined_dataset(
+        datasets: list[Dataset], log_chunk_switch: bool = False, log_initial_pos: bool = False
+    ) -> Dataset:
         """Factory method for creating a combined datset .
 
         Args:
             datasets (list[Dataset]): List of datasets to combine.
+            log_chunk_switch (bool, optional): Whether to log when switching between dataset chunks. Defaults to False.
+            log_initial_pos (bool, optional): Whether to log the initial position of at the beginning of a
+                training or warmstart. Defaults to False.
 
         Returns:
             Dataset: CombinedDataset object.
         """
-        return CombinedDataset(datasets=datasets)
+        return CombinedDataset(datasets=datasets, log_chunk_switch=log_chunk_switch, log_initial_pos=log_initial_pos)