                            data_per_class_fraction,
                            extensions=extensions,
                            is_valid_file=is_valid_file)

    def _find_classes(self, dir: str) -> Tuple[List[str], Dict[str, int]]:
        """
        Finds the class folders in a dataset.
        Args:
            dir (string): Root directory path.
        Returns:
            tuple: (classes, class_to_idx) where classes are relative to (dir), and class_to_idx is a dictionary.
        Ensures:
            No class is a subdirectory of another.
        """
        all_classes = [d.name for d in os.scandir(dir) if d.is_dir()]
        classes = all_classes[0:int(len(all_classes) * self.classes_fraction)]
        classes.sort()
        class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)}
        return classes, class_to_idx

    def __getitem__(self, index: int) -> Tuple[Any, Any]:
        """
        Args:
            index (int): Index
        Returns:
            tuple: (sample, target) where target is class_index of the target class.
        """
        curr_index = index
        for x in range(self.total):
            try:
                path, target = self.samples[curr_index]
                sample = self.loader(path)
                break
            except Exception as e:
                curr_index = np.random.randint(0, self.total)

        if self.transform is not None:
            sample = self.transform(sample)
        if self.target_transform is not None:
            target = self.target_transform(target)

        return sample, target

    def __len__(self) -> int:
        return len(self.samples)

IMG_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp')

def pil_loader(path: str) -> Image.Image:
    # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
    with open(path, 'rb') as f:
        img = Image.open(f)
        return img.convert('RGB')

# TODO: specify the return type
def accimage_loader(path: str) -> Any:
    import accimage
    try:
        return accimage.Image(path)


            index (int): Index
        Returns:
            tuple: (sample, target) where target is class_index of the target class.
        """
        curr_index = index
        for x in range(self.total):
            try:
                path, target = self.samples[curr_index]
                sample = self.loader(path)
                break
            except Exception as e:
                curr_index = np.random.randint(0, self.total)

        if self.transform is not None:
            sample = self.transform(sample)
        if self.target_transform is not None:
            target = self.target_transform(target)

        return sample, target

    def __len__(self) -> int:
        return len(self.samples)

IMG_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp')

def pil_loader(path: str) -> Image.Image:
    # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
    with open(path, 'rb') as f:
        img = Image.open(f)
        return img.convert('RGB')

# TODO: specify the return type
def accimage_loader(path: str) -> Any:
    import accimage
    try:
        return accimage.Image(path)
    except IOError:
        # Potentially a decoding problem, fall back to PIL.Image
        return pil_loader(path)


        directory (str): root dataset directory
        class_to_idx (Dict[str, int]): dictionary mapping class name to class index
        extensions (optional): A list of allowed extensions.
            Either extensions or is_valid_file should be passed. Defaults to None.
        is_valid_file (optional): A function that takes path of a file
            and checks if the file is a valid file
            (used to check of corrupt files) both extensions and
            is_valid_file should not be passed. Defaults to None.
    Raises:
        ValueError: In case ``extensions`` and ``is_valid_file`` are None or both are not None.
    Returns:
        List[Tuple[str, int]]: samples of a form (path_to_sample, class)
    """
    instances = []
    directory = os.path.expanduser(directory)
    both_none = extensions is None and is_valid_file is None
    both_something = extensions is not None and is_valid_file is not None
    if both_none or both_something:
        raise ValueError("Both extensions and is_valid_file cannot be None or not None at the same time")
    if extensions is not None:
        def is_valid_file(x: str) -> bool:
            return has_file_allowed_extension(x, cast(Tuple[str, ...], extensions))
    is_valid_file = cast(Callable[[str], bool], is_valid_file)
    for target_class in sorted(class_to_idx.keys()):
        class_index = class_to_idx[target_class]
        target_dir = os.path.join(directory, target_class)
        if not os.path.isdir(target_dir):
            continue
        local_instances = []
        for root, _, fnames in sorted(os.walk(target_dir, followlinks=True)):
            for fname in sorted(fnames):
                path = os.path.join(root, fname)
                if is_valid_file(path):
                    item = path, class_index
                    local_instances.append(item)

        instances.extend(local_instances[0:int(len(local_instances) * data_per_class_fraction)])

    return instances

class DatasetFolder(VisionDataset):


# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
import os
import random
import numpy as np
import torch
import torchvision.transforms as T
from torchvision import datasets
from megatron import get_args
from megatron.data.image_folder import ImageFolder
from megatron.data.autoaugment import ImageNetPolicy
from megatron.data.data_samplers import RandomSeedDataset
from PIL import Image, ImageFilter, ImageOps

class GaussianBlur(object):
    """
    Apply Gaussian Blur to the PIL image.
    """
    def __init__(self, p=0.5, radius_min=0.1, radius_max=2.):
        self.prob = p
        self.radius_min = radius_min
        self.radius_max = radius_max

    def __call__(self, img):
        do_it = random.random() <= self.prob
        if not do_it:
            return img

        return img.filter(
            ImageFilter.GaussianBlur(
                radius=random.uniform(self.radius_min, self.radius_max)
            )
        )

class Solarization(object):
    """
    Apply Solarization to the PIL image.
    """
    def __init__(self, p):
        self.p = p

    def __call__(self, img):
        if random.random() < self.p:
            return ImageOps.solarize(img)
        else:
            return img

class ClassificationTransform():
    def __init__(self, image_size, train=True):
        args = get_args()
        assert args.fp16 or args.bf16
        self.data_type = torch.half if args.fp16 else torch.bfloat16
        if train:
            self.transform = T.Compose([
                T.RandomResizedCrop(image_size),
                T.RandomHorizontalFlip(),
                T.ColorJitter(0.4, 0.4, 0.4, 0.1),
                ImageNetPolicy(),
                T.ToTensor(),
                T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
                T.ConvertImageDtype(self.data_type)
            ])
        else:
            self.transform = T.Compose([
                T.Resize(image_size),
                T.CenterCrop(image_size),
                T.ToTensor(),
                T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
                T.ConvertImageDtype(self.data_type)
            ])

    def __call__(self, input):
        output = self.transform(input)
        return output

class InpaintingTransform():
    def __init__(self, image_size, train=True):

        args = get_args()
        self.mask_factor = args.mask_factor
        self.mask_type = args.mask_type
        self.image_size = image_size
        self.patch_size = args.patch_dim
        self.mask_size = int(self.mask_factor*(image_size[0]/self.patch_size)*(image_size[1]/self.patch_size))
        self.train = train
        assert args.fp16 or args.bf16
        self.data_type = torch.half if args.fp16 else torch.bfloat16

        if self.train:
            self.transform = T.Compose([
                T.RandomResizedCrop(self.image_size),
                T.RandomHorizontalFlip(),
                T.ColorJitter(0.4, 0.4, 0.4, 0.1),
                ImageNetPolicy(),
                T.ToTensor(),
                T.ConvertImageDtype(self.data_type)
            ])
        else:
            self.transform = T.Compose([
                T.Resize(self.image_size, interpolation=2),
                T.CenterCrop(self.image_size),
                T.ToTensor(),
                T.ConvertImageDtype(self.data_type)
            ])

    def gen_mask(self, image_size, mask_size, mask_type, patch_size):
        # output: mask as a list with indices for missing patches
        action_list = [[0, 1], [0, -1], [1, 0], [-1, 0]]
        assert image_size[0] == image_size[1]
        img_size_patch = image_size[0] // patch_size

        # drop masked patches
        mask = torch.zeros((image_size[0], image_size[1]), dtype=torch.float)

        if mask_type == 'random':
            x = torch.randint(0, img_size_patch, ())
            y = torch.randint(0, img_size_patch, ())
            for i in range(mask_size):
                r = torch.randint(0, len(action_list), ())
                x = torch.clamp(x + action_list[r][0], min=0, max=img_size_patch - 1)
                y = torch.clamp(y + action_list[r][1], min=0, max=img_size_patch - 1)
                x_offset = x * patch_size
                y_offset = y * patch_size
                mask[x_offset:x_offset+patch_size, y_offset:y_offset+patch_size] = 1
        else:
            assert mask_type == 'row'
            count = 0
            for x in reversed(range(img_size_patch)):
                for y in reversed(range(img_size_patch)):
                    if (count < mask_size):
                        count += 1
                        x_offset = x * patch_size
                        y_offset = y * patch_size
                        mask[x_offset:x_offset+patch_size, y_offset:y_offset+patch_size] = 1
        return mask

    def __call__(self, input):
        trans_input = self.transform(input)
        mask = self.gen_mask(self.image_size, self.mask_size,
			     self.mask_type, self.patch_size)
        mask = mask.unsqueeze(dim=0)
        return trans_input, mask

