vision/m4/training/dataset.py

"""
This file defines the dataloader logic.
"""
import copy
import inspect
import logging
import os
from dataclasses import asdict
from functools import partial
from pathlib import Path
from typing import Optional

import datasets
import numpy as np
import torch
import webdataset as wds
from accelerate.state import AcceleratorState
from PIL import Image, ImageFile
from torch.utils.data import Sampler

from m4.training.config import DataParams, DatasetParams, Parameters
from m4.training.dataset_utils import check_webdataset_command, get_webdataset
from m4.training.packing import (
    split_pack_and_pad_iqa_finetuning,
    split_pack_and_pad_ocr,
    split_pack_and_pad_pairs,
    split_pack_and_pad_sft,
    split_pack_and_pad_webdocs,
)
from m4.training.types import DatasetNames, DatasetTypes


Image.MAX_IMAGE_PIXELS = None
ImageFile.LOAD_TRUNCATED_IMAGES = True

logger = logging.getLogger(__name__)

"""

Possible dataloader/dataset nestings

DataLoaderForIterableWrapperDataset
  CustomChainDataset
    IterableWrapperDataset

DataLoaderForMappedWrapperDataset
  MapWrapperDataset

"""


# TODO(siddk) :: This file needs to be cleaned up a bit?
"""
dataset[idx]:
    {
        'images':[
            <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=256x256 at 0x7FD503595350>,
            None,
            None,
            None,
            None,
            <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=256x256 at 0x7FD503595D90>,
            <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=256x256 at 0x7FD503595E90>,
            <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=256x256 at 0x7FD503595F90>,
            <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=256x256 at 0x7FD5035930D0>,
            <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=256x256 at 0x7FD503595D50>,
            <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=256x256 at 0x7FD503593290>,
            ],
        'texts':[
            None,
            'VALMAS COMMERCE',
            'VALMAS SA imports top quality products and materials that meet the requirements of modern construction. The goal of our company is to provide material and technical know-how for reliable solutions in the field of construction, renovation and repair. Raw materials of high-quality for every particular need, which also remain environmentally sensitive. Investing continually in enhancing our infrastructure and with a steadily growing track record over the years, we are in a position to meet your every need. We will be glad to welcome you to our 500m² exhibition to choose from a range of quality materials the ones that best match your own requirements.',
            'FLOORING – PAINTS\n\nQuality and great variety in paneling materials for classic and traditional applications of modern residential aesthetics.\n\nFLOORING: With knowledge and experience in construction, we are able to propose solutions that combine the available budget with the specific demands for functionality and aesthetics. Wooden floors, ceramic tiles, marble, stone, as well as construction flooring, such as colored or special cement mortars, offer classic or traditional choices, following the original and modern residential aesthetic result.\n\n·'
            'Partners List',
            None,
            None,
            None,
            None,
            None,
            None,
        ]
    }
"""


def to_tensor(batch):
    for key in batch:
        if not torch.is_tensor(batch[key]):
            batch[key] = torch.tensor(batch[key])
    return batch


# We don't want collate_fn to add extra dim --> can't use a lambda because of `multiprocessing` pickle requirements!
def simple_collate(x):
    return x[0]


def get_mapper(
    tokenizer,
    image_transform,
    dataset_type: DatasetTypes,
    image_seq_len: int,
    max_seq_len: int = 256,
    max_num_images: int = 5,
    max_image_size: int = 384,
    vision_encoder_max_image_size: int = 384,
    pre_split_scale_up_max=1.0,
    pre_split_scale_up_frequency=0.0,
    is_t5: bool = False,
    pad_dataset: bool = True,
    max_num_samples_per_document: int = 1,
    t5_mlm_noise_density: Optional[float] = None,
    t5_mlm_mean_noise_span_length: Optional[int] = None,
    add_begin_of_doc_token: bool = True,
    add_end_of_doc_token: bool = True,
    max_num_images_per_document: Optional[int] = None,
):
    mapper_kwargs = {
        "tokenizer": tokenizer,
        "image_transform": image_transform,
        "max_seq_len": max_seq_len,
        "max_num_images": max_num_images,
        "max_image_size": max_image_size,
        "vision_encoder_max_image_size": vision_encoder_max_image_size,
        "pre_split_scale_up_max": pre_split_scale_up_max,
        "pre_split_scale_up_frequency": pre_split_scale_up_frequency,
        "image_seq_len": image_seq_len,
        "add_begin_of_doc_token": add_begin_of_doc_token,
        "add_end_of_doc_token": add_end_of_doc_token,
    }

    if not pad_dataset:
        raise ValueError("This feature has been deprecated. The dataset must be padded")

    if is_t5:
        mapper_kwargs["noise_density"] = t5_mlm_noise_density
        mapper_kwargs["mean_noise_span_length"] = t5_mlm_mean_noise_span_length
        raise ValueError("This feature has been deprecated. We can't pack for t5")
    elif dataset_type == DatasetTypes.IMAGE_CAPTION_PAIRS:
        split_fn = split_pack_and_pad_pairs
    elif dataset_type == DatasetTypes.OCR:
        split_fn = split_pack_and_pad_ocr
    elif (dataset_type == DatasetTypes.VQAV2_TASK_FINETUNING) or (dataset_type == DatasetTypes.DOCVQA):
        split_fn = split_pack_and_pad_iqa_finetuning
    elif dataset_type == DatasetTypes.SFT:
        split_fn = split_pack_and_pad_sft
    elif dataset_type == DatasetTypes.WEB_DOCUMENTS:
        split_fn = split_pack_and_pad_webdocs
        mapper_kwargs["max_num_samples_per_document"] = max_num_samples_per_document
        mapper_kwargs["max_num_images_per_document"] = max_num_images_per_document

    mapper_with_args = partial(split_fn, **mapper_kwargs)
    return mapper_with_args


def get_dataloaders(
    config: Parameters,
    rank: int,
    world_size: int,
    tokenizer,
    train_image_transforms,
    val_image_transforms,
    image_seq_len: int,
):
    logger.info("Getting the train dataloader")
    train_loader = get_dataloader_from_config(
        tokenizer=tokenizer,
        image_transforms=train_image_transforms,
        seed=config.data_param.train_seed,
        config=config,
        is_train=True,
        rank=rank,
        world_size=world_size,
        image_seq_len=image_seq_len,
    )

    if config.hparams.do_validation:
        logger.info("Getting the validation dataloader")
        val_loader = get_dataloader_from_config(
            tokenizer=tokenizer,
            image_transforms=val_image_transforms,
            seed=config.data_param.val_seed,
            config=config,
            is_train=False,
            rank=rank,
            world_size=world_size,
            image_seq_len=image_seq_len,
        )
    else:
        val_loader = None

    return train_loader, val_loader


