created a helper function to convert the existing huggingface dataset to an iterable dataset (#64)

Author: ChaoPang

src/cehrbert/runners/hf_cehrbert_finetune_runner.py

@@ -27,6 +27,7 @@
 from cehrbert.models.hf_models.tokenization_hf_cehrbert import CehrBertTokenizer
 from cehrbert.runners.hf_runner_argument_dataclass import FineTuneModelType, ModelArguments
 from cehrbert.runners.runner_util import (
+    convert_dataset_to_iterable_dataset,
     generate_prepared_ds_path,
     get_last_hf_checkpoint,
     get_meds_extension_path,
@@ -99,7 +100,9 @@ def main():
         LOG.info(f"Loading prepared dataset from disk at {prepared_ds_path}...")
         processed_dataset = load_from_disk(str(prepared_ds_path))
         if data_args.streaming:
-            processed_dataset = processed_dataset.to_iterable_dataset(num_shards=training_args.dataloader_num_workers)
+            processed_dataset = convert_dataset_to_iterable_dataset(
+                processed_dataset, num_shards=training_args.dataloader_num_workers
+            )
         LOG.info("Prepared dataset loaded from disk...")
     else:
         # If the data is in the MEDS format, we need to convert it to the CEHR-BERT format
@@ -112,13 +115,9 @@ def main():
                 LOG.info(f"Trying to load the MEDS extension from disk at {meds_extension_path}...")
                 dataset = load_from_disk(meds_extension_path)
                 if data_args.streaming:
-                    if isinstance(dataset, DatasetDict):
-                        dataset = {
-                            k: v.to_iterable_dataset(num_shards=training_args.dataloader_num_workers)
-                            for k, v in dataset.items()
-                        }
-                    else:
-                        dataset = dataset.to_iterable_dataset(num_shards=training_args.dataloader_num_workers)
+                    dataset = convert_dataset_to_iterable_dataset(
+                        dataset, num_shards=training_args.dataloader_num_workers
+                    )
             except Exception as e:
                 LOG.exception(e)
                 dataset = create_dataset_from_meds_reader(data_args, is_pretraining=False)

src/cehrbert/runners/hf_cehrbert_pretrain_runner.py

@@ -14,6 +14,7 @@
 from cehrbert.models.hf_models.tokenization_hf_cehrbert import CehrBertTokenizer
 from cehrbert.runners.hf_runner_argument_dataclass import DataTrainingArguments, ModelArguments
 from cehrbert.runners.runner_util import (
+    convert_dataset_to_iterable_dataset,
     generate_prepared_ds_path,
     get_last_hf_checkpoint,
     get_meds_extension_path,
@@ -160,7 +161,9 @@ def main():
         LOG.info("Loading prepared dataset from disk at %s...", prepared_ds_path)
         processed_dataset = load_from_disk(str(prepared_ds_path))
         if data_args.streaming:
-            processed_dataset = processed_dataset.to_iterable_dataset(num_shards=training_args.dataloader_num_workers)
+            processed_dataset = convert_dataset_to_iterable_dataset(
+                processed_dataset, num_shards=training_args.dataloader_num_workers
+            )
         LOG.info("Prepared dataset loaded from disk...")
         # If the data has been processed in the past, it's assume the tokenizer has been created
         # before. We load the CEHR-BERT tokenizer from the output folder.
@@ -179,13 +182,9 @@ def main():
                 )
                 dataset = load_from_disk(meds_extension_path)
                 if data_args.streaming:
-                    if isinstance(dataset, DatasetDict):
-                        dataset = {
-                            k: v.to_iterable_dataset(num_shards=training_args.dataloader_num_workers)
-                            for k, v in dataset.items()
-                        }
-                    else:
-                        dataset = dataset.to_iterable_dataset(num_shards=training_args.dataloader_num_workers)
+                    dataset = convert_dataset_to_iterable_dataset(
+                        dataset, num_shards=training_args.dataloader_num_workers
+                    )
             except FileNotFoundError as e:
                 LOG.exception(e)
                 dataset = create_dataset_from_meds_reader(data_args, is_pretraining=True)

src/cehrbert/runners/runner_util.py

@@ -4,10 +4,10 @@
 import re
 import sys
 from pathlib import Path
-from typing import Tuple, Union
+from typing import Dict, Tuple, Union
 
 import torch
-from datasets import Dataset, IterableDataset, load_dataset
+from datasets import Dataset, DatasetDict, IterableDataset, IterableDatasetDict, load_dataset
 from torch.nn import functional as F
 from transformers import EvalPrediction, HfArgumentParser, TrainingArguments
 from transformers.trainer_utils import get_last_checkpoint
@@ -337,3 +337,53 @@ def get_meds_extension_path(data_folder: str, dataset_prepared_path: str):
     basename = os.path.basename(data_folder)
     meds_extension_path = os.path.join(dataset_prepared_path, f"{basename}_meds_extension")
     return meds_extension_path
+
+
+def convert_dataset_to_iterable_dataset(
+    dataset: Union[Dataset, DatasetDict], num_shards: int = 1
+) -> Union[IterableDataset, Dict[str, IterableDataset]]:
+    """
+    Converts a Hugging Face `Dataset` or `DatasetDict` into an `IterableDataset` or.
+
+    a dictionary of `IterableDataset` objects, enabling efficient parallel processing
+    using multiple workers in a data loader.
+
+    Parameters
+    ----------
+    dataset : Union[Dataset, DatasetDict]
+        The input dataset, which can be either:
+        - A single `Dataset` object
+        - A `DatasetDict` (containing multiple datasets, such as train, validation, and test splits)
+
+    num_shards : int
+        The number of workers (shards) to split the dataset into for parallel data loading.
+        This allows efficient sharding of the dataset across multiple workers.
+
+    Returns
+    -------
+    Union[IterableDataset, Dict[str, IterableDataset]]
+        The converted dataset, either as:
+        - A single `IterableDataset` if the input was a `Dataset`
+        - A dictionary of `IterableDataset` objects if the input was a `DatasetDict` or `IterableDatasetDict`
+
+    Notes
+    -----
+    - If the input `dataset` is a `DatasetDict` (or `IterableDatasetDict`), each dataset split
+      (e.g., train, validation, test) is converted into an `IterableDataset`.
+    - If the input `dataset` is a single `Dataset`, it is directly converted into an `IterableDataset`.
+    - The `num_shards` parameter in `to_iterable_dataset` allows splitting the dataset for parallel
+      data loading with multiple workers.
+
+    Example
+    -------
+    # Convert a standard Dataset to an IterableDataset for parallel processing
+    iterable_dataset = convert_dataset_to_iterable_dataset(my_dataset, dataloader_num_workers=4)
+
+    # Convert a DatasetDict (e.g., train, validation splits) into IterableDataset objects
+    iterable_dataset_dict = convert_dataset_to_iterable_dataset(my_dataset_dict, dataloader_num_workers=4)
+    """
+    if isinstance(dataset, DatasetDict) or isinstance(dataset, IterableDatasetDict):
+        dataset = {k: v.to_iterable_dataset(num_shards=num_shards) for k, v in dataset.items()}
+    else:
+        dataset = dataset.to_iterable_dataset(num_shards=num_shards)
+    return dataset