class DinoTransform(object):
    def __init__(self, image_size, train=True):
        args = get_args()
        self.data_type = torch.half if args.fp16 else torch.bfloat16

        flip_and_color_jitter = T.Compose([
            T.RandomHorizontalFlip(p=0.5),
            T.RandomApply(
                [T.ColorJitter(brightness=0.4, contrast=0.4,
			       saturation=0.2, hue=0.1)],
                p=0.8
            ),
            T.RandomGrayscale(p=0.2),
        ])

        if args.fp16 or args.bf16:
            normalize = T.Compose([
                T.ToTensor(),
                T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
                T.ConvertImageDtype(self.data_type)
            ])
        else:
            normalize = T.Compose([
                T.ToTensor(),
                T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
            ])

        # first global crop
        scale_const = 0.4
        self.global_transform1 = T.Compose([
            T.RandomResizedCrop(image_size,
                                scale=(scale_const, 1),
                                interpolation=Image.BICUBIC),
            flip_and_color_jitter,
            GaussianBlur(1.0),
            normalize
        ])
        # second global crop
        self.global_transform2 = T.Compose([
            T.RandomResizedCrop(image_size,
                                scale=(scale_const, 1),
                                interpolation=Image.BICUBIC),
            flip_and_color_jitter,
            GaussianBlur(0.1),
            Solarization(0.2),
            normalize
        ])
        # transformation for the local small crops
        self.local_crops_number = args.dino_local_crops_number
        self.local_transform = T.Compose([
            T.RandomResizedCrop(args.dino_local_img_size,
                                scale=(0.05, scale_const),
                                interpolation=Image.BICUBIC),
            flip_and_color_jitter,
            GaussianBlur(p=0.5),
            normalize
        ])

    def __call__(self, image):
        crops = []
        crops.append(self.global_transform1(image))
        crops.append(self.global_transform2(image))
        for _ in range(self.local_crops_number):
            crops.append(self.local_transform(image))
        return crops

def build_train_valid_datasets(data_path, image_size=224):
    args = get_args()

    if args.vision_pretraining_type == 'classify':
        train_transform = ClassificationTransform(image_size)
        val_transform = ClassificationTransform(image_size, train=False)
    elif args.vision_pretraining_type == 'inpaint':
        train_transform = InpaintingTransform(image_size, train=False)
        val_transform = InpaintingTransform(image_size, train=False)
    elif args.vision_pretraining_type == 'dino':
        train_transform = DinoTransform(image_size, train=True)
        val_transform = ClassificationTransform(image_size, train=False)
    else:
        raise Exception('{} vit pretraining type is not supported.'.format(
                args.vit_pretraining_type))

    # training dataset
    train_data_path = data_path[0] if len(data_path) <= 2 else data_path[2]
    train_data = ImageFolder(
        root=train_data_path,
        transform=train_transform,
        classes_fraction=args.classes_fraction,
        data_per_class_fraction=args.data_per_class_fraction
    )
    train_data = RandomSeedDataset(train_data)

    # validation dataset
    val_data_path = data_path[1]
    val_data = ImageFolder(
        root=val_data_path,
        transform=val_transform
    )
    val_data = RandomSeedDataset(val_data)

    return train_data, val_data


GaussianBlur

# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
import os
import random
import numpy as np
import torch
import torchvision.transforms as T
from torchvision import datasets
from megatron import get_args
from megatron.data.image_folder import ImageFolder
from megatron.data.autoaugment import ImageNetPolicy
from megatron.data.data_samplers import RandomSeedDataset
from PIL import Image, ImageFilter, ImageOps

class GaussianBlur(object):
    """
    Apply Gaussian Blur to the PIL image.
    """
    def __init__(self, p=0.5, radius_min=0.1, radius_max=2.):
        self.prob = p
        self.radius_min = radius_min
        self.radius_max = radius_max

    def __call__(self, img):
        do_it = random.random() <= self.prob
        if not do_it:
            return img

        return img.filter(
            ImageFilter.GaussianBlur(
                radius=random.uniform(self.radius_min, self.radius_max)
            )
        )

class Solarization(object):
    """
    Apply Solarization to the PIL image.
    """
    def __init__(self, p):
        self.p = p

    def __call__(self, img):
        if random.random() < self.p:
            return ImageOps.solarize(img)
        else:
            return img

class ClassificationTransform():
    def __init__(self, image_size, train=True):
        args = get_args()
        assert args.fp16 or args.bf16


Solarization

    """
    Apply Gaussian Blur to the PIL image.
    """
    def __init__(self, p=0.5, radius_min=0.1, radius_max=2.):
        self.prob = p
        self.radius_min = radius_min
        self.radius_max = radius_max

    def __call__(self, img):
        do_it = random.random() <= self.prob
        if not do_it:
            return img

        return img.filter(
            ImageFilter.GaussianBlur(
                radius=random.uniform(self.radius_min, self.radius_max)
            )
        )

class Solarization(object):
    """
    Apply Solarization to the PIL image.
    """
    def __init__(self, p):
        self.p = p

    def __call__(self, img):
        if random.random() < self.p:
            return ImageOps.solarize(img)
        else:
            return img