def load_hf_dataset(dataset_path):
    split_name = None
    config_name = None
    if ":" in dataset_path:
        dataset_path_splitted = dataset_path.split(":")
        if len(dataset_path_splitted) == 2:
            dataset_path, split_name = dataset_path_splitted
        elif len(dataset_path_splitted) == 3:
            dataset_path, config_name, split_name = dataset_path_splitted

    if os.path.exists(dataset_path):
        # a local path dataset can be of two kinds
        # 1. generated by `save_to_disk` and thus must be loaded with `load_from_disk`
        # 2. hub-like dataset, but which is not online
        # so we try the first and if it fails with `FileNotFoundError` (despite the path existing) we try the second
        try:
            hf_dataset = datasets.load_from_disk(dataset_path)
        except FileNotFoundError:
            if config_name is not None:
                hf_dataset = datasets.load_dataset(dataset_path, name=config_name)
            else:
                hf_dataset = datasets.load_dataset(dataset_path)
    else:
        if config_name is not None:
            hf_dataset = datasets.load_dataset(
                dataset_path, name=config_name, use_auth_token=os.environ.get("HF_TOKEN", True)
            )
        else:
            hf_dataset = datasets.load_dataset(dataset_path, use_auth_token=os.environ.get("HF_TOKEN", True))

    if split_name is not None:
        hf_dataset = hf_dataset[split_name]

    return hf_dataset


def get_dataset_hf(
    dataset_config: DatasetParams,
    tokenizer,
    image_transform,
    is_train: bool = True,
    realtime_processing: bool = True,
    is_t5: bool = False,
):
    dataset_list = []
    hf_datasets_paths = (
        dataset_config.training_datasets_paths if is_train else dataset_config.validation_datasets_paths
    )
    # hf_datasets_paths can be a list of paths, or a .txt file path that contains the paths
    if len(hf_datasets_paths) == 1 and str(hf_datasets_paths[0]).endswith(".txt"):
        with open(hf_datasets_paths[0], "r") as file_shards:
            hf_datasets_paths = [path for path in file_shards.read().split("\n") if path]

    for dataset_path in hf_datasets_paths:
        hf_dataset = load_hf_dataset(dataset_path=str(dataset_path))

        is_paired_dataset = "meta" in hf_dataset[0] and "source" in hf_dataset[0]

        optional_kwargs_defaults = [
            ("pad_dataset", True),
            ("max_num_samples_per_document", 1),
            ("t5_mlm_noise_density", 0.15),
            ("t5_mlm_mean_noise_span_length", 3),
            ("add_begin_of_doc_token", True),
            ("add_end_of_doc_token", True),
            ("max_num_images_per_document", None),
        ]
        optional_kwargs = {}
        for key, default in optional_kwargs_defaults:
            optional_kwargs[key] = getattr(dataset_config, key, default)

        if not realtime_processing:
            mapper_with_args = get_mapper(
                tokenizer=tokenizer,
                image_transform=image_transform,
                image_seq_len=dataset_config.image_seq_len,
                max_seq_len=dataset_config.max_seq_len,
                max_num_images=dataset_config.max_num_images,
                max_image_size=dataset_config.max_image_size,
                vision_encoder_max_image_size=dataset_config.vision_encoder_max_image_size,
                pre_split_scale_up_max=dataset_config.pre_split_scale_up_max,
                pre_split_scale_up_frequency=dataset_config.pre_split_scale_up_frequency,
                dataset_type=DatasetTypes.IMAGE_CAPTION_PAIRS if is_paired_dataset else DatasetTypes.WEB_DOCUMENTS,
                is_t5=is_t5,
                **optional_kwargs,
            )
            hf_dataset = hf_dataset.map(
                mapper_with_args,
                batched=True,
                batch_size=dataset_config.map_batch_size,
                remove_columns=hf_dataset.column_names,
                num_proc=dataset_config.map_num_proc,
            )
        dataset_list.append(hf_dataset)
    return dataset_list


def get_dataset_webdataset(
    dataset_config: DatasetParams,
    is_train: bool = True,
    realtime_processing: bool = True,
):
    if not realtime_processing:
        raise NotImplementedError("WebDataset is only supported for realtime processing")

    webdataset_paths = dataset_config.training_datasets_paths if is_train else dataset_config.validation_datasets_paths

    if len(webdataset_paths) == 0:
        return None

    # webdataset_paths can be a list of paths/commands, or a .txt file path that contains the paths
    if len(webdataset_paths) == 1 and str(webdataset_paths[0]).endswith(".txt"):
        with open(webdataset_paths[0], "r") as file_shards:
            webdataset_paths = [path for path in file_shards.read().split("\n") if path]
    else:
        raise ValueError("WebDataset only supports a .txt file with the paths or the commands.")

    # Check if the paths/commands are valid
    checks = all([check_webdataset_command(path) for path in webdataset_paths])
    if not checks:
        raise ValueError("WebDataset paths/commands are not valid. Please check the paths/commands.")

    combined_dataset = get_webdataset(
        urls=webdataset_paths,
        ds_type=dataset_config.dataset_type,
        batch_size=dataset_config.map_batch_size,
        shuffle_initial_urls_list=dataset_config.shuffle_initial_urls_list if is_train else False,
        shuffle_before_split_by_node_buffer_size=(
            dataset_config.shuffle_before_split_by_node_buffer_size if is_train else None
        ),
        shuffle_before_split_by_worker_buffer_size=(
            dataset_config.shuffle_before_split_by_worker_buffer_size if is_train else None
        ),
        shuffle_after_tarfile_to_samples_buffer_size=(
            dataset_config.shuffle_after_tarfile_to_samples_buffer_size if is_train else None
        ),
        shuffle_after_batching_buffer_size=dataset_config.shuffle_after_batching_buffer_size if is_train else None,
    )
    return combined_dataset


def get_dataset(
    dataset_config: DatasetParams,
    tokenizer=None,
    image_transform=None,
    is_train: bool = True,
    realtime_processing: bool = True,
    is_t5: bool = False,
    use_webdataset: bool = False,
):
    if use_webdataset:
        return get_dataset_webdataset(
            dataset_config=dataset_config,
            is_train=is_train,
            realtime_processing=realtime_processing,
        )
    else:
        return get_dataset_hf(
            dataset_config=dataset_config,
            tokenizer=tokenizer,
            image_transform=image_transform,
            is_train=is_train,
            realtime_processing=realtime_processing,
            is_t5=is_t5,
        )


