feat: added logging for chunk switches and initial positions in CombinedDataset

Author: le1nux

src/modalities/config/config.py

@@ -448,6 +448,8 @@ class PackedMemMapDatasetMegatronConfig(BaseModel):
 
 class CombinedDatasetConfig(BaseModel):
     datasets: list[PydanticDatasetIFType]
+    log_chunk_switch: bool = False
+    log_initial_pos: bool = False
 
 
 class BatchSamplerConfig(BaseModel):

src/modalities/dataloader/dataset.py

@@ -14,6 +14,7 @@
 from modalities.dataloader.create_packed_data import EmbeddedStreamData
 from modalities.dataloader.large_file_lines_reader import LargeFileLinesReader
 from modalities.tokenization.tokenizer_wrapper import TokenizerWrapper
+from modalities.utils.logger_utils import get_logger
 
 
 class Dataset(TorchdataSet):
@@ -445,20 +446,28 @@ class CombinedDataset(Dataset):
     In the Dataloader, a batch will still contain packed samples from different datasets.
     """
 
-    def __init__(self, datasets: list[Dataset]):
+    def __init__(self, datasets: list[Dataset], log_chunk_switch: bool = False, log_initial_pos: bool = False):
         """Initializes the CombinedDataset object, combining multiple datasets.
 
         Args:
             datasets (list[Dataset]): A list of datasets to combine.
         """
+        self.log_chunk_switch = log_chunk_switch
+        self.log_initial_pos = log_initial_pos
         self.datasets = datasets
         self.cumulative_sizes = np.cumsum([len(ds) for ds in datasets], dtype=np.int64)
+        self.logger = get_logger(__name__)
 
     def __len__(self) -> int:
         return self.cumulative_sizes[-1]
 
     def __getitem__(self, idx: int) -> dict:
         dataset_idx = np.searchsorted(self.cumulative_sizes, idx, side="right")
         local_idx = idx - (self.cumulative_sizes[dataset_idx - 1] if dataset_idx > 0 else 0)
+        if self.log_chunk_switch and local_idx == 0:
+            self.logger.info(f"global_index={idx} chunk index={dataset_idx}, local index={local_idx}")
+
+        if self.log_initial_pos:
+            self.logger.info(f"global_index={idx} chunk index={dataset_idx}, local index={local_idx}")
 
         return self.datasets[dataset_idx][local_idx]