class ClassificationTransform():
    def __init__(self, image_size, train=True):
        args = get_args()
        assert args.fp16 or args.bf16
        self.data_type = torch.half if args.fp16 else torch.bfloat16
        if train:
            self.transform = T.Compose([
                T.RandomResizedCrop(image_size),
                T.RandomHorizontalFlip(),
                T.ColorJitter(0.4, 0.4, 0.4, 0.1),
                ImageNetPolicy(),
                T.ToTensor(),
                T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
                T.ConvertImageDtype(self.data_type)
            ])
        else:
            self.transform = T.Compose([
                T.Resize(image_size),


ClassificationTransform

            ImageFilter.GaussianBlur(
                radius=random.uniform(self.radius_min, self.radius_max)
            )
        )

class Solarization(object):
    """
    Apply Solarization to the PIL image.
    """
    def __init__(self, p):
        self.p = p

    def __call__(self, img):
        if random.random() < self.p:
            return ImageOps.solarize(img)
        else:
            return img

class ClassificationTransform():
    def __init__(self, image_size, train=True):
        args = get_args()
        assert args.fp16 or args.bf16
        self.data_type = torch.half if args.fp16 else torch.bfloat16
        if train:
            self.transform = T.Compose([
                T.RandomResizedCrop(image_size),
                T.RandomHorizontalFlip(),
                T.ColorJitter(0.4, 0.4, 0.4, 0.1),
                ImageNetPolicy(),
                T.ToTensor(),
                T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
                T.ConvertImageDtype(self.data_type)
            ])
        else:
            self.transform = T.Compose([
                T.Resize(image_size),
                T.CenterCrop(image_size),
                T.ToTensor(),
                T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
                T.ConvertImageDtype(self.data_type)
            ])

    def __call__(self, input):
        output = self.transform(input)
        return output

class InpaintingTransform():
    def __init__(self, image_size, train=True):

        args = get_args()
        self.mask_factor = args.mask_factor
        self.mask_type = args.mask_type
        self.image_size = image_size
        self.patch_size = args.patch_dim
        self.mask_size = int(self.mask_factor*(image_size[0]/self.patch_size)*(image_size[1]/self.patch_size))
        self.train = train
        assert args.fp16 or args.bf16
        self.data_type = torch.half if args.fp16 else torch.bfloat16

        if self.train:
            self.transform = T.Compose([
                T.RandomResizedCrop(self.image_size),
                T.RandomHorizontalFlip(),
                T.ColorJitter(0.4, 0.4, 0.4, 0.1),


InpaintingTransform

                T.ColorJitter(0.4, 0.4, 0.4, 0.1),
                ImageNetPolicy(),
                T.ToTensor(),
                T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
                T.ConvertImageDtype(self.data_type)
            ])
        else:
            self.transform = T.Compose([
                T.Resize(image_size),
                T.CenterCrop(image_size),
                T.ToTensor(),
                T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
                T.ConvertImageDtype(self.data_type)
            ])

    def __call__(self, input):
        output = self.transform(input)
        return output

class InpaintingTransform():
    def __init__(self, image_size, train=True):

        args = get_args()
        self.mask_factor = args.mask_factor
        self.mask_type = args.mask_type
        self.image_size = image_size
        self.patch_size = args.patch_dim
        self.mask_size = int(self.mask_factor*(image_size[0]/self.patch_size)*(image_size[1]/self.patch_size))
        self.train = train
        assert args.fp16 or args.bf16
        self.data_type = torch.half if args.fp16 else torch.bfloat16

        if self.train:
            self.transform = T.Compose([
                T.RandomResizedCrop(self.image_size),
                T.RandomHorizontalFlip(),
                T.ColorJitter(0.4, 0.4, 0.4, 0.1),
                ImageNetPolicy(),
                T.ToTensor(),
                T.ConvertImageDtype(self.data_type)
            ])
        else:
            self.transform = T.Compose([
                T.Resize(self.image_size, interpolation=2),
                T.CenterCrop(self.image_size),
                T.ToTensor(),
                T.ConvertImageDtype(self.data_type)
            ])

    def gen_mask(self, image_size, mask_size, mask_type, patch_size):
        # output: mask as a list with indices for missing patches
        action_list = [[0, 1], [0, -1], [1, 0], [-1, 0]]
        assert image_size[0] == image_size[1]
        img_size_patch = image_size[0] // patch_size

        # drop masked patches
        mask = torch.zeros((image_size[0], image_size[1]), dtype=torch.float)

        if mask_type == 'random':
            x = torch.randint(0, img_size_patch, ())
            y = torch.randint(0, img_size_patch, ())
            for i in range(mask_size):
                r = torch.randint(0, len(action_list), ())
                x = torch.clamp(x + action_list[r][0], min=0, max=img_size_patch - 1)
                y = torch.clamp(y + action_list[r][1], min=0, max=img_size_patch - 1)
                x_offset = x * patch_size
                y_offset = y * patch_size
                mask[x_offset:x_offset+patch_size, y_offset:y_offset+patch_size] = 1
        else:
            assert mask_type == 'row'
            count = 0
            for x in reversed(range(img_size_patch)):
                for y in reversed(range(img_size_patch)):
                    if (count < mask_size):
                        count += 1
                        x_offset = x * patch_size
                        y_offset = y * patch_size
                        mask[x_offset:x_offset+patch_size, y_offset:y_offset+patch_size] = 1
        return mask

    def __call__(self, input):
        trans_input = self.transform(input)
        mask = self.gen_mask(self.image_size, self.mask_size,
			     self.mask_type, self.patch_size)
        mask = mask.unsqueeze(dim=0)
        return trans_input, mask

class DinoTransform(object):
    def __init__(self, image_size, train=True):
        args = get_args()
        self.data_type = torch.half if args.fp16 else torch.bfloat16

        flip_and_color_jitter = T.Compose([
            T.RandomHorizontalFlip(p=0.5),
            T.RandomApply(
                [T.ColorJitter(brightness=0.4, contrast=0.4,
			       saturation=0.2, hue=0.1)],
                p=0.8
            ),
            T.RandomGrayscale(p=0.2),
        ])

        if args.fp16 or args.bf16:
            normalize = T.Compose([
                T.ToTensor(),


DinoTransform

        else:
            assert mask_type == 'row'
            count = 0
            for x in reversed(range(img_size_patch)):
                for y in reversed(range(img_size_patch)):
                    if (count < mask_size):
                        count += 1
                        x_offset = x * patch_size
                        y_offset = y * patch_size
                        mask[x_offset:x_offset+patch_size, y_offset:y_offset+patch_size] = 1
        return mask

    def __call__(self, input):
        trans_input = self.transform(input)
        mask = self.gen_mask(self.image_size, self.mask_size,
			     self.mask_type, self.patch_size)
        mask = mask.unsqueeze(dim=0)
        return trans_input, mask

class DinoTransform(object):
    def __init__(self, image_size, train=True):
        args = get_args()
        self.data_type = torch.half if args.fp16 else torch.bfloat16

        flip_and_color_jitter = T.Compose([
            T.RandomHorizontalFlip(p=0.5),
            T.RandomApply(
                [T.ColorJitter(brightness=0.4, contrast=0.4,
			       saturation=0.2, hue=0.1)],
                p=0.8
            ),
            T.RandomGrayscale(p=0.2),
        ])

        if args.fp16 or args.bf16:
            normalize = T.Compose([
                T.ToTensor(),
                T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
                T.ConvertImageDtype(self.data_type)
            ])
        else:
            normalize = T.Compose([
                T.ToTensor(),
                T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
            ])

        # first global crop
        scale_const = 0.4
        self.global_transform1 = T.Compose([
            T.RandomResizedCrop(image_size,
                                scale=(scale_const, 1),
                                interpolation=Image.BICUBIC),
            flip_and_color_jitter,
            GaussianBlur(1.0),
            normalize
        ])
        # second global crop
        self.global_transform2 = T.Compose([
            T.RandomResizedCrop(image_size,
                                scale=(scale_const, 1),
                                interpolation=Image.BICUBIC),
            flip_and_color_jitter,
            GaussianBlur(0.1),
            Solarization(0.2),
            normalize
        ])
        # transformation for the local small crops
        self.local_crops_number = args.dino_local_crops_number
        self.local_transform = T.Compose([
            T.RandomResizedCrop(args.dino_local_img_size,
                                scale=(0.05, scale_const),
                                interpolation=Image.BICUBIC),
            flip_and_color_jitter,
            GaussianBlur(p=0.5),
            normalize
        ])

    def __call__(self, image):
        crops = []
        crops.append(self.global_transform1(image))
        crops.append(self.global_transform2(image))
        for _ in range(self.local_crops_number):
            crops.append(self.local_transform(image))
        return crops

def build_train_valid_datasets(data_path, image_size=224):
    args = get_args()

    if args.vision_pretraining_type == 'classify':
        train_transform = ClassificationTransform(image_size)
        val_transform = ClassificationTransform(image_size, train=False)
    elif args.vision_pretraining_type == 'inpaint':
        train_transform = InpaintingTransform(image_size, train=False)
        val_transform = InpaintingTransform(image_size, train=False)
    elif args.vision_pretraining_type == 'dino':
        train_transform = DinoTransform(image_size, train=True)
        val_transform = ClassificationTransform(image_size, train=False)
    else:
        raise Exception('{} vit pretraining type is not supported.'.format(
                args.vit_pretraining_type))

    # training dataset
    train_data_path = data_path[0] if len(data_path) <= 2 else data_path[2]


        # transformation for the local small crops
        self.local_crops_number = args.dino_local_crops_number
        self.local_transform = T.Compose([
            T.RandomResizedCrop(args.dino_local_img_size,
                                scale=(0.05, scale_const),
                                interpolation=Image.BICUBIC),
            flip_and_color_jitter,
            GaussianBlur(p=0.5),
            normalize
        ])

    def __call__(self, image):
        crops = []
        crops.append(self.global_transform1(image))
        crops.append(self.global_transform2(image))
        for _ in range(self.local_crops_number):
            crops.append(self.local_transform(image))
        return crops

def build_train_valid_datasets(data_path, image_size=224):
    args = get_args()

    if args.vision_pretraining_type == 'classify':
        train_transform = ClassificationTransform(image_size)
        val_transform = ClassificationTransform(image_size, train=False)
    elif args.vision_pretraining_type == 'inpaint':
        train_transform = InpaintingTransform(image_size, train=False)
        val_transform = InpaintingTransform(image_size, train=False)
    elif args.vision_pretraining_type == 'dino':
        train_transform = DinoTransform(image_size, train=True)
        val_transform = ClassificationTransform(image_size, train=False)
    else:
        raise Exception('{} vit pretraining type is not supported.'.format(
                args.vit_pretraining_type))

    # training dataset
    train_data_path = data_path[0] if len(data_path) <= 2 else data_path[2]
    train_data = ImageFolder(
        root=train_data_path,
        transform=train_transform,
        classes_fraction=args.classes_fraction,
        data_per_class_fraction=args.data_per_class_fraction
    )
    train_data = RandomSeedDataset(train_data)

    # validation dataset
    val_data_path = data_path[1]
    val_data = ImageFolder(
        root=val_data_path,
        transform=val_transform
    )
    val_data = RandomSeedDataset(val_data)

    return train_data, val_data


            assert mask_type == 'row'
            count = 0
            for x in reversed(range(img_size_patch)):
                for y in reversed(range(img_size_patch)):
                    if (count < mask_size):
                        count += 1
                        x_offset = x * patch_size
                        y_offset = y * patch_size
                        mask[x_offset:x_offset+patch_size, y_offset:y_offset+patch_size] = 1
        return mask

    def __call__(self, input):
        trans_input = self.transform(input)
        mask = self.gen_mask(self.image_size, self.mask_size,
			     self.mask_type, self.patch_size)
        mask = mask.unsqueeze(dim=0)
        return trans_input, mask

class DinoTransform(object):
    def __init__(self, image_size, train=True):
        args = get_args()
        self.data_type = torch.half if args.fp16 else torch.bfloat16

        flip_and_color_jitter = T.Compose([
            T.RandomHorizontalFlip(p=0.5),
            T.RandomApply(
                [T.ColorJitter(brightness=0.4, contrast=0.4,
			       saturation=0.2, hue=0.1)],
                p=0.8
            ),
            T.RandomGrayscale(p=0.2),
        ])

        if args.fp16 or args.bf16:
            normalize = T.Compose([
                T.ToTensor(),
                T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
                T.ConvertImageDtype(self.data_type)
            ])
        else:
            normalize = T.Compose([
                T.ToTensor(),
                T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
            ])

        # first global crop
        scale_const = 0.4
        self.global_transform1 = T.Compose([
            T.RandomResizedCrop(image_size,
                                scale=(scale_const, 1),
                                interpolation=Image.BICUBIC),
            flip_and_color_jitter,
            GaussianBlur(1.0),
            normalize
        ])
        # second global crop
        self.global_transform2 = T.Compose([
            T.RandomResizedCrop(image_size,
                                scale=(scale_const, 1),
                                interpolation=Image.BICUBIC),
            flip_and_color_jitter,
            GaussianBlur(0.1),
            Solarization(0.2),
            normalize
        ])
        # transformation for the local small crops
        self.local_crops_number = args.dino_local_crops_number
        self.local_transform = T.Compose([
            T.RandomResizedCrop(args.dino_local_img_size,
                                scale=(0.05, scale_const),
                                interpolation=Image.BICUBIC),
            flip_and_color_jitter,
            GaussianBlur(p=0.5),
            normalize
        ])

    def __call__(self, image):
        crops = []
        crops.append(self.global_transform1(image))
        crops.append(self.global_transform2(image))
        for _ in range(self.local_crops_number):
            crops.append(self.local_transform(image))
        return crops

def build_train_valid_datasets(data_path, image_size=224):
    args = get_args()

    if args.vision_pretraining_type == 'classify':
        train_transform = ClassificationTransform(image_size)
        val_transform = ClassificationTransform(image_size, train=False)
    elif args.vision_pretraining_type == 'inpaint':
        train_transform = InpaintingTransform(image_size, train=False)
        val_transform = InpaintingTransform(image_size, train=False)
    elif args.vision_pretraining_type == 'dino':


        self.global_transform2 = T.Compose([
            T.RandomResizedCrop(image_size,
                                scale=(scale_const, 1),
                                interpolation=Image.BICUBIC),
            flip_and_color_jitter,
            GaussianBlur(0.1),
            Solarization(0.2),
            normalize
        ])
        # transformation for the local small crops
        self.local_crops_number = args.dino_local_crops_number
        self.local_transform = T.Compose([
            T.RandomResizedCrop(args.dino_local_img_size,
                                scale=(0.05, scale_const),
                                interpolation=Image.BICUBIC),
            flip_and_color_jitter,
            GaussianBlur(p=0.5),
            normalize
        ])

    def __call__(self, image):
        crops = []
        crops.append(self.global_transform1(image))
        crops.append(self.global_transform2(image))
        for _ in range(self.local_crops_number):
            crops.append(self.local_transform(image))
        return crops

def build_train_valid_datasets(data_path, image_size=224):
    args = get_args()

    if args.vision_pretraining_type == 'classify':
        train_transform = ClassificationTransform(image_size)
        val_transform = ClassificationTransform(image_size, train=False)
    elif args.vision_pretraining_type == 'inpaint':
        train_transform = InpaintingTransform(image_size, train=False)
        val_transform = InpaintingTransform(image_size, train=False)
    elif args.vision_pretraining_type == 'dino':
        train_transform = DinoTransform(image_size, train=True)
        val_transform = ClassificationTransform(image_size, train=False)
    else:
        raise Exception('{} vit pretraining type is not supported.'.format(
                args.vit_pretraining_type))

    # training dataset
    train_data_path = data_path[0] if len(data_path) <= 2 else data_path[2]


        assert args.fp16 or args.bf16
        self.data_type = torch.half if args.fp16 else torch.bfloat16

        if self.train:
            self.transform = T.Compose([
                T.RandomResizedCrop(self.image_size),
                T.RandomHorizontalFlip(),
                T.ColorJitter(0.4, 0.4, 0.4, 0.1),
                ImageNetPolicy(),
                T.ToTensor(),
                T.ConvertImageDtype(self.data_type)
            ])
        else:
            self.transform = T.Compose([
                T.Resize(self.image_size, interpolation=2),
                T.CenterCrop(self.image_size),
                T.ToTensor(),
                T.ConvertImageDtype(self.data_type)
            ])

    def gen_mask(self, image_size, mask_size, mask_type, patch_size):
        # output: mask as a list with indices for missing patches
        action_list = [[0, 1], [0, -1], [1, 0], [-1, 0]]
        assert image_size[0] == image_size[1]
        img_size_patch = image_size[0] // patch_size

        # drop masked patches
        mask = torch.zeros((image_size[0], image_size[1]), dtype=torch.float)

        if mask_type == 'random':
            x = torch.randint(0, img_size_patch, ())
            y = torch.randint(0, img_size_patch, ())
            for i in range(mask_size):
                r = torch.randint(0, len(action_list), ())
                x = torch.clamp(x + action_list[r][0], min=0, max=img_size_patch - 1)
                y = torch.clamp(y + action_list[r][1], min=0, max=img_size_patch - 1)
                x_offset = x * patch_size
                y_offset = y * patch_size
                mask[x_offset:x_offset+patch_size, y_offset:y_offset+patch_size] = 1
        else:
            assert mask_type == 'row'
            count = 0
            for x in reversed(range(img_size_patch)):
                for y in reversed(range(img_size_patch)):
                    if (count < mask_size):
                        count += 1
                        x_offset = x * patch_size
                        y_offset = y * patch_size
                        mask[x_offset:x_offset+patch_size, y_offset:y_offset+patch_size] = 1
        return mask

    def __call__(self, input):
        trans_input = self.transform(input)
        mask = self.gen_mask(self.image_size, self.mask_size,
			     self.mask_type, self.patch_size)
        mask = mask.unsqueeze(dim=0)
        return trans_input, mask

class DinoTransform(object):
    def __init__(self, image_size, train=True):
        args = get_args()
        self.data_type = torch.half if args.fp16 else torch.bfloat16

        flip_and_color_jitter = T.Compose([
            T.RandomHorizontalFlip(p=0.5),
            T.RandomApply(
                [T.ColorJitter(brightness=0.4, contrast=0.4,
			       saturation=0.2, hue=0.1)],
                p=0.8


# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.

"""Wikipedia dataset from DPR code for ORQA."""

from abc import ABC
import csv
import numpy as np
import random
import torch
from torch.utils.data import Dataset

from megatron import print_rank_0, get_args, get_tokenizer
from megatron.core import tensor_parallel
from megatron.data.biencoder_dataset_utils import make_attention_mask

def get_open_retrieval_wiki_dataset():
    args = get_args()
    tokenizer = get_tokenizer()

    dataset = OpenRetrievalEvidenceDataset('2018 Wikipedia from DPR codebase',
                                           'evidence',
                                           args.evidence_data_path,
                                           tokenizer,
                                           args.retriever_seq_length)
    return dataset

def get_open_retrieval_batch(data_iterator):
    # Items and their type.
    keys = ['row_id', 'context', 'context_mask', 'context_types',
        'context_pad_mask']
    datatype = torch.int64

    # Broadcast data.
    data = None if data_iterator is None else next(data_iterator)
    data_b = tensor_parallel.broadcast_data(keys, data, datatype)

    # Unpack.
    row_id = data_b['row_id'].long()
    context = data_b['context'].long()

    # TODO: make the context mask a binary one
    context_mask = (data_b['context_mask'] < 0.5)

    context_types = data_b['context_types'].long()
    context_pad_mask = data_b['context_pad_mask'].long()

    return row_id, context, context_mask, context_types, context_pad_mask

def build_tokens_types_paddings_from_text(row, tokenizer, max_seq_length):
    """Build token types and paddings, trim if needed, and pad if needed."""

    title_ids = tokenizer.tokenize(row['title'])
    context_ids = tokenizer.tokenize(row['text'])

    # Appending the title of the context at front
    extended_context_ids = title_ids + [tokenizer.sep_id] + context_ids

    context_ids, context_types, context_pad_mask = \
        build_tokens_types_paddings_from_ids(extended_context_ids,
            max_seq_length, tokenizer.cls, tokenizer.sep, tokenizer.pad)

    return context_ids, context_types, context_pad_mask

# noinspection DuplicatedCode
def build_tokens_types_paddings_from_ids(text_ids, max_seq_length,
                                         cls_id, sep_id, pad_id):
    """Build token types and paddings, trim if needed, and pad if needed."""
    enc_ids = []
    tokentypes_enc = []

    # [CLS].
    enc_ids.append(cls_id)
    tokentypes_enc.append(0)

    # A.
    len_src = len(text_ids)
    enc_ids.extend(text_ids)
    tokentypes_enc.extend([0] * len_src)

    # Cap the size.
    if len(enc_ids) > max_seq_length - 1:
        enc_ids = enc_ids[0: max_seq_length - 1]
        tokentypes_enc = tokentypes_enc[0: max_seq_length - 1]

    # [SEP].
    enc_ids.append(sep_id)
    tokentypes_enc.append(0)

    num_tokens_enc = len(enc_ids)
    # Padding.
    padding_length = max_seq_length - len(enc_ids)
    if padding_length > 0:
        enc_ids.extend([pad_id] * padding_length)
        tokentypes_enc.extend([pad_id] * padding_length)

    pad_mask = ([1] * num_tokens_enc) + ([0] * padding_length)
    pad_mask = np.array(pad_mask, dtype=np.int64)

    return enc_ids, tokentypes_enc, pad_mask

def build_sample(row_id, context_ids, context_types, context_pad_mask):
    """Convert to numpy and return a sample consumed by the batch producer."""

    context_ids = np.array(context_ids, dtype=np.int64)
    context_types = np.array(context_types, dtype=np.int64)
    context_mask = make_attention_mask(context_ids, context_ids)

    sample = ({
        'row_id': row_id,
        'context': context_ids,
        'context_mask': context_mask,
        'context_types': context_types,
        'context_pad_mask': context_pad_mask
    })
    return sample

class OpenRetrievalEvidenceDataset(ABC, Dataset):
    """Open Retrieval Evidence dataset class."""

    def __init__(self, task_name, dataset_name, datapath, tokenizer,
            max_seq_length):
        # Store inputs.
        self.task_name = task_name
        self.dataset_name = dataset_name
        self.tokenizer = tokenizer
        self.max_seq_length = max_seq_length
        print_rank_0(' > building {} dataset for {}:'.format(self.task_name,
                                                            self.dataset_name))
        # Process the files.
        print_rank_0(datapath)
        self.samples, self.id2text = self.process_samples_from_single_path(
                                        datapath)

        args = get_args()
        if args.sample_rate < 1:  # subsample
            k = int(len(self.samples) * args.sample_rate)
            self.samples = random.sample(self.samples, k)

        print_rank_0('  >> total number of samples: {}'.format(
            len(self.samples)))

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

    def __getitem__(self, idx):
        row = self.samples[idx]

        context_ids, context_types, context_pad_mask = \
            build_tokens_types_paddings_from_text(row, self.tokenizer,
                self.max_seq_length)

        sample = build_sample(row['doc_id'],
                              context_ids,
                              context_types,
                              context_pad_mask)
        return sample

    @staticmethod
    def process_samples_from_single_path(filename):
        print_rank_0(' > Processing {} ...'.format(filename))
        total = 0

        rows = []
        id2text = {}

        with open(filename) as tsvfile:
            reader = csv.reader(tsvfile, delimiter='\t')
            next(reader, None)  # skip the headers
            for row in reader:
                # file format: doc_id, doc_text, title
                doc_id = int(row[0])
                text = row[1]
                title = row[2]

                rows.append({'doc_id': doc_id,
                             'text': text,
                             'title': title})

                assert doc_id not in id2text
                id2text[doc_id] = (text, title)

                total += 1
                if total % 100000 == 0:
                    print_rank_0('  > processed {} rows so far ...'.format(
                        total))

        print_rank_0(' >> processed {} samples.'.format(len(rows)))
        return rows, id2text


# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.

"""Wikipedia dataset from DPR code for ORQA."""

from abc import ABC
import csv
import numpy as np
import random
import torch
from torch.utils.data import Dataset

from megatron import print_rank_0, get_args, get_tokenizer
from megatron.core import tensor_parallel
from megatron.data.biencoder_dataset_utils import make_attention_mask

def get_open_retrieval_wiki_dataset():
    args = get_args()
    tokenizer = get_tokenizer()

    dataset = OpenRetrievalEvidenceDataset('2018 Wikipedia from DPR codebase',
                                           'evidence',
                                           args.evidence_data_path,
                                           tokenizer,
                                           args.retriever_seq_length)
    return dataset

def get_open_retrieval_batch(data_iterator):
    # Items and their type.
    keys = ['row_id', 'context', 'context_mask', 'context_types',
        'context_pad_mask']
    datatype = torch.int64

    # Broadcast data.
    data = None if data_iterator is None else next(data_iterator)
    data_b = tensor_parallel.broadcast_data(keys, data, datatype)

    # Unpack.
    row_id = data_b['row_id'].long()
    context = data_b['context'].long()

    # TODO: make the context mask a binary one
    context_mask = (data_b['context_mask'] < 0.5)

    context_types = data_b['context_types'].long()


import random
import torch
from torch.utils.data import Dataset

from megatron import print_rank_0, get_args, get_tokenizer
from megatron.core import tensor_parallel
from megatron.data.biencoder_dataset_utils import make_attention_mask

def get_open_retrieval_wiki_dataset():
    args = get_args()
    tokenizer = get_tokenizer()

    dataset = OpenRetrievalEvidenceDataset('2018 Wikipedia from DPR codebase',
                                           'evidence',
                                           args.evidence_data_path,
                                           tokenizer,
                                           args.retriever_seq_length)
    return dataset

def get_open_retrieval_batch(data_iterator):
    # Items and their type.
    keys = ['row_id', 'context', 'context_mask', 'context_types',
        'context_pad_mask']
    datatype = torch.int64

    # Broadcast data.
    data = None if data_iterator is None else next(data_iterator)
    data_b = tensor_parallel.broadcast_data(keys, data, datatype)

    # Unpack.
    row_id = data_b['row_id'].long()
    context = data_b['context'].long()

    # TODO: make the context mask a binary one
    context_mask = (data_b['context_mask'] < 0.5)

    context_types = data_b['context_types'].long()
    context_pad_mask = data_b['context_pad_mask'].long()

    return row_id, context, context_mask, context_types, context_pad_mask

def build_tokens_types_paddings_from_text(row, tokenizer, max_seq_length):
    """Build token types and paddings, trim if needed, and pad if needed."""

    title_ids = tokenizer.tokenize(row['title'])
    context_ids = tokenizer.tokenize(row['text'])

    # Appending the title of the context at front
    extended_context_ids = title_ids + [tokenizer.sep_id] + context_ids

    context_ids, context_types, context_pad_mask = \
        build_tokens_types_paddings_from_ids(extended_context_ids,
            max_seq_length, tokenizer.cls, tokenizer.sep, tokenizer.pad)

    return context_ids, context_types, context_pad_mask

# noinspection DuplicatedCode
def build_tokens_types_paddings_from_ids(text_ids, max_seq_length,


        'context_pad_mask']
    datatype = torch.int64

    # Broadcast data.
    data = None if data_iterator is None else next(data_iterator)
    data_b = tensor_parallel.broadcast_data(keys, data, datatype)

    # Unpack.
    row_id = data_b['row_id'].long()
    context = data_b['context'].long()

    # TODO: make the context mask a binary one
    context_mask = (data_b['context_mask'] < 0.5)

    context_types = data_b['context_types'].long()
    context_pad_mask = data_b['context_pad_mask'].long()

    return row_id, context, context_mask, context_types, context_pad_mask

def build_tokens_types_paddings_from_text(row, tokenizer, max_seq_length):
    """Build token types and paddings, trim if needed, and pad if needed."""

    title_ids = tokenizer.tokenize(row['title'])
    context_ids = tokenizer.tokenize(row['text'])

    # Appending the title of the context at front
    extended_context_ids = title_ids + [tokenizer.sep_id] + context_ids

    context_ids, context_types, context_pad_mask = \
        build_tokens_types_paddings_from_ids(extended_context_ids,
            max_seq_length, tokenizer.cls, tokenizer.sep, tokenizer.pad)

    return context_ids, context_types, context_pad_mask

# noinspection DuplicatedCode
def build_tokens_types_paddings_from_ids(text_ids, max_seq_length,
                                         cls_id, sep_id, pad_id):
    """Build token types and paddings, trim if needed, and pad if needed."""
    enc_ids = []
    tokentypes_enc = []

    # [CLS].
    enc_ids.append(cls_id)
    tokentypes_enc.append(0)

    # A.
    len_src = len(text_ids)
    enc_ids.extend(text_ids)
    tokentypes_enc.extend([0] * len_src)

    # Cap the size.
    if len(enc_ids) > max_seq_length - 1:


    return row_id, context, context_mask, context_types, context_pad_mask

def build_tokens_types_paddings_from_text(row, tokenizer, max_seq_length):
    """Build token types and paddings, trim if needed, and pad if needed."""

    title_ids = tokenizer.tokenize(row['title'])
    context_ids = tokenizer.tokenize(row['text'])

    # Appending the title of the context at front
    extended_context_ids = title_ids + [tokenizer.sep_id] + context_ids

    context_ids, context_types, context_pad_mask = \
        build_tokens_types_paddings_from_ids(extended_context_ids,
            max_seq_length, tokenizer.cls, tokenizer.sep, tokenizer.pad)

    return context_ids, context_types, context_pad_mask

# noinspection DuplicatedCode
def build_tokens_types_paddings_from_ids(text_ids, max_seq_length,
                                         cls_id, sep_id, pad_id):
    """Build token types and paddings, trim if needed, and pad if needed."""
    enc_ids = []
    tokentypes_enc = []

    # [CLS].
    enc_ids.append(cls_id)
    tokentypes_enc.append(0)

    # A.
    len_src = len(text_ids)
    enc_ids.extend(text_ids)
    tokentypes_enc.extend([0] * len_src)

    # Cap the size.
    if len(enc_ids) > max_seq_length - 1:
        enc_ids = enc_ids[0: max_seq_length - 1]
        tokentypes_enc = tokentypes_enc[0: max_seq_length - 1]

    # [SEP].
    enc_ids.append(sep_id)
    tokentypes_enc.append(0)

    num_tokens_enc = len(enc_ids)
    # Padding.
    padding_length = max_seq_length - len(enc_ids)
    if padding_length > 0:
        enc_ids.extend([pad_id] * padding_length)
        tokentypes_enc.extend([pad_id] * padding_length)

    pad_mask = ([1] * num_tokens_enc) + ([0] * padding_length)
    pad_mask = np.array(pad_mask, dtype=np.int64)

    return enc_ids, tokentypes_enc, pad_mask

def build_sample(row_id, context_ids, context_types, context_pad_mask):
    """Convert to numpy and return a sample consumed by the batch producer."""

    context_ids = np.array(context_ids, dtype=np.int64)
    context_types = np.array(context_types, dtype=np.int64)
    context_mask = make_attention_mask(context_ids, context_ids)

    sample = ({
        'row_id': row_id,
        'context': context_ids,
        'context_mask': context_mask,
        'context_types': context_types,
        'context_pad_mask': context_pad_mask
    })
    return sample

class OpenRetrievalEvidenceDataset(ABC, Dataset):


        enc_ids = enc_ids[0: max_seq_length - 1]
        tokentypes_enc = tokentypes_enc[0: max_seq_length - 1]

    # [SEP].
    enc_ids.append(sep_id)
    tokentypes_enc.append(0)

    num_tokens_enc = len(enc_ids)
    # Padding.
    padding_length = max_seq_length - len(enc_ids)
    if padding_length > 0:
        enc_ids.extend([pad_id] * padding_length)
        tokentypes_enc.extend([pad_id] * padding_length)

    pad_mask = ([1] * num_tokens_enc) + ([0] * padding_length)
    pad_mask = np.array(pad_mask, dtype=np.int64)

    return enc_ids, tokentypes_enc, pad_mask

def build_sample(row_id, context_ids, context_types, context_pad_mask):
    """Convert to numpy and return a sample consumed by the batch producer."""

    context_ids = np.array(context_ids, dtype=np.int64)
    context_types = np.array(context_types, dtype=np.int64)
    context_mask = make_attention_mask(context_ids, context_ids)

    sample = ({
        'row_id': row_id,
        'context': context_ids,
        'context_mask': context_mask,
        'context_types': context_types,
        'context_pad_mask': context_pad_mask
    })
    return sample

class OpenRetrievalEvidenceDataset(ABC, Dataset):
    """Open Retrieval Evidence dataset class."""

    def __init__(self, task_name, dataset_name, datapath, tokenizer,
            max_seq_length):
        # Store inputs.
        self.task_name = task_name
        self.dataset_name = dataset_name
        self.tokenizer = tokenizer
        self.max_seq_length = max_seq_length
        print_rank_0(' > building {} dataset for {}:'.format(self.task_name,
                                                            self.dataset_name))
        # Process the files.
        print_rank_0(datapath)
        self.samples, self.id2text = self.process_samples_from_single_path(
                                        datapath)

        args = get_args()


OpenRetrievalEvidenceDataset

    return enc_ids, tokentypes_enc, pad_mask

def build_sample(row_id, context_ids, context_types, context_pad_mask):
    """Convert to numpy and return a sample consumed by the batch producer."""

    context_ids = np.array(context_ids, dtype=np.int64)
    context_types = np.array(context_types, dtype=np.int64)
    context_mask = make_attention_mask(context_ids, context_ids)

    sample = ({
        'row_id': row_id,
        'context': context_ids,
        'context_mask': context_mask,
        'context_types': context_types,
        'context_pad_mask': context_pad_mask
    })
    return sample

class OpenRetrievalEvidenceDataset(ABC, Dataset):
    """Open Retrieval Evidence dataset class."""

    def __init__(self, task_name, dataset_name, datapath, tokenizer,
            max_seq_length):
        # Store inputs.
        self.task_name = task_name
        self.dataset_name = dataset_name
        self.tokenizer = tokenizer
        self.max_seq_length = max_seq_length
        print_rank_0(' > building {} dataset for {}:'.format(self.task_name,
                                                            self.dataset_name))
        # Process the files.
        print_rank_0(datapath)
        self.samples, self.id2text = self.process_samples_from_single_path(
                                        datapath)

        args = get_args()
        if args.sample_rate < 1:  # subsample
            k = int(len(self.samples) * args.sample_rate)
            self.samples = random.sample(self.samples, k)

        print_rank_0('  >> total number of samples: {}'.format(
            len(self.samples)))

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

    def __getitem__(self, idx):
        row = self.samples[idx]

        context_ids, context_types, context_pad_mask = \
            build_tokens_types_paddings_from_text(row, self.tokenizer,
                self.max_seq_length)

        sample = build_sample(row['doc_id'],
                              context_ids,
                              context_types,
                              context_pad_mask)
        return sample

    @staticmethod
    def process_samples_from_single_path(filename):
        print_rank_0(' > Processing {} ...'.format(filename))
        total = 0

        rows = []
        id2text = {}

        with open(filename) as tsvfile:
            reader = csv.reader(tsvfile, delimiter='\t')
            next(reader, None)  # skip the headers
            for row in reader:
                # file format: doc_id, doc_text, title
                doc_id = int(row[0])
                text = row[1]
                title = row[2]

                rows.append({'doc_id': doc_id,
                             'text': text,
                             'title': title})

                assert doc_id not in id2text
                id2text[doc_id] = (text, title)

                total += 1
                if total % 100000 == 0:
                    print_rank_0('  > processed {} rows so far ...'.format(
                        total))

        print_rank_0(' >> processed {} samples.'.format(len(rows)))
        return rows, id2text


def build_sample(row_id, context_ids, context_types, context_pad_mask):
    """Convert to numpy and return a sample consumed by the batch producer."""

    context_ids = np.array(context_ids, dtype=np.int64)
    context_types = np.array(context_types, dtype=np.int64)
    context_mask = make_attention_mask(context_ids, context_ids)

    sample = ({
        'row_id': row_id,
        'context': context_ids,
        'context_mask': context_mask,
        'context_types': context_types,
        'context_pad_mask': context_pad_mask
    })
    return sample

class OpenRetrievalEvidenceDataset(ABC, Dataset):
    """Open Retrieval Evidence dataset class."""

    def __init__(self, task_name, dataset_name, datapath, tokenizer,
            max_seq_length):
        # Store inputs.
        self.task_name = task_name
        self.dataset_name = dataset_name
        self.tokenizer = tokenizer
        self.max_seq_length = max_seq_length
        print_rank_0(' > building {} dataset for {}:'.format(self.task_name,
                                                            self.dataset_name))
        # Process the files.
        print_rank_0(datapath)
        self.samples, self.id2text = self.process_samples_from_single_path(
                                        datapath)

        args = get_args()
        if args.sample_rate < 1:  # subsample
            k = int(len(self.samples) * args.sample_rate)
            self.samples = random.sample(self.samples, k)

        print_rank_0('  >> total number of samples: {}'.format(
            len(self.samples)))

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

    def __getitem__(self, idx):
        row = self.samples[idx]

        context_ids, context_types, context_pad_mask = \
            build_tokens_types_paddings_from_text(row, self.tokenizer,
                self.max_seq_length)

        sample = build_sample(row['doc_id'],
                              context_ids,
                              context_types,
                              context_pad_mask)
        return sample

    @staticmethod
    def process_samples_from_single_path(filename):
        print_rank_0(' > Processing {} ...'.format(filename))


        # Store inputs.
        self.task_name = task_name
        self.dataset_name = dataset_name
        self.tokenizer = tokenizer
        self.max_seq_length = max_seq_length
        print_rank_0(' > building {} dataset for {}:'.format(self.task_name,
                                                            self.dataset_name))
        # Process the files.
        print_rank_0(datapath)
        self.samples, self.id2text = self.process_samples_from_single_path(
                                        datapath)

        args = get_args()
        if args.sample_rate < 1:  # subsample
            k = int(len(self.samples) * args.sample_rate)
            self.samples = random.sample(self.samples, k)

        print_rank_0('  >> total number of samples: {}'.format(
            len(self.samples)))

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

    def __getitem__(self, idx):
        row = self.samples[idx]

        context_ids, context_types, context_pad_mask = \
            build_tokens_types_paddings_from_text(row, self.tokenizer,
                self.max_seq_length)

        sample = build_sample(row['doc_id'],
                              context_ids,
                              context_types,
                              context_pad_mask)
        return sample

    @staticmethod
    def process_samples_from_single_path(filename):
        print_rank_0(' > Processing {} ...'.format(filename))
        total = 0

        rows = []


        self.tokenizer = tokenizer
        self.max_seq_length = max_seq_length
        print_rank_0(' > building {} dataset for {}:'.format(self.task_name,
                                                            self.dataset_name))
        # Process the files.
        print_rank_0(datapath)
        self.samples, self.id2text = self.process_samples_from_single_path(
                                        datapath)

        args = get_args()
        if args.sample_rate < 1:  # subsample
            k = int(len(self.samples) * args.sample_rate)
            self.samples = random.sample(self.samples, k)

        print_rank_0('  >> total number of samples: {}'.format(
            len(self.samples)))

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

    def __getitem__(self, idx):
        row = self.samples[idx]

        context_ids, context_types, context_pad_mask = \
            build_tokens_types_paddings_from_text(row, self.tokenizer,
                self.max_seq_length)

        sample = build_sample(row['doc_id'],
                              context_ids,
                              context_types,
                              context_pad_mask)
        return sample

    @staticmethod
    def process_samples_from_single_path(filename):
        print_rank_0(' > Processing {} ...'.format(filename))
        total = 0

        rows = []
        id2text = {}

        with open(filename) as tsvfile:
            reader = csv.reader(tsvfile, delimiter='\t')
            next(reader, None)  # skip the headers
            for row in reader:
                # file format: doc_id, doc_text, title
                doc_id = int(row[0])
                text = row[1]
                title = row[2]

                rows.append({'doc_id': doc_id,
                             'text': text,


        print_rank_0('  >> total number of samples: {}'.format(
            len(self.samples)))

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

    def __getitem__(self, idx):
        row = self.samples[idx]

        context_ids, context_types, context_pad_mask = \
            build_tokens_types_paddings_from_text(row, self.tokenizer,
                self.max_seq_length)

        sample = build_sample(row['doc_id'],
                              context_ids,
                              context_types,
                              context_pad_mask)
        return sample

    @staticmethod
    def process_samples_from_single_path(filename):
        print_rank_0(' > Processing {} ...'.format(filename))
        total = 0

        rows = []
        id2text = {}

        with open(filename) as tsvfile:
            reader = csv.reader(tsvfile, delimiter='\t')
            next(reader, None)  # skip the headers
            for row in reader:
                # file format: doc_id, doc_text, title
                doc_id = int(row[0])
                text = row[1]
                title = row[2]

                rows.append({'doc_id': doc_id,
                             'text': text,
                             'title': title})

                assert doc_id not in id2text
                id2text[doc_id] = (text, title)

                total += 1
                if total % 100000 == 0:
                    print_rank_0('  > processed {} rows so far ...'.format(
                        total))

        print_rank_0(' >> processed {} samples.'.format(len(rows)))
        return rows, id2text


# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.

from PIL import Image, UnidentifiedImageError
import numpy as np
import io
import torch

try:
    from torchvision.transforms import InterpolationMode
    BICUBIC = InterpolationMode.BICUBIC
except ImportError:
    BICUBIC = Image.BICUBIC

from torchvision.transforms import Compose, ToTensor, Normalize, ToPILImage, RandomResizedCrop, Resize

def _convert_image_to_rgb(image):
    return image.convert("RGB")

def _transform(img_h, img_w):
    return Compose([
        ToPILImage(),
        RandomResizedCrop((img_h, img_w), scale=(0.5, 1.0), interpolation=BICUBIC),
        _convert_image_to_rgb,
        ToTensor(),
        Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
    ])

class MultiModalDataset(torch.utils.data.Dataset):

    def __init__(self, name, data_prefix, indexed_dataset,
                 num_samples, seq_length, seed, img_h, img_w):

        self.name = name
        self.indexed_dataset = indexed_dataset
        self.doc_idx = indexed_dataset.get_doc_idx()
        self.visual_transform = _transform(img_h, img_w)

    def __len__(self):
        return self.indexed_dataset.sizes.shape[0]

    def __getitem__(self, idx):
        text_sample, mode = self.indexed_dataset.get(self.doc_idx[idx])
        assert mode == 0
        img_sample, mode = self.indexed_dataset.get(self.doc_idx[idx]+1)
        assert mode == 1
        img_pad = img_sample[0].item()
        xs = img_sample[1:].tobytes(order='C')
        xs = xs[:len(xs)-img_pad]

        img_sample = np.array(Image.open(io.BytesIO(xs)))
        img_sample = self.visual_transform(img_sample).reshape(-1)

        return {'text': np.array(text_sample, dtype=np.int64),
                'img': np.array(img_sample, dtype=np.float32)}


# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.

from PIL import Image, UnidentifiedImageError
import numpy as np
import io
import torch

try:
    from torchvision.transforms import InterpolationMode
    BICUBIC = InterpolationMode.BICUBIC
except ImportError:
    BICUBIC = Image.BICUBIC

from torchvision.transforms import Compose, ToTensor, Normalize, ToPILImage, RandomResizedCrop, Resize

def _convert_image_to_rgb(image):
    return image.convert("RGB")

def _transform(img_h, img_w):
    return Compose([
        ToPILImage(),
        RandomResizedCrop((img_h, img_w), scale=(0.5, 1.0), interpolation=BICUBIC),
        _convert_image_to_rgb,
        ToTensor(),
        Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
    ])

class MultiModalDataset(torch.utils.data.Dataset):

    def __init__(self, name, data_prefix, indexed_dataset,
                 num_samples, seq_length, seed, img_h, img_w):

        self.name = name
        self.indexed_dataset = indexed_dataset
        self.doc_idx = indexed_dataset.get_doc_idx()
        self.visual_transform = _transform(img_h, img_w)


# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.

from PIL import Image, UnidentifiedImageError
import numpy as np
import io
import torch

try:
    from torchvision.transforms import InterpolationMode
    BICUBIC = InterpolationMode.BICUBIC
except ImportError:
    BICUBIC = Image.BICUBIC

from torchvision.transforms import Compose, ToTensor, Normalize, ToPILImage, RandomResizedCrop, Resize

def _convert_image_to_rgb(image):
    return image.convert("RGB")

def _transform(img_h, img_w):
    return Compose([
        ToPILImage(),
        RandomResizedCrop((img_h, img_w), scale=(0.5, 1.0), interpolation=BICUBIC),
        _convert_image_to_rgb,
        ToTensor(),
        Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
    ])

class MultiModalDataset(torch.utils.data.Dataset):

    def __init__(self, name, data_prefix, indexed_dataset,
                 num_samples, seq_length, seed, img_h, img_w):

        self.name = name
        self.indexed_dataset = indexed_dataset
        self.doc_idx = indexed_dataset.get_doc_idx()
        self.visual_transform = _transform(img_h, img_w)

    def __len__(self):
        return self.indexed_dataset.sizes.shape[0]

    def __getitem__(self, idx):
        text_sample, mode = self.indexed_dataset.get(self.doc_idx[idx])
        assert mode == 0
        img_sample, mode = self.indexed_dataset.get(self.doc_idx[idx]+1)
        assert mode == 1
        img_pad = img_sample[0].item()


MultiModalDataset

try:
    from torchvision.transforms import InterpolationMode
    BICUBIC = InterpolationMode.BICUBIC
except ImportError:
    BICUBIC = Image.BICUBIC

from torchvision.transforms import Compose, ToTensor, Normalize, ToPILImage, RandomResizedCrop, Resize

def _convert_image_to_rgb(image):
    return image.convert("RGB")

def _transform(img_h, img_w):
    return Compose([
        ToPILImage(),
        RandomResizedCrop((img_h, img_w), scale=(0.5, 1.0), interpolation=BICUBIC),
        _convert_image_to_rgb,
        ToTensor(),
        Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
    ])

class MultiModalDataset(torch.utils.data.Dataset):

    def __init__(self, name, data_prefix, indexed_dataset,
                 num_samples, seq_length, seed, img_h, img_w):

        self.name = name
        self.indexed_dataset = indexed_dataset
        self.doc_idx = indexed_dataset.get_doc_idx()
        self.visual_transform = _transform(img_h, img_w)

    def __len__(self):
        return self.indexed_dataset.sizes.shape[0]

    def __getitem__(self, idx):
        text_sample, mode = self.indexed_dataset.get(self.doc_idx[idx])
        assert mode == 0
        img_sample, mode = self.indexed_dataset.get(self.doc_idx[idx]+1)
        assert mode == 1
        img_pad = img_sample[0].item()
        xs = img_sample[1:].tobytes(order='C')
        xs = xs[:len(xs)-img_pad]

        img_sample = np.array(Image.open(io.BytesIO(xs)))
        img_sample = self.visual_transform(img_sample).reshape(-1)

        return {'text': np.array(text_sample, dtype=np.int64),
                'img': np.array(img_sample, dtype=np.float32)}


    BICUBIC = InterpolationMode.BICUBIC
except ImportError:
    BICUBIC = Image.BICUBIC

from torchvision.transforms import Compose, ToTensor, Normalize, ToPILImage, RandomResizedCrop, Resize

def _convert_image_to_rgb(image):
    return image.convert("RGB")

def _transform(img_h, img_w):
    return Compose([
        ToPILImage(),
        RandomResizedCrop((img_h, img_w), scale=(0.5, 1.0), interpolation=BICUBIC),
        _convert_image_to_rgb,
        ToTensor(),
        Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
    ])

class MultiModalDataset(torch.utils.data.Dataset):

    def __init__(self, name, data_prefix, indexed_dataset,
                 num_samples, seq_length, seed, img_h, img_w):

        self.name = name
        self.indexed_dataset = indexed_dataset
        self.doc_idx = indexed_dataset.get_doc_idx()
        self.visual_transform = _transform(img_h, img_w)

    def __len__(self):
        return self.indexed_dataset.sizes.shape[0]

    def __getitem__(self, idx):
        text_sample, mode = self.indexed_dataset.get(self.doc_idx[idx])
        assert mode == 0
        img_sample, mode = self.indexed_dataset.get(self.doc_idx[idx]+1)
        assert mode == 1
        img_pad = img_sample[0].item()
        xs = img_sample[1:].tobytes(order='C')
        xs = xs[:len(xs)-img_pad]

        img_sample = np.array(Image.open(io.BytesIO(xs)))
        img_sample = self.visual_transform(img_sample).reshape(-1)

        return {'text': np.array(text_sample, dtype=np.int64),
                'img': np.array(img_sample, dtype=np.float32)}


def _transform(img_h, img_w):
    return Compose([
        ToPILImage(),
        RandomResizedCrop((img_h, img_w), scale=(0.5, 1.0), interpolation=BICUBIC),
        _convert_image_to_rgb,
        ToTensor(),
        Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
    ])

class MultiModalDataset(torch.utils.data.Dataset):

    def __init__(self, name, data_prefix, indexed_dataset,
                 num_samples, seq_length, seed, img_h, img_w):

        self.name = name
        self.indexed_dataset = indexed_dataset
        self.doc_idx = indexed_dataset.get_doc_idx()
        self.visual_transform = _transform(img_h, img_w)

    def __len__(self):
        return self.indexed_dataset.sizes.shape[0]

    def __getitem__(self, idx):
        text_sample, mode = self.indexed_dataset.get(self.doc_idx[idx])
        assert mode == 0
        img_sample, mode = self.indexed_dataset.get(self.doc_idx[idx]+1)
        assert mode == 1
        img_pad = img_sample[0].item()
        xs = img_sample[1:].tobytes(order='C')
        xs = xs[:len(xs)-img_pad]

        img_sample = np.array(Image.open(io.BytesIO(xs)))
        img_sample = self.visual_transform(img_sample).reshape(-1)

        return {'text': np.array(text_sample, dtype=np.int64),
                'img': np.array(img_sample, dtype=np.float32)}


        ToPILImage(),
        RandomResizedCrop((img_h, img_w), scale=(0.5, 1.0), interpolation=BICUBIC),
        _convert_image_to_rgb,
        ToTensor(),
        Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
    ])

class MultiModalDataset(torch.utils.data.Dataset):

    def __init__(self, name, data_prefix, indexed_dataset,
                 num_samples, seq_length, seed, img_h, img_w):

        self.name = name
        self.indexed_dataset = indexed_dataset
        self.doc_idx = indexed_dataset.get_doc_idx()
        self.visual_transform = _transform(img_h, img_w)

    def __len__(self):
        return self.indexed_dataset.sizes.shape[0]

    def __getitem__(self, idx):
        text_sample, mode = self.indexed_dataset.get(self.doc_idx[idx])
        assert mode == 0
        img_sample, mode = self.indexed_dataset.get(self.doc_idx[idx]+1)
        assert mode == 1
        img_pad = img_sample[0].item()
        xs = img_sample[1:].tobytes(order='C')
        xs = xs[:len(xs)-img_pad]

        img_sample = np.array(Image.open(io.BytesIO(xs)))
        img_sample = self.visual_transform(img_sample).reshape(-1)

        return {'text': np.array(text_sample, dtype=np.int64),
                'img': np.array(img_sample, dtype=np.float32)}


# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.

"""Blendable dataset."""

import hashlib
import os
import time

import numpy as np
import torch

from megatron import print_rank_0
from megatron.core import mpu

class BlendableDataset(torch.utils.data.Dataset):

    def __init__(self, datasets, weights, size, *,
                 data_cache_path=None):

        self.datasets = datasets
        num_datasets = len(datasets)
        assert num_datasets == len(weights)

        self.size = size

        # Normalize weights.
        weights = np.array(weights, dtype=np.float64)
        sum_weights = np.sum(weights)
        assert sum_weights > 0.0
        weights /= sum_weights

        # Build indicies.
        def _build_indices():
            start_time = time.time()
            assert num_datasets < 32767
            # Dataset index is a 16-bit integer to alow at least 2^15 datasets.
            # PyTorch isn't happy casting numpy uint16 to a Torch Tensor,
            # so we use int16 although a dataset_index can never be negative.
            dataset_index = np.zeros(self.size, dtype=np.int16)
            dataset_sample_index = np.zeros(self.size, dtype=np.int64)

            from megatron.data import helpers
            helpers.build_blending_indices(dataset_index, dataset_sample_index,
                                           weights, num_datasets, self.size,
                                           torch.distributed.get_rank() == 0)
            print_rank_0('> elapsed time for building blendable dataset indices: '
                         '{:.2f} (sec)'.format(time.time() - start_time))
            return dataset_index, dataset_sample_index

        desc = "Blendable dataset\n\n"
        desc += "Datasets:\n"
        for dataset in datasets:
            desc += dataset.desc + "\n\n"
        desc += f"Weights: {weights}\n"
        desc += f"Size: {size}\n"
        self.desc = desc

        if data_cache_path:
            desc_hash = hashlib.md5(desc.encode('utf-8')).hexdigest()
            desc_path = os.path.join(data_cache_path, desc_hash + ".dsc")
            index_path = os.path.join(data_cache_path, desc_hash + "_index.npy")
            sample_index_path = os.path.join(data_cache_path, desc_hash + "_sample_index.npy")
            cache_hit = os.path.isfile(index_path) and os.path.isfile(sample_index_path)
            cache_success = True
            if torch.distributed.get_rank() == 0 and not cache_hit:
                print(' > WARNING: could not find index map files for blendable'
                      ' dataset, building indices on rank 0 ...', flush=True)
                dataset_index, dataset_sample_index = _build_indices()
                try:
                    os.makedirs(os.path.dirname(index_path), exist_ok=True)
                    with open(desc_path, 'wt') as fd:
                        fd.write(desc)
                        np.save(index_path, dataset_index, allow_pickle=True)
                        np.save(sample_index_path, dataset_sample_index,
                                allow_pickle=True)
                except OSError:
                    print(f'There was an error trying to create the data cache directory ({data_cache_path})')
                    print('or a file in it. This is set with the --data-cache-path argument. Please')
                    print('ensure you have write access to this directory or specify one that you do have')
                    print('write access to.')
                    cache_success = False

            counts = torch.cuda.LongTensor([cache_success])
            torch.distributed.all_reduce(counts, group=mpu.get_data_parallel_group())
            torch.distributed.all_reduce(counts, group=mpu.get_pipeline_model_parallel_group())
            if counts[0].item() != (
                torch.distributed.get_world_size() //
                torch.distributed.get_world_size(group=mpu.get_tensor_model_parallel_group())):
                print_rank_0("Data index creation unsuccessful, exiting.")
                exit()

            # Load on all ranks.
            print_rank_0(f'> loading blendable dataset index: {index_path}')
            self.dataset_index = np.load(index_path, allow_pickle=True, mmap_mode='r')
            assert self.dataset_index.size == self.size

            print_rank_0(f'> loading blendable dataset sample index: {sample_index_path}')
            self.dataset_sample_index = np.load(sample_index_path, allow_pickle=True, mmap_mode='r')
            assert self.dataset_sample_index.size == self.size
        else:
            self.dataset_index, self.dataset_sample_index = _build_indices()

        # Check size
        _ = self.__getitem__(self.size - 1)
        try:
            _ = self.__getitem__(self.size)
            raise RuntimeError('BlendedDataset size is improperly bounded')
        except IndexError:
            pass
        print_rank_0('> size of blendable dataset: '
                     '{} samples'.format(self.size))

    def __len__(self):
        return self.size

    def __getitem__(self, idx):
        dataset_idx = self.dataset_index[idx]
        sample_idx = self.dataset_sample_index[idx]
        return {
            "dataset_idx" : dataset_idx,
            **self.datasets[dataset_idx][sample_idx],
        }


            assert self.dataset_index.size == self.size

            print_rank_0(f'> loading blendable dataset sample index: {sample_index_path}')
            self.dataset_sample_index = np.load(sample_index_path, allow_pickle=True, mmap_mode='r')
            assert self.dataset_sample_index.size == self.size
        else:
            self.dataset_index, self.dataset_sample_index = _build_indices()

        # Check size
        _ = self.__getitem__(self.size - 1)
        try:
            _ = self.__getitem__(self.size)
            raise RuntimeError('BlendedDataset size is improperly bounded')
        except IndexError:
            pass
        print_rank_0('> size of blendable dataset: '
                     '{} samples'.format(self.size))

    def __len__(self):
        return self.size

    def __getitem__(self, idx):
        dataset_idx = self.dataset_index[idx]
        sample_idx = self.dataset_sample_index[idx]
        return {
            "dataset_idx" : dataset_idx,
            **self.datasets[dataset_idx][sample_idx],
        }


            assert self.dataset_sample_index.size == self.size
        else:
            self.dataset_index, self.dataset_sample_index = _build_indices()

        # Check size
        _ = self.__getitem__(self.size - 1)
        try:
            _ = self.__getitem__(self.size)
            raise RuntimeError('BlendedDataset size is improperly bounded')
        except IndexError:
            pass
        print_rank_0('> size of blendable dataset: '
                     '{} samples'.format(self.size))

    def __len__(self):
        return self.size

    def __getitem__(self, idx):
        dataset_idx = self.dataset_index[idx]
        sample_idx = self.dataset_sample_index[idx]
        return {
            "dataset_idx" : dataset_idx,
            **self.datasets[dataset_idx][sample_idx],
        }


class BlendableDataset(torch.utils.data.Dataset):

    def __init__(self, datasets, weights, size, *,
                 data_cache_path=None):

        self.datasets = datasets
        num_datasets = len(datasets)
        assert num_datasets == len(weights)

        self.size = size

        # Normalize weights.
        weights = np.array(weights, dtype=np.float64)
        sum_weights = np.sum(weights)
        assert sum_weights > 0.0
        weights /= sum_weights

        # Build indicies.
        def _build_indices():
            start_time = time.time()
            assert num_datasets < 32767
            # Dataset index is a 16-bit integer to alow at least 2^15 datasets.
            # PyTorch isn't happy casting numpy uint16 to a Torch Tensor,
            # so we use int16 although a dataset_index can never be negative.
            dataset_index = np.zeros(self.size, dtype=np.int16)
            dataset_sample_index = np.zeros(self.size, dtype=np.int64)

            from megatron.data import helpers
            helpers.build_blending_indices(dataset_index, dataset_sample_index,
                                           weights, num_datasets, self.size,
                                           torch.distributed.get_rank() == 0)
            print_rank_0('> elapsed time for building blendable dataset indices: '
                         '{:.2f} (sec)'.format(time.time() - start_time))
            return dataset_index, dataset_sample_index

        desc = "Blendable dataset\n\n"
        desc += "Datasets:\n"
        for dataset in datasets:
            desc += dataset.desc + "\n\n"
        desc += f"Weights: {weights}\n"
        desc += f"Size: {size}\n"
        self.desc = desc

        if data_cache_path:
            desc_hash = hashlib.md5(desc.encode('utf-8')).hexdigest()
            desc_path = os.path.join(data_cache_path, desc_hash + ".dsc")
            index_path = os.path.join(data_cache_path, desc_hash + "_index.npy")
            sample_index_path = os.path.join(data_cache_path, desc_hash + "_sample_index.npy")
            cache_hit = os.path.isfile(index_path) and os.path.isfile(sample_index_path)
            cache_success = True
            if torch.distributed.get_rank() == 0 and not cache_hit:
                print(' > WARNING: could not find index map files for blendable'
                      ' dataset, building indices on rank 0 ...', flush=True)
                dataset_index, dataset_sample_index = _build_indices()