def get_dataloader(
    tokenizer,
    image_transforms,
    seed,
    num_workers=1,
    pin_memory=False,
    batch_size=10,
    is_train=True,
    persistent_workers=True,
    realtime_processing=False,
    # The following arguments only used for iterable dataset
    rank=None,
    world_size=None,
    # This argument is for controlling sample order randomness for Map-style Datasets when resuming a run
    sampler_rng=None,
    accumulate_datasets=False,
    model_name=None,
    data_param: Optional[DataParams] = None,
    image_seq_len=None,
):
    if is_train:
        select_n_examples = data_param.select_n_examples_train
    else:
        select_n_examples = data_param.select_n_examples_validation

    if data_param is None:
        raise ValueError("data_param must be provided")

    is_t5 = "t5" in model_name if model_name is not None else False
    dataset_list_map = {}

    # Try all possible datasets, if they don't have datasets paths in the config, they
    # will end up with an empty list
    for dataset_name in DatasetNames:
        curr_dataset_config = getattr(data_param, dataset_name.name.lower())
        dataset_list_map[dataset_name.name.lower()] = get_dataset(
            dataset_config=curr_dataset_config,
            tokenizer=tokenizer,
            image_transform=image_transforms[dataset_name.name.lower()],
            is_train=is_train,
            realtime_processing=realtime_processing,
            is_t5=is_t5,
            use_webdataset=data_param.use_webdataset,
        )

    if not realtime_processing:
        # => Important & gnarly: Set image transform based on a novel instance of `np.default_rng` seeded by the parent
        #   seed and the current rank; because of the way DataLoader `num_workers` multiprocessing works in tandem with
        #   the default `hf.dataset` Arrow backend + normal PyTorch `DataLoader, Sampler, and BatchSampler` behavior,
        #   any "global" reference to an rng object that gets to this transform will get "pickled and copied over" to
        #   each separate worker process.
        #
        #   This wouldn't be a problem if we could simply just "reset" the randomness of the Dataset, but that's opaque
        #   given the `hf.Dataset` wrapper; as such, we need to just handle the randomness ourselves by advancing the
        #   random state the appropriate amount in the `__getitem__` of the MapWrapperDataset (as that's only other
        #   place where we're sure we're in scope for a given worker's process block).
        full_dataset = datasets.concatenate_datasets(
            [dataset for dataset_list in dataset_list_map.values() for dataset in dataset_list]
        )
        if select_n_examples is not None:
            full_dataset = full_dataset.select(range(select_n_examples))
        transform_rng = np.random.default_rng(seed=[seed, rank])

        # Wrap `full_dataset` in custom MapWrapperDataset, and initialize a ResumableSampler
        full_dataset = MapWrapperDataset(full_dataset, transform_rng)
        resume_sampler = ResumableSampler(full_dataset, sampler_rng)

        return DataLoaderForMappedWrapperDataset(
            full_dataset,
            batch_size=batch_size,
            sampler=resume_sampler,
            num_workers=num_workers,
            pin_memory=pin_memory,
            drop_last=True,
            persistent_workers=persistent_workers,
        )
    else:
        realtime_processing_datasets = []
        for dataset_name in DatasetNames:
            dataset_list_or_combined = dataset_list_map.get(dataset_name.name.lower(), [])

            if dataset_list_or_combined is None or (
                isinstance(dataset_list_or_combined, list) and len(dataset_list_or_combined) == 0
            ):
                continue

            if isinstance(dataset_list_or_combined, list):
                # If we have a list of datasets, we know those are hf datasets, so we can concatenate them
                combined_dataset = datasets.concatenate_datasets(dataset_list_or_combined)
                if len(combined_dataset) // max(num_workers, 1) < batch_size:
                    raise ValueError(
                        f"For real-time processing, len(dataset) [={len(combined_dataset)}] // num_workers"
                        f" [={num_workers}] must be >= batch_size [={batch_size}]!"
                    )
                if select_n_examples is not None:
                    combined_dataset = combined_dataset.select(range(select_n_examples))
                wrapper_dataset_class = IterableWrapperHFDataset
            elif isinstance(dataset_list_or_combined, wds.pipeline.DataPipeline):
                combined_dataset = dataset_list_or_combined
                wrapper_dataset_class = IterableWrapperWebdataset
            else:
                raise ValueError("Type unrecognized")

            dataset_config: DatasetParams = getattr(data_param, dataset_name.name.lower())
            dataset_kwargs = asdict(dataset_config)
            signature = inspect.signature(wrapper_dataset_class.__init__)
            dataset_kwargs = {k: v for k, v in dataset_kwargs.items() if k in signature.parameters}
            iterable_dataset_instance = wrapper_dataset_class(
                combined_dataset,
                tokenizer=tokenizer,
                image_transform=image_transforms[dataset_name.name.lower()],
                batch_size=batch_size,
                seed=seed,
                shuffle=is_train,
                rank=rank,
                world_size=world_size,
                drop_last=True,
                is_t5=is_t5,
                image_seq_len=image_seq_len,
                **dataset_kwargs,
            )
            realtime_processing_datasets.append(iterable_dataset_instance)

        full_dataset = CustomChainDataset(
            realtime_processing_datasets,
            num_workers,
            rank,
            accumulate_datasets=accumulate_datasets,
            proba_interleaving_dataset=data_param.proba_interleaving_dataset,
            is_train=is_train,
        )

        return DataLoaderForIterableWrapperDataset(
            full_dataset,
            seed=seed,
            num_workers=num_workers,
            pin_memory=pin_memory,
            persistent_workers=persistent_workers,
            collate_fn=simple_collate,
            drop_last=True,
            rank=rank,
            world_size=world_size,
        )


def get_dataloader_from_config(
    tokenizer,
    image_transforms,
    image_seq_len,
    seed,
    config,
    is_train,
    rank=None,
    world_size=None,
    sampler_rng=None,
):
    dataloader = get_dataloader(
        tokenizer=tokenizer,
        image_transforms=image_transforms,
        seed=seed,
        num_workers=config.data_param.num_workers,
        pin_memory=config.data_param.pin_memory,
        batch_size=config.data_param.batch_size,
        is_train=is_train,
        persistent_workers=config.data_param.persistent_workers,
        realtime_processing=config.data_param.realtime_processing,
        rank=rank,
        world_size=world_size,
        sampler_rng=sampler_rng,
        accumulate_datasets=True if config.hparams.loss_weights_per_dataset is not None else False,
        model_name=config.hparams.model_name,
        data_param=config.data_param,
        image_seq_len=image_seq_len,
    )
    return dataloader


# Creates a Resumable Sampler for Map-style Datasets that predefines the set of indices to retrieve for
# the given epoch (seeded with a generator!). We are using this instead of just seeding the "default Sampler"
# so we can quickly "short-circuit" and bypass any and all "seen" indices (see `__iter__`).
class ResumableSampler(Sampler):
    def __init__(self, data_source, sample_generator):
        super().__init__(data_source)
        self.data_source, self.indices = data_source, None

        # Note: `accelerate` hardcodes a search for an attribute named `generator`, which it then uses to synchronize
        # generators across all ranks on each call to `__iter__`; in our case, this is really bad as we want full
        # control over "sample" randomness (only for Map-style datasets). To get around this, it suffices to just
        # name the instance attributes anything other than `generator`... so that's what we do!
        self.sample_generator = sample_generator

        # For "instant" resuming
        self.n_seen_examples = 0

    def set_state(self, n_seen_examples_per_worker):
        self.n_seen_examples = sum(n_seen_examples_per_worker.values())

    def get_index_order(self):
        return torch.randperm(len(self.data_source), generator=self.sample_generator).tolist()

    def __iter__(self):
        self.indices = self.get_index_order()

        # Resume logic -> advance by `n_seen_examples`
        self.indices = self.indices[self.n_seen_examples :]

        yield from self.indices

    def __len__(self):
        return len(self.data_source)


# Super simple wrapper around dataset; mostly for compatibility with IterableWrapperDataset
class MapWrapperDataset(torch.utils.data.Dataset):
    def __init__(self, dataset, transform_rng):
        self.wrapped_dataset, self.transform_rng = dataset, transform_rng

        # For "instant" resuming
        self.n_seen_examples_per_worker = {}

    def state_dict(self):
        state_dict = {}

        # recurse into its dataset
        state_dict["wrapped_dataset"] = self.wrapped_dataset.state_dict()

        return state_dict

    def load_state_dict(self, state_dict):
        # recurse into its dataset
        self.wrapped_dataset.load_state_dict(state_dict["wrapped_dataset"])

    @property
    def dataset(self):
        return self.wrapped_dataset

    def set_state(self, n_seen_examples_per_worker):
        self.n_seen_examples_per_worker = n_seen_examples_per_worker

    def __getitem__(self, idx):
        # Dummy dataset idx used for compatibility with CustomChainDataset
        dummy_dataset_idx = 0
        worker_info = torch.utils.data.get_worker_info()
        n_seen_examples = self.n_seen_examples_per_worker.get(worker_info.id, 0)

        # If `n_seen_examples` is non-zero --> advance random state "quickly" then return new example!
        if n_seen_examples > 0:
            for _ in range(n_seen_examples):
                self.transform_rng.random()
            self.n_seen_examples_per_worker[worker_info.id] = 0

        return dummy_dataset_idx, worker_info.id, self.wrapped_dataset[idx]

    def __len__(self):
        return len(self.wrapped_dataset)


class IterableWrapperHFDataset(torch.utils.data.IterableDataset):
    def __init__(
        self,
        dataset,
        tokenizer,
        image_transform,
        batch_size,
        seed,
        dataset_type,
        dataset_name,
        image_seq_len,
        shuffle=True,
        rank=None,
        world_size=None,
        drop_last=True,
        is_t5=False,
        # Dataset specific params
        max_num_images=5,
        max_seq_len=256,
        max_image_size=384,
        vision_encoder_max_image_size=384,
        pre_split_scale_up_max=1.0,
        pre_split_scale_up_frequency=0.0,
        pad_dataset=True,
        mapper_batch_size=128,
        # Setting default to 0 as PMD doesn't need it and it is set to 0.5 in CM4 config by default
        max_num_samples_per_document=1,
        t5_mlm_noise_density=None,
        t5_mlm_mean_noise_span_length=None,
        add_begin_of_doc_token=True,
        add_end_of_doc_token=True,
        shuffle_after_packing=False,
        max_num_images_per_document=None,
    ):
        self.dataset = dataset
        self.mapper = get_mapper(
            tokenizer=tokenizer,
            image_transform=image_transform,
            image_seq_len=image_seq_len,
            max_seq_len=max_seq_len,
            max_num_images=max_num_images,
            max_image_size=max_image_size,
            vision_encoder_max_image_size=vision_encoder_max_image_size,
            pre_split_scale_up_max=pre_split_scale_up_max,
            pre_split_scale_up_frequency=pre_split_scale_up_frequency,
            pad_dataset=pad_dataset,
            max_num_samples_per_document=max_num_samples_per_document,
            dataset_type=dataset_type,
            is_t5=is_t5,
            t5_mlm_noise_density=t5_mlm_noise_density,
            t5_mlm_mean_noise_span_length=t5_mlm_mean_noise_span_length,
            add_begin_of_doc_token=add_begin_of_doc_token,
            add_end_of_doc_token=add_end_of_doc_token,
            max_num_images_per_document=max_num_images_per_document,
        )
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.seed = seed
        self.epoch = 0
        self.rank = rank
        self.world_size = world_size
        self.mapper_batch_size = mapper_batch_size
        self.drop_last = drop_last

        self.shuffle_after_packing = shuffle_after_packing

        # To be initialized later in __iter__
        self.rng = None

        # Resume Tracking --> Dict[worker_idx] -> Tuple[map_idx, key_idx]; `map_idx` lets us jumpstart!
        self.worker_idx_tracker = {}
        self.start_worker_idx = 0

        self.dataset_name = dataset_name

    def set_state(self, worker_idx_tracker, start_worker_idx):
        self.worker_idx_tracker = worker_idx_tracker
        self.start_worker_idx = start_worker_idx

    def set_epoch(self, epoch):
        self.epoch = epoch

    def state_dict(self):
        state_dict = {}
        return state_dict

    def load_state_dict(self, state_dict):
        pass

    def _get_worker_id_and_worker_total_num(self):
        worker_info = torch.utils.data.get_worker_info()
        if worker_info is None:
            worker_total_num = 1
            worker_id = 0
        else:
            worker_total_num = worker_info.num_workers
            worker_id = worker_info.id

        worker_id = (worker_id + self.start_worker_idx) % worker_total_num
        return worker_id, worker_total_num

    def _get_worker_indices(self):
        sampler = torch.utils.data.DistributedSampler(
            self.dataset,
            num_replicas=self.world_size,
            rank=self.rank,
            shuffle=self.shuffle,
            seed=self.seed,
            drop_last=self.drop_last,
        )
        sampler.set_epoch(self.epoch)

        # Get list of indices for current rank from the distributed sampler
        indices = list(iter(sampler))

        worker_id, worker_total_num = self._get_worker_id_and_worker_total_num()
        # Take the subset of indices that belong to this worker
        #   => It will look something like 0, 2, 4 when worker_id is 0 and num_workers is 2
        worker_indices = indices[worker_id::worker_total_num]
        return worker_indices, worker_id

    def __iter__(self):
        # Dummy dataset idx used for compatibility with CustomChainDataset
        dummy_dataset_idx = 0

        if self.rank is None or self.world_size is None:
            raise ValueError("rank and world_size must be provided")

        # Get worker indices & details for resuming...
        worker_indices, worker_id = self._get_worker_indices()
        num_worker_indices = len(worker_indices)

        # Set start idx of loop based on `self.worker_resume_idxs`
        map_start_idx, last_key_idx, overflow_batch = self.worker_idx_tracker.get(worker_id, (0, -1, {}))
        self.rng_seed = [self.seed, self.epoch, self.rank, worker_id, map_start_idx]
        self.rng = np.random.default_rng(seed=self.rng_seed)
        for i in range(map_start_idx, num_worker_indices, self.mapper_batch_size):
            # Set seed for the worker according to worker index and the index and then reset it work
            # This needs to be done so that torch random crop is deterministic
            rng_state = torch.get_rng_state()
            torch.manual_seed(f"{self.seed}{worker_id}{i}")
            # Feed `worker_indices[i]` to mapper to ensure "deterministic randomness" that we don't have to track...
            curr_mapped_batch = self.mapper(
                self.dataset[worker_indices[i : i + self.mapper_batch_size]],
                prefix_seed=(self.seed, self.epoch, self.rank, worker_id, i),
            )
            torch.set_rng_state(rng_state)
            keys = list(curr_mapped_batch.keys())
            overflow_batch_keys = overflow_batch.keys()

            # Check if overflow from previous batches is left, if yes, add it to the current batch
            # Specifically, we should prepend this overflow batch so as it goes out first and
            # current batch possibly becomes next overflow batch
            if len(overflow_batch_keys) > 0:
                if sorted(overflow_batch_keys) != sorted(keys):
                    raise ValueError(
                        "Overflow batch keys not equal to current keys. Make sure mapper is always returning"
                        "  dictionary with the same keys. "
                        f"Overflow: {sorted(overflow_batch_keys)}, Mapping: {sorted(keys)}"
                    )
                else:
                    mapped_batch = {}

                    if "pixel_values" in overflow_batch or "pixel_values" in curr_mapped_batch:
                        total_batch_size = overflow_batch["input_ids"].size(0) + curr_mapped_batch["input_ids"].size(0)
                        max_num_images = max(
                            overflow_batch["pixel_values"].size(1) if "pixel_values" in overflow_batch else 0,
                            curr_mapped_batch["pixel_values"].size(1) if "pixel_values" in curr_mapped_batch else 0,
                        )
                        max_height = max(
                            overflow_batch["pixel_values"].size(3) if "pixel_values" in overflow_batch else 0,
                            curr_mapped_batch["pixel_values"].size(3) if "pixel_values" in curr_mapped_batch else 0,
                        )
                        max_width = max(
                            overflow_batch["pixel_values"].size(4) if "pixel_values" in overflow_batch else 0,
                            curr_mapped_batch["pixel_values"].size(4) if "pixel_values" in curr_mapped_batch else 0,
                        )
                        padded_image_tensor = torch.zeros(total_batch_size, max_num_images, 3, max_height, max_width)
                        padded_pixel_attention_masks = torch.zeros(
                            total_batch_size, max_num_images, max_height, max_width, dtype=torch.bool
                        )

                        start = 0
                        for batch in [overflow_batch, curr_mapped_batch]:
                            if "pixel_values" not in batch:
                                continue
                            px = batch["pixel_values"]
                            px_attn_mask = batch["pixel_attention_mask"]
                            end = start + px.size(0)
                            padded_image_tensor[start:end, :, :, : px.size(3), : px.size(4)] = px
                            padded_pixel_attention_masks[start:end, :, : px.size(3), : px.size(4)] = px_attn_mask
                            start += px.size(0)

                        mapped_batch["pixel_values"] = padded_image_tensor.contiguous()
                        mapped_batch["pixel_attention_mask"] = padded_pixel_attention_masks.contiguous()

                    for key in keys:
                        if key in ["pixel_values", "pixel_attention_mask"]:
                            continue
                        mapped_batch[key] = torch.cat([overflow_batch[key], curr_mapped_batch[key]], dim=0)
                    previous_overflow_batch = copy.deepcopy(overflow_batch)
                    overflow_batch = {}
            else:
                previous_overflow_batch = {}
                mapped_batch = curr_mapped_batch

            first_key = keys[0]
            mapped_batch_length = len(mapped_batch[first_key])

            if self.shuffle_after_packing:
                indices = list(range(mapped_batch_length))
                self.rng.shuffle(indices)

                for key in mapped_batch.keys():
                    mapped_batch[key] = mapped_batch[key][indices, ...]

            if mapped_batch_length < self.batch_size:
                # We need to add more data to this batch to make it of size `self.batch_size`
                # Just setting mapped_batch to overflow_batch should be enough as the next iteration
                # will add more data to it
                overflow_batch = mapped_batch
            else:
                # Now, yield batches of size batch_size from the mapped batch
                for key_idx in range(0, mapped_batch_length, self.batch_size):
                    # Set "reproducible" randomness
                    self.rng_seed = [self.seed, self.epoch, self.rank, worker_id, i, key_idx]
                    self.rng = np.random.default_rng(seed=self.rng_seed)

                    if i == map_start_idx and key_idx <= last_key_idx:
                        # Handle Resume (only for "first" loop iteration) advance random state until `last_key_idx`
                        self.rng.random()
                    else:
                        overflow_batch = {key: mapped_batch[key][key_idx : key_idx + self.batch_size] for key in keys}
                        if len(overflow_batch[first_key]) != self.batch_size:
                            # Last batch
                            break
                        else:
                            dataset_state = {
                                "worker_idx": worker_id,
                                "map_start_idx": i,
                                "last_key_idx": key_idx,
                                "previous_overflow_batch": previous_overflow_batch,
                            }
                            yield dummy_dataset_idx, self.dataset_name.name.lower(), dataset_state, overflow_batch
                            overflow_batch = {}


class IterableWrapperWebdataset(torch.utils.data.IterableDataset):
    def __init__(
        self,
        dataset,
        tokenizer,
        image_transform,
        batch_size,
        seed,
        dataset_type,
        dataset_name,
        image_seq_len,
        shuffle=True,
        rank=None,
        world_size=None,
        drop_last=True,
        is_t5=False,
        # Dataset specific params
        max_num_images=5,
        max_seq_len=256,
        max_image_size=384,
        vision_encoder_max_image_size=384,
        pre_split_scale_up_max=1.0,
        pre_split_scale_up_frequency=0.0,
        pad_dataset=True,
        mapper_batch_size=128,
        # Setting default to 0 as PMD doesn't need it and it is set to 0.5 in CM4 config by default
        max_num_samples_per_document=1,
        t5_mlm_noise_density=None,
        t5_mlm_mean_noise_span_length=None,
        add_begin_of_doc_token=True,
        add_end_of_doc_token=True,
        shuffle_after_packing=False,
        max_num_images_per_document=None,
    ):
        self._webdataset = dataset
        self.dataset = iter(self._webdataset)
        self.mapper = get_mapper(
            tokenizer=tokenizer,
            image_transform=image_transform,
            image_seq_len=image_seq_len,
            max_seq_len=max_seq_len,
            max_num_images=max_num_images,
            max_image_size=max_image_size,
            vision_encoder_max_image_size=vision_encoder_max_image_size,
            pre_split_scale_up_max=pre_split_scale_up_max,
            pre_split_scale_up_frequency=pre_split_scale_up_frequency,
            pad_dataset=pad_dataset,
            max_num_samples_per_document=max_num_samples_per_document,
            dataset_type=dataset_type,
            is_t5=is_t5,
            t5_mlm_noise_density=t5_mlm_noise_density,
            t5_mlm_mean_noise_span_length=t5_mlm_mean_noise_span_length,
            add_begin_of_doc_token=add_begin_of_doc_token,
            add_end_of_doc_token=add_end_of_doc_token,
            max_num_images_per_document=max_num_images_per_document,
        )
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.seed = seed
        self.epoch = 0
        self.rank = rank
        self.world_size = world_size
        self.mapper_batch_size = mapper_batch_size
        self.drop_last = drop_last
        self.shuffle_after_packing = shuffle_after_packing

        # To be initialized later in __iter__
        self.rng = None

        # Resume Tracking --> Dict[worker_idx] -> Tuple[map_idx, key_idx]; `map_idx` lets us jumpstart!
        self.worker_idx_tracker = {}
        self.start_worker_idx = 0

        self.dataset_name = dataset_name

    def set_state(self, worker_idx_tracker, start_worker_idx):
        self.worker_idx_tracker = worker_idx_tracker
        self.start_worker_idx = start_worker_idx

    def reset_state(self):
        for key in self.worker_idx_tracker.keys():
            self.worker_idx_tracker[key] = (0, -1, {})
        self.start_worker_idx = 0

    def set_epoch(self, epoch):
        self.epoch = epoch
        self.dataset = iter(self._webdataset)  # Reset dataset iterator

    def state_dict(self):
        state_dict = {}
        return state_dict

    def load_state_dict(self, state_dict):
        pass

    def _get_worker_id_and_worker_total_num(self):
        worker_info = torch.utils.data.get_worker_info()
        if worker_info is None:
            worker_total_num = 1
            worker_id = 0
        else:
            worker_total_num = worker_info.num_workers
            worker_id = worker_info.id

        worker_id = (worker_id + self.start_worker_idx) % worker_total_num
        return worker_id, worker_total_num

    def _get_worker_indices(self):
        sampler = torch.utils.data.DistributedSampler(
            self.dataset,
            num_replicas=self.world_size,
            rank=self.rank,
            shuffle=self.shuffle,
            seed=self.seed,
            drop_last=self.drop_last,
        )
        sampler.set_epoch(self.epoch)

        # Get list of indices for current rank from the distributed sampler
        indices = list(iter(sampler))

        worker_id, worker_total_num = self._get_worker_id_and_worker_total_num()
        # Take the subset of indices that belong to this worker
        #   => It will look something like 0, 2, 4 when worker_id is 0 and num_workers is 2
        worker_indices = indices[worker_id::worker_total_num]
        return worker_indices, worker_id

    def __iter__(self):
        # Dummy dataset idx used for compatibility with CustomChainDataset
        dummy_dataset_idx = 0

        if self.rank is None or self.world_size is None:
            raise ValueError("rank and world_size must be provided")

        # Relic from previous implementation - but needed for rng seed
        worker_id, worker_total_num = self._get_worker_id_and_worker_total_num()

        # Relic from previous implementation - but needed for rng seed
        map_start_idx, last_key_idx, overflow_batch = self.worker_idx_tracker.get(worker_id, (0, -1, {}))

        # Relic from previous implementation - but needed for rng seed
        i = map_start_idx

        # Initialize rng_seed
        self.rng_seed = [self.seed, self.epoch, self.rank, worker_id, i]
        self.rng = np.random.default_rng(seed=self.rng_seed)
        while True:
            # Set seed for the worker according to worker index and the index and then reset it work
            # This needs to be done so that torch random crop is deterministic
            rng_state = torch.get_rng_state()
            torch.manual_seed(f"{self.seed}{worker_id}{i}")
            try:
                next_batch = next(self.dataset)
                i += 1
            except StopIteration:
                logger.info(
                    f"{self.dataset_name.name.lower()} has finished one epoch and is moving on to the next one."
                    f" (epoch={self.epoch} - rank={self.rank} - worker_id={worker_id})"
                )
                break
            curr_mapped_batch = self.mapper(
                next_batch,
                prefix_seed=(self.seed, self.epoch, self.rank, worker_id, i),
            )
            torch.set_rng_state(rng_state)
            keys = list(curr_mapped_batch.keys())
            overflow_batch_keys = overflow_batch.keys()

            # Check if overflow from previous batches is left, if yes, add it to the current batch
            # Specifically, we should prepend this overflow batch so as it goes out first and
            # current batch possibly becomes next overflow batch
            if len(overflow_batch_keys) > 0:
                if sorted(overflow_batch_keys) != sorted(keys):
                    raise ValueError(
                        "Overflow batch keys not equal to current keys. Make sure mapper is always returning"
                        "  dictionary with the same keys. "
                        f"Overflow: {sorted(overflow_batch_keys)}, Mapping: {sorted(keys)}"
                    )
                else:
                    mapped_batch = {}

                    if "pixel_values" in overflow_batch or "pixel_values" in curr_mapped_batch:
                        total_batch_size = overflow_batch["input_ids"].size(0) + curr_mapped_batch["input_ids"].size(0)
                        max_num_images = max(
                            overflow_batch["pixel_values"].size(1) if "pixel_values" in overflow_batch else 0,
                            curr_mapped_batch["pixel_values"].size(1) if "pixel_values" in curr_mapped_batch else 0,
                        )
                        max_height = max(
                            overflow_batch["pixel_values"].size(3) if "pixel_values" in overflow_batch else 0,
                            curr_mapped_batch["pixel_values"].size(3) if "pixel_values" in curr_mapped_batch else 0,
                        )
                        max_width = max(
                            overflow_batch["pixel_values"].size(4) if "pixel_values" in overflow_batch else 0,
                            curr_mapped_batch["pixel_values"].size(4) if "pixel_values" in curr_mapped_batch else 0,
                        )
                        padded_image_tensor = torch.zeros(total_batch_size, max_num_images, 3, max_height, max_width)
                        padded_pixel_attention_masks = torch.zeros(
                            total_batch_size, max_num_images, max_height, max_width, dtype=torch.bool
                        )

                        start = 0
                        for batch in [overflow_batch, curr_mapped_batch]:
                            if "pixel_values" not in batch:
                                continue
                            px = batch["pixel_values"]
                            px_attn_mask = batch["pixel_attention_mask"]
                            end = start + px.size(0)
                            padded_image_tensor[start:end, :, :, : px.size(3), : px.size(4)] = px
                            padded_pixel_attention_masks[start:end, :, : px.size(3), : px.size(4)] = px_attn_mask
                            start += px.size(0)

                        mapped_batch["pixel_values"] = padded_image_tensor.contiguous()
                        mapped_batch["pixel_attention_mask"] = padded_pixel_attention_masks.contiguous()

                    for key in keys:
                        if key in ["pixel_values", "pixel_attention_mask"]:
                            continue
                        mapped_batch[key] = torch.cat([overflow_batch[key], curr_mapped_batch[key]], dim=0)
                    overflow_batch = {}
            else:
                mapped_batch = curr_mapped_batch

            first_key = keys[0]
            mapped_batch_length = len(mapped_batch[first_key])

            if self.shuffle_after_packing:
                indices = list(range(mapped_batch_length))
                self.rng.shuffle(indices)

                for key in mapped_batch.keys():
                    mapped_batch[key] = mapped_batch[key][indices, ...]

            if mapped_batch_length < self.batch_size:
                # We need to add more data to this batch to make it of size `self.batch_size`
                # Just setting mapped_batch to overflow_batch should be enough as the next iteration
                # will add more data to it
                overflow_batch = mapped_batch
            else:
                # Now, yield batches of size batch_size from the mapped batch
                for key_idx in range(0, mapped_batch_length, self.batch_size):
                    # Set "reproducible" randomness
                    self.rng_seed = [self.seed, self.epoch, self.rank, worker_id, i, key_idx]
                    self.rng = np.random.default_rng(seed=self.rng_seed)

                    overflow_batch = {key: mapped_batch[key][key_idx : key_idx + self.batch_size] for key in keys}
                    if len(overflow_batch[first_key]) != self.batch_size:
                        # Last batch
                        break
                    else:
                        dataset_state = {
                            "worker_idx": worker_id,
                            "map_start_idx": i,
                            "last_key_idx": key_idx,
                            "previous_overflow_batch": {},
                        }
                        yield dummy_dataset_idx, self.dataset_name.name.lower(), dataset_state, overflow_batch
                        overflow_batch = {}


class CustomChainDataset(torch.utils.data.IterableDataset):
    r"""Dataset for chaining multiple :class:`IterableDataset` s.

    This class is useful to assemble different existing dataset streams. The
    chaining operation is done on-the-fly, so concatenating large-scale
    datasets with this class will be efficient.

    For validation set:
        This class iterates over each dataset one by one.
        One dataset is iterated over completely before moving to the next one.

    Args:
        datasets (iterable of IterableDataset): datasets to be chained together
        num_workers (int): number of workers to use for loading data
        rank (int): rank of the current process
        accumulate_datasets (bool): whether to accumulate the datasets or not
        proba_interleaving_dataset (list of float): probability of interleaving each dataset, first probability
            is of the PMD, second one of CM4
        is_train (bool): whether the dataset is used for training or not. Setting this to false will ignore
            accumulate_datasets and proba_interleaving_dataset parameters and set them to False and None.
            See commentary above for validation set to understand the behavior of this class when
            `is_train` is False.
    """

    def __init__(
        self, datasets, num_workers, rank, accumulate_datasets=False, proba_interleaving_dataset=None, is_train=True
    ):
        super(CustomChainDataset, self).__init__()

        for d in datasets:
            if not isinstance(d, torch.utils.data.IterableDataset):
                raise ValueError(f"CustomChainDataset only supports IterableDataset, but got {type(d)}")

        self.datasets = datasets
        self.num_workers = num_workers if num_workers > 1 else 1
        self.num_datasets = len(self.datasets)
        self.is_train = is_train

        if self.is_train is False:
            if accumulate_datasets is True or proba_interleaving_dataset is not None:
                logger.warn("accumulate_datasets and proba_interleaving_dataset are ignored when is_train is False")
            self.accumulate_datasets = False
            self.dataset_proba = None
        else:
            self.accumulate_datasets = accumulate_datasets
            if not self.accumulate_datasets:
                if proba_interleaving_dataset is not None:
                    self.dataset_proba = np.asarray(proba_interleaving_dataset)
                else:
                    self.dataset_proba = np.full((self.num_datasets), 1 / self.num_datasets, dtype=float)

                if abs(self.dataset_proba.sum() - 1) > 0.001:
                    # Allow a tolerance for floating points rounding errors.
                    raise ValueError("proba_interleaving_dataset must sum to 1")

        self.epoch = 0
        self.seed = sum(dataset.seed for dataset in self.datasets)
        self.rank = rank

        # state-related attributes
        self.start_worker_id = 0
        self.chain_dataset_last_idx_tracker = {}
        self.reset_state()

    def reset_state(self):
        self.chain_dataset_last_idx_tracker = {worker_id: 0 for worker_id in range(self.num_workers)}

    def update_state(self, dataset_state):
        self.chain_dataset_last_idx_tracker[dataset_state["worker_idx"]] = dataset_state["chain_dataset_last_idx"]

    def load_resume_states(self, resumable_states):
        for idx, d in enumerate(self.datasets):
            worker_idx_tracker, start_worker_id = resumable_states[idx]
            d.set_state(worker_idx_tracker, start_worker_id)
            self.start_worker_id = start_worker_id

    def state_dict(self):
        state_dict = {}

        state_dict["chain_dataset_last_idx_tracker"] = self.chain_dataset_last_idx_tracker

        # recurse into its datasets
        state_dict["datasets"] = [d.state_dict() for d in self.datasets]

        return state_dict

    def load_state_dict(self, state_dict):
        for key in state_dict["chain_dataset_last_idx_tracker"].keys():
            if key in self.chain_dataset_last_idx_tracker:
                self.chain_dataset_last_idx_tracker[key] = state_dict["chain_dataset_last_idx_tracker"][key]
            else:
                self.chain_dataset_last_idx_tracker[key] = 0

        # recurse into its datasets
        for idx, d in enumerate(self.datasets):
            d.load_state_dict(state_dict["datasets"][idx])

    def set_epoch(self, epoch):
        #  TODO: change this and fix trainer as well when epoch logic
        #  described in iter is implemented
        self.epoch = epoch
        for d in self.datasets:
            d.set_epoch(epoch)

    def _get_worker_id(self):
        worker_info = torch.utils.data.get_worker_info()
        if worker_info is None:
            worker_total_num = 1
            worker_id = 0
        else:
            worker_total_num = worker_info.num_workers
            worker_id = worker_info.id

        worker_id = (worker_id + self.start_worker_id) % worker_total_num
        return worker_id

    def __iter__(self):
        ds_iterators = [iter(d) for d in self.datasets]
        worker_id = self._get_worker_id()

        # Needed for dataset accumulation. Allows chain dataset to start at the right
        # dataset_idx, as it needs to know what was the last idx was.
        # Note that for validation set, this will always start from zero.
        dataset_idx = self.chain_dataset_last_idx_tracker[worker_id]
        while True:
            rng_seed = [
                self.seed,
                self.epoch,
                self.rank,
                worker_id,
                self.chain_dataset_last_idx_tracker[worker_id],
            ]
            rng = np.random.default_rng(seed=rng_seed)

            if self.is_train:
                if self.accumulate_datasets:
                    dataset_idx = (dataset_idx + 1) % self.num_datasets
                else:
                    dataset_idx = rng.choice(np.arange(0, self.num_datasets), p=self.dataset_proba)

            try:
                _, dataset_name, dataset_state, batch = next(ds_iterators[dataset_idx])
                self.chain_dataset_last_idx_tracker[worker_id] += 1
                dataset_state["chain_dataset_last_idx"] = self.chain_dataset_last_idx_tracker[worker_id]
                yield dataset_idx, dataset_name, dataset_state, batch
            except StopIteration:
                if not self.is_train and dataset_idx < self.num_datasets - 1:
                    dataset_idx += 1
                else:
                    self.epoch += 1
                    self.reset_state()

                    for d in self.datasets:
                        d.reset_state()
                        d.set_epoch(self.epoch)

                    ds_iterators = [iter(d) for d in self.datasets]

                    # TODO: Move epoch logic here instead of training loop
                    # ie: make an infinite dataloader that keeps track of the epochs of
                    # each dataset


# this class is currently not being maintained
class DataLoaderForMappedWrapperDataset(torch.utils.data.DataLoader):
    def state_dict(self):
        state_dict = {}

        # recurse into its dataset
        state_dict["dataset"] = self.dataset.state_dict()

        return state_dict

    def load_state_dict(self, state_dict):
        # recurse into its dataset
        self.dataset.load_state_dict(state_dict["dataset"])


# These are notes from IterableResumableState that got folded into
# DataLoaderForIterableWrapperDataset, but the comments are still relevant:
#
# **IMPORTANT**: Serializing "transform" randomness is really really difficult to get right because any "global"
#   generators we initialize are not actually going to be respected in any DataLoader with `num_workers > 1`. When you
#   run a DataLoader w/ multiprocessing, each "rng" generator is "copied" over to the separate process; setting that
#   generator only works _within_ the parts of the Dataset/Sampler/BatchSampler/DataLoader that are actually in scope
#   while an individual worker process is live... and there's no way I can currently think of to cleanly do that...


class DataLoaderForIterableWrapperDataset(torch.utils.data.DataLoader):
    def __init__(self, dataset, seed, **kwargs):
        self.seed = seed
        self.rank = kwargs.pop("rank", None)
        self.world_size = kwargs.pop("world_size", None)
        self.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

        if self.rank is None or self.world_size is None:
            try:
                state = AcceleratorState()
                self.rank = state.process_index
                self.device = state.device
                self.world_size = state.num_processes
            except ValueError:
                # This will fail later on when we try to create the IterableWrapperDataset
                pass

        if self.world_size == 1:
            # we normally don't use one process, but when we do in testing there can be a problem
            # with Too many open files issue, and switching to "file_system" strategy overcomes it.
            torch.multiprocessing.set_sharing_strategy("file_system")

        dataset.rank = self.rank
        dataset.world_size = self.world_size
        super().__init__(dataset, **kwargs)
        self.dataset_count = len(self.dataset.datasets)

        self.worker_idx_tracker = [{} for i in range(self.dataset_count)]
        self.next_worker_idx = [0 for i in range(self.dataset_count)]
        self.reset_state()

    def reset_state(self):
        for dataset_idx in range(self.dataset_count):
            self.worker_idx_tracker[dataset_idx] = {}
            self.next_worker_idx[dataset_idx] = 0

        # TODO: Aman: It looks like this has somewhat changed from last I checked. On slack, we
        # discussed to not have the states in the dataset themselves but just on the dataloader and
        # have current state returned from __iter__ of the dataset.
        #
        # The reason for this is what we discussed on slack related to dataset object being unique
        # to each worker and being passed around. So, the current dataset object in your dataloader
        # possibly won't give you correct state rather the state it had when it was initialized. Are
        # we confident that it will work this way?

        self.dataset.reset_state()

    def update_state(self, dataset_idx, dataset_state):
        self.dataset.update_state(dataset_state)

        self.worker_idx_tracker[dataset_idx][dataset_state["worker_idx"]] = (
            dataset_state["map_start_idx"],
            dataset_state["last_key_idx"],
            dataset_state["previous_overflow_batch"],
        )

        for dataset_idx in range(self.dataset_count):
            self.next_worker_idx[dataset_idx] = (dataset_state["worker_idx"] + 1) % self.num_workers

    def set_epoch(self, epoch):
        # Ensure each epoch has a different sequence to sample from
        self.dataset.set_epoch(epoch)

    def load_resume_states(self):
        self.dataset.load_resume_states(self.get_resume_states())

    def get_resume_state(self, dataset_idx):
        return self.worker_idx_tracker[dataset_idx], self.next_worker_idx[dataset_idx]

    def get_resume_states(self):
        return [
            [self.worker_idx_tracker[dataset_idx], self.next_worker_idx[dataset_idx]]
            for dataset_idx in range(self.dataset_count)
        ]

    def state_dict(self):
        state_dict = {}

        state_dict["worker_idx_tracker"] = self.worker_idx_tracker
        state_dict["next_worker_idx"] = self.next_worker_idx

        # recurse into its dataset
        state_dict["dataset"] = self.dataset.state_dict()

        return state_dict

    def load_state_dict(self, state_dict):
        self.worker_idx_tracker = state_dict["worker_idx_tracker"]
        self.next_worker_idx = state_dict["next_worker_idx"]

        # recurse into its dataset
        self.dataset.load_state_dict(state_dict["dataset"])

    def save_state(self, path):
        """
        Saves state_dict to `m4_states_{process_index}.pkl`
        """
        path = Path(path)
        path.mkdir(parents=True, exist_ok=True)
        location = path / f"m4_states_{self.rank}.pkl"
        logger.info(f"Saving the DL states to {location}")
        torch.save(self.state_dict(), location)

    def load_state(self, path):
        """
        Loads the state_dict at `m4_states_{process_index}.pkl`
        """
        location = Path(path) / f"m4_states_{self.rank}.pkl"
        logger.info(f"Loading the DL states from {location}")
        self.load_state_dict(torch.load(location))

    def __iter__(self):
        main_iterator = super().__iter__()

        stopped = torch.tensor(0.0, device=self.device)
        while True:
            try:
                batch = next(main_iterator)
            except StopIteration:
                stopped += 1

            # check that dist is initialized in case this DL w/ world_size>1 was used w/o distributed environment
            if self.world_size > 1 and torch.distributed.is_initialized():
                torch.distributed.all_reduce(stopped, op=torch.distributed.ReduceOp.SUM)

            # stop iterating if one or more processes stopped to avoid blocking
            if stopped > 0:
                break

            yield batch


# This class isn't being maintained (or tested) at the moment
class MappedResumableState:
    def __init__(self, train_seed, val_seed, world_size, rank, dataloader_num_workers, using_iterable_dataset=False):
        self.train_seed, self.val_seed, self.world_size, self.rank = train_seed, val_seed, world_size, rank
        self.dataloader_num_workers, self.using_iterable_dataset = dataloader_num_workers, using_iterable_dataset
        self.epoch = None

        self.n_seen_examples_per_worker = {w: 0 for w in range(self.dataloader_num_workers)}

        # Create generators for the various sources of randomness
        self.sampler_rng, self.sampler_rng_val = torch.Generator(), torch.Generator()
        self.sampler_rng.manual_seed(self.train_seed)
        self.sampler_rng_val.manual_seed(self.val_seed)

        # Create instance variables to track the `sampler_rng` states (see comment in `set_epoch` for why!)
        self.sampler_rng_state = self.sampler_rng.get_state()
        self.sampler_rng_val_state = self.sampler_rng_val.get_state()

    def set_epoch(self, epoch):
        # At the start of each epoch, the iter(sampler) gets called, which advances the random state; this isn't
        #   great, as it means that if we save the random state post-hoc, when we resume, we're "skipping" to the
        #   next epoch. To preserve the _same_ order within an epoch, we actually need to save the random state
        #   prior to the call to __iter__!
        if self.epoch != epoch:
            self.epoch = epoch

            self.sampler_rng_state = self.sampler_rng.get_state()
            self.sampler_rng_val_state = self.sampler_rng_val.get_state()

    def update_state(self, state_trackers):
        # Gate based on `self.using_iterable_dataset`

        # For a Map-style dataset, `state_trackers` is a Tensor containing all worker_ids responsible for fetching
        #   the given examples of the batch; we'll be using this to increment `n_seen_examples_per_worker`. We'll
        #   be using the key assumption that a single worker generates all elements for the batch!
        self.n_seen_examples_per_worker[state_trackers[0].item()] += state_trackers.numel()

    def update_next_worker_idx(self, curr_worker_idx):
        self.next_worker_idx = (curr_worker_idx + 1) % self.dataloader_num_workers

    def get_resume_state(self):
        # Return minimal `state` to "fast-forward" a given Dataset/DataLoader
        return self.n_seen_examples_per_worker

    def state_dict(self):
        return {
            "epoch": self.epoch,
            "world_size": self.world_size,
            "dataloader_num_workers": self.dataloader_num_workers,
            "sampler_rng_state": self.sampler_rng_state,
            "sampler_rng_state_checksum": self.sampler_rng_state.sum(),
            "sampler_rng_val_state": self.sampler_rng_val_state,
            "sampler_rng_val_state_checksum": self.sampler_rng_val_state.sum(),
            "n_seen_examples_per_worker": self.n_seen_examples_per_worker,
        }

    def load_state_dict(self, state_dict):
        # Verify same "world size" and "num workers" on a resume!
        if self.world_size != state_dict["world_size"]:
            raise ValueError(f"Current world_size {self.world_size} != Loaded world_size {state_dict['world_size']}")

        if self.dataloader_num_workers != state_dict["dataloader_num_workers"]:
            raise ValueError(
                f"Current num_workers `{self.dataloader_num_workers}` != Loaded num_workers"
                f" `{state_dict['dataloader_num_workers']}`"
            )

        # Set epoch
        self.epoch = state_dict["epoch"]

        # Set `sampler_rng`
        self.sampler_rng.set_state(state_dict["sampler_rng_state"])
        self.sampler_rng_val.set_state(state_dict["sampler_rng_val_state"])

        # Set `n_seen_examples_per_worker`
        self.n_seen_examples_per_worker = state_dict["n_seen_examples_per_worker"]