fourm/data/unified_datasets.py

# Copyright 2024 EPFL and Apple Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import copy
import io
import itertools
import os
import re
from functools import partial
from typing import Any, Callable, Dict, Iterable, List, Optional

import braceexpand
import numpy as np
import torch
import webdataset as wds
from PIL import Image
from torch.utils.data import IterableDataset
from torch.utils.data._utils.collate import default_collate
from torchvision import transforms
from webdataset.filters import pipelinefilter, reraise_exception
from webdataset.handlers import warn_and_continue

try:
    # Optionally load huggingface datasets
    from datasets import load_dataset
    from datasets.distributed import split_dataset_by_node
except ImportError:
    print("Huggingface datasets not installed. Please install with `pip install datasets`.")

from fourm.data.masking import TransferMasking, UnifiedMasking
from fourm.data.modality_transforms import (CropSettingsTransform, IdentityTransform,
                                      MaskTransform, UnifiedDataTransform,
                                      get_transform_key)
from fourm.data.multimodal_dataset_folder import MultiModalDatasetFolder
from fourm.utils.dist import get_rank, get_world_size


def build_fm_pretraining_dataset(
        data_path, all_domains, modality_info, modality_transforms, 
        image_augmenter, text_tokenizer, 
        input_tokens_range, target_tokens_range,
        sampling_weights=None):
    """Builds the FourM pre-training dataset based on the given arguments.
    This function should mainly used for smaller datasets (e.g. validation sets), 
    while large training sets should be loaded with build_wds_fm_pretraining_dataloader in webdataset format.
    
    Args:
        data_path: Path to the dataset.
        all_domains: List of all modalities to be used.
        modality_info: Dictionary containing information about the modalities.
        modality_transforms: Dictionary containing the transforms for each modality.
        image_augmenter: Image augmenter.
        text_tokenizer: Text tokenizer (for sequence modalities).
        input_tokens_range: Range of the input token budget.
        target_tokens_range: Range of the target token budget.
        sampling_weights: Sampling weights for the mixture of Dirichlet distributions.

    Returns:
        FourM pre-training dataset as a PyTorch Dataset.
    """

    transform = transforms.Compose([
        UnifiedDataTransform(transforms_dict=modality_transforms, image_augmenter=image_augmenter),
        UnifiedMasking(modality_info=modality_info, text_tokenizer=text_tokenizer,
                       input_tokens_range=input_tokens_range, target_tokens_range=target_tokens_range,
                       sampling_weights=sampling_weights),
         ])

    # Remove vq domains that require a tokenizer
    modalities_without_vq = [mod for mod in all_domains if not modality_info[mod].get("requires_tokenizer", False)]
    # If we are using a pre-tokenized modality, we default to pre-computed crop settings
    if any([modality_info[domain].get("pretokenized", False) for domain in all_domains]):
        modalities_without_vq.append("crop_settings")
        modality_transforms = copy.deepcopy(modality_transforms)
        modality_transforms["crop_settings"] = CropSettingsTransform()

    modality_paths = {mod: modality_info[mod]['path'] for mod in modality_info if modality_info[mod].get('path', None) is not None}
    
    return MultiModalDatasetFolder(root=data_path, modalities=modalities_without_vq, modality_paths=modality_paths,
                                   modality_transforms=modality_transforms, transform=transform)


def build_fm_transfer_dataset(
    data_path, modality_info, transform, modality_transforms, all_domains, 
    load_mask_valid: bool = False, max_samples: Optional[int] = None, 
    pre_shuffle: bool = False, cache: bool = False):
    """Builds the FourM transfer dataset based on the given arguments.
    
    Args:
        data_path: Path to the dataset.
        modality_info: Dictionary containing information about the modalities.
        transform: Transform to be applied to the dataset.
        modality_transforms: Dictionary containing the transforms for each modality.
        all_domains: List of all modalities to be used.
        load_mask_valid: Whether to load the mask_valid "modality".
        max_samples: Maximum number of samples to load.
        pre_shuffle: Whether to shuffle the dataset before loading.
        cache: Whether to cache the dataset in memory.

    Returns:
        FourM transfer dataset as a PyTorch Dataset.
    """

    # Remove vq domains that require a tokenizer
    modalities_without_vq = [mod for mod in all_domains if not modality_info[mod].get("requires_tokenizer", False)]
    # If we are using a pre-tokenized modality, we default to pre-computed crop settings
    if any([modality_info[domain].get("pretokenized", False) for domain in all_domains]):
        modalities_without_vq.append("crop_settings")
        modality_transforms = copy.deepcopy(modality_transforms)
        modality_transforms["crop_settings"] = CropSettingsTransform()

    if load_mask_valid:
        modalities_without_vq.append("mask_valid")
        modality_transforms = copy.deepcopy(modality_transforms)
        modality_transforms["mask_valid"] = MaskTransform()

    modality_paths = {mod: modality_info[mod]['path'] for mod in modality_info if modality_info[mod].get('path', None) is not None}

    return MultiModalDatasetFolder(root=data_path, modalities=modalities_without_vq, modality_paths=modality_paths,
                                   modality_transforms=modality_transforms, transform=transform, max_samples=max_samples, 
                                   pre_shuffle=pre_shuffle, cache=cache)


### Webdatasets (wds) functions

def _keyless_map(data, f, handler=reraise_exception):
    """Map samples without adding __key__."""
    for sample in data:
        try:
            result = f(sample)
        except Exception as exn:
            if handler(exn):
                continue
            else:
                break
        if result is None:
            continue
        yield result

map = pipelinefilter(_keyless_map)

def check_dots(s):
    if '.gz' in s:
        return s.count('.') == 2
    return s.count('.') == 1

def remove_ext_with_gz(s):
    if s.endswith('.gz'):
        s = s.replace(".gz", "")
    return os.path.splitext(s)[0]

def wds_decoder(key, value):
    if key == "png" or key.endswith(".png"):
        img = Image.open(io.BytesIO(value))
        return img
    elif key == "jpg" or key.endswith(".jpg"):
        img = Image.open(io.BytesIO(value))
        return img
    elif key == "jpeg" or key.endswith(".jpeg"):
        img = Image.open(io.BytesIO(value))
        return img
    elif key == 'npy' or key.endswith("npy"):
        content = np.load(io.BytesIO(value), allow_pickle=True)
        # try:
        #     content = np.load(io.BytesIO(value))
        # except:
        #     content = np.load(io.BytesIO(value), allow_pickle=True)
        return content
    elif key == "jpx" or key.endswith('.jpx'):
        img = Image.open(io.BytesIO(value))
        return img
    elif 'output' in key:
        return int(value)
    else:
        # If not an image, use the basic handlers (.txt, .json, .pickle, .npz, ...)
        # See https://github.com/webdataset/webdataset/blob/main/webdataset/autodecode.py
        return None

def repeat_fn(src, n_repeats=5):
    """
    Repeat each sample n_repeats times.
    E.g. A B C ... repeated 3 times becomes A A A B B B C C C ...
    Depending on the downstream application, a shuffle should be added after this.
    """
    for sample in src:
        for _ in range(n_repeats):
            yield sample
            
def remove_extensions(sample):
    """
    In webdatasets, we identify the type of a given modality by adding an extension
    in the form f"{modality_name}.{modality_extension}", e.g. "rgb.jpg" or "caption.json".
    This function removes them and returns a dictionary of {f"{modality_name}": modality}.
    """
    return {remove_ext_with_gz(k): v for k, v in sample.items()}

def filter_metadata(sample, metadata=['__key__', '__url__', 'file_name', 'class_name', 'class_idx']):
    """ Filters out non-modality entries specified in metadata when loading tar files with webdatasets. """
    return {k: v for k, v in sample.items() if k not in metadata}

def apply_modality_transforms(sample, modality_transforms):
    """ Applies a dictionary of modality-specific transforms to a dictionary of modalities. """
    return {k: (modality_transforms[get_transform_key(k)](v) if k in modality_transforms else v) for k, v in sample.items() }

def tok_to_int64(sample):
    """
    Pre-computed tokens are saved as int16, but we need them as int64 instead.
    """
    return {k: (v.astype('int64') if 'tok_' in k else v) for k, v in sample.items()}

def rename_modalities(sample, modality_paths):
    """
    Renames modalities to their corresponding names in modality_paths.
    """
    return {out_path: sample[loaded_path] for out_path, loaded_path in modality_paths.items()}

def extract_modality_names(s):
    # Regular expression pattern to match anything enclosed in '{' and '}', and comma separated
    pattern = r'\{([^}]*)\}'
    match = re.search(pattern, s)
    return match.group(1).split(',') if match else []

def identity(sample):
    """ Identity function that does nothing. """
    return sample

def multi_tarfile_samples(src_iter: Iterable[Dict[str, Any]], 
                          modality_name_map: Dict[str, str] = None, 
                          handler: Callable[[Exception], bool] = warn_and_continue):
    """Webdataset does not support splitting up shards by modality, so we need to do this manually.
    Usually, we would need to save all modalities in the same tar file, e.g. shard_root_train/{00000..12345}.tar, 
    where each shard contains 1000 samples and each sample contains all modalities.
    This is not flexible when adding new modalities, so we instead save each modality in a separate tar file,
    e.g. shard_root_train_rgb/{00000..12345}.tar, shard_root_train_caption/{00000..12345}.tar, etc., where each shard contains
    again 1000 samples, but each sample contains only one modality. All samples in all shards have to be aligned.

    This function takes an iterator over shard URLs, where we use brace expansion to specify multiple tar files per modality.
    E.g. shard_root_train_[rgb,caption]/00123.tar will be expanded to shard_root_train_rgb/00123.tar and shard_root_train_caption/00123.tar,
    and the samples from these two tar files will be combined into a single sample.

    Args:
        src_iter: Iterator over shards that *already brace expanded the shard numbers*, 
            e.g. {'url': 'shard_root_train_[rgb,caption]/00000.tar'}, {'url': 'shard_root_train_[rgb,caption]/00001.tar'}, ...
            This function will also work when no square braces for multiple modalities are used, e.g. {'url': 'shard_root_train/00000.tar'}, ...
            It can be a drop-in replacement for wds.tarfile_samples.
        modality_name_map: Optional dictionary specifying a mapping from modality folder names to arbitrary other names.
        handler: Function that handles exceptions. If it returns True, the shard is skipped. If it returns False, the function exits.

    Yields:
        Dictionary of aligned samples from all modalities.
    """
    for src in src_iter:
        
        # Multi tar file URLs use brace expansion with square braces
        multi_tar_urls = src['url'].translate(str.maketrans('[]', '{}'))
        modality_names = extract_modality_names(multi_tar_urls)
        if len(modality_names) == 0:
            # Case where multi-modal braceexpand is not used, e.g. shard_dir/shard00000.tar
            modality_names = [None]
            multi_tar_urls = [multi_tar_urls]
        elif len(modality_names) == 1:
            # Brace expand doesn't work with a single entry, e.g. shard_dir/[foo]/shard00000.tar
            multi_tar_urls = [multi_tar_urls.replace("{", "").replace("}", "")]
        else:
            # Remaining cases where multiple modalities are specified, e.g. shard_dir/[foo,bar]/shard00000.tar
            multi_tar_urls = list(braceexpand.braceexpand(multi_tar_urls))

        # Create tar iterators for shards of all modalities
        tar_iters = [wds.tarfile_samples([{'url': tar_url}]) for tar_url in multi_tar_urls]
        
        try:
            # Loop over these iterators in parallel and combine the tar files from different modalities
            for multi_tar_files in zip(*tar_iters):
                
                merged_dict = {}
                merged_dict['__key__'] = multi_tar_files[0]['__key__']
                merged_dict['__url__'] = src['url']
                
                for modality_name, modality_dict in zip(modality_names, multi_tar_files):
                    _key = modality_dict.pop('__key__')
                    _url = modality_dict.pop('__url__')

                    if _key != merged_dict['__key__']:
                        raise ValueError(f"Divergence detected! Trying to merge keys {_key} of {modality_name} and {merged_dict['__key__']} of merged_dict with modalities {merged_dict.keys()}.")
                        
                    tar_is_multimodal = len(modality_dict) > 1
                    for k, v in modality_dict.items():
                        if tar_is_multimodal or check_dots(k) or modality_name is None:
                            # We don't change the keys in the following cases:
                            # 1. The shard contains multiple modalities. Then they *have* to follow the idx.modality_id.ext convention
                            # 2. If any key contains a dot, this means it already has the idx.modality_id.ext format (idx. is already removed at this stage)
                            # 3. If the modality name is None, no modality folder was specified (see beginning of function)
                            merged_dict[k] = v
                        else:
                            mapped_name = modality_name if modality_name_map is None else modality_name_map.get(modality_name, modality_name)
                            merged_dict[f'{mapped_name}.{k}'] = v

                yield merged_dict

        except Exception as e:
            print(e)
            print(f"Exception occurred while processing {src['url']}.")
            if handler(e):
                print('Skipping shard...')
                continue
            else:
                break

def build_wds_fm_pretraining_dataloader(
        data_path, all_domains, modality_info, modality_transforms, image_augmenter, 
        text_tokenizer, input_tokens_range, target_tokens_range,
        num_gpus, num_workers, batch_size, epoch_size, sampling_weights=None, modality_name_map=None,
        shuffle_buffer_load=1000, shuffle_buffer_repeat=5000, n_repeats=5):
    """Builds the WebDataset FourM pre-training dataloader based on the given arguments.
    
    Args:
        data_path: Path to the dataset.
        all_domains: List of all modalities to be used.
        modality_info: Dictionary containing information about the modalities.
        modality_transforms: Dictionary containing the transforms for each modality.
        image_augmenter: Image augmenter.
        text_tokenizer: Text tokenizer (for sequence modalities).
        input_tokens_range: Range of the input token budget.
        target_tokens_range: Range of the target token budget.
        num_gpus: Number of GPUs.
        num_workers: Number of workers.
        batch_size: Batch size.
        epoch_size: Number of samples per "epoch". (Here, epoch refers to an interrupted training loop without evaluation or checkpointing).
        sampling_weights: Sampling weights for the mixture of Dirichlet distributions.
        modality_name_map: Optional dictionary specifying a mapping from modality folder names to arbitrary other names.
        shuffle_buffer_load: Shuffle buffer size when loading samples from tar files (first shuffle).
        shuffle_buffer_repeat: Shuffle buffer size after repeating samples (second shuffle).
        n_repeats: Number of times to repeat each sample.

    Returns:
        FourM pre-training dataloader as a WebDataset DataLoader.
    """

    modality_paths = {mod: modality_info[mod].get('path', None) or mod for mod in modality_info}

    # Remove vq domains that require a tokenizer
    modalities_without_vq = [mod for mod in all_domains if not modality_info[mod].get("requires_tokenizer", False)]
    # If we are using a pre-tokenized modality, we default to pre-computed crop settings
    if any([modality_info[domain].get("pretokenized", False) for domain in all_domains]):
        modalities_without_vq.append("crop_settings")
        modality_transforms = copy.deepcopy(modality_transforms)
        modality_transforms["crop_settings"] = CropSettingsTransform()
        modality_paths["crop_settings"] = "crop_settings"

    # Webdatasets always adds __key__ to the dictionary, so we add a transform that does nothing with it
    modality_transforms["__key__"] = IdentityTransform()

    transform = transforms.Compose([
        UnifiedDataTransform(transforms_dict=modality_transforms, image_augmenter=image_augmenter),
        UnifiedMasking(modality_info=modality_info, text_tokenizer=text_tokenizer,
                       input_tokens_range=input_tokens_range, target_tokens_range=target_tokens_range,
                       sampling_weights=sampling_weights)
    ])
    
    datapipe = wds.DataPipeline(
        # Infinitely sample shards from the shard list with replacement. Each worker is seeded independently.
        wds.ResampledShards(data_path),
        partial(multi_tarfile_samples, modality_name_map=modality_name_map), # Extract individual samples from single or multi-modal tar files
        wds.shuffle(shuffle_buffer_load), # Shuffle with a buffer of given size
        wds.decode(wds_decoder), # Decode from bytes to PIL images, numpy arrays, etc.
        wds.filters.compose(partial(repeat_fn, n_repeats=n_repeats)), # Repeats each sample n times -> A A A B B B C C C ...
        wds.shuffle(shuffle_buffer_repeat), # Shuffle again with a buffer of given size
        wds.map(remove_extensions), # Remove "file extensions" from dictionary keys
        map(filter_metadata), # Remove non-task keys
        map(tok_to_int64), # Convert pre-computed tokens to int64
        map(partial(rename_modalities, modality_paths=modality_paths)), # Rename modalities to their corresponding names in modality_paths
        map(transform), # Apply data augmentation and masking
        wds.batched(batch_size, collation_fn=default_collate, partial=False)
            if batch_size is not None else map(identity), # Batching
    )

    if epoch_size is not None:
        batch_size_iter = batch_size if batch_size is not None else 1
        datapipe = datapipe.with_epoch(epoch_size // (num_gpus * num_workers * batch_size_iter)) # Pre-define iterator length
    
    if batch_size is not None:
        # Perform multi-threaded dataloading
        return wds.WebLoader(datapipe, num_workers=num_workers, batch_size=None)
    else:
        return datapipe


def build_wds_divae_dataloader(
    data_path, modality_info, modality_transforms, image_augmenter,  
    num_gpus, num_workers, batch_size, epoch_size, shuffle_buffer_load=1000, 
    shuffle_buffer_repeat=5000, n_repeats=1):

    modality_paths = {mod: modality_info[mod].get('path', None) or mod for mod in modality_info}

    # Webdatasets always adds __key__ to the dictionary, so we add a transform that does nothing with it
    modality_transforms["__key__"] = IdentityTransform()

    transform = UnifiedDataTransform(transforms_dict=modality_transforms, image_augmenter=image_augmenter)

    datapipe = wds.DataPipeline(
        # Infinitely sample shards from the shard list with replacement. Each worker is seeded independently.
        wds.ResampledShards(data_path),
        multi_tarfile_samples, # Extract individual samples from single or multi-modal tar files
        wds.shuffle(shuffle_buffer_load), # Shuffle with a buffer of given size
        wds.decode(wds_decoder), # Decode from bytes to PIL images, numpy arrays, etc.
        wds.filters.compose(partial(repeat_fn, n_repeats=n_repeats)), # Repeats each sample n times -> A A A B B B C C C ...
        wds.shuffle(shuffle_buffer_repeat), # Shuffle again with a buffer of given size
        map(remove_extensions), # Remove "file extensions" from dictionary keys
        map(filter_metadata), # Remove non-task keys
        map(tok_to_int64), # Convert pre-computed tokens to int64
        map(partial(rename_modalities, modality_paths=modality_paths)), # Rename modalities to their corresponding names in modality_paths
        map(transform), # Apply data augmentation and masking
        wds.batched(batch_size, collation_fn=default_collate, partial=False)
            if batch_size is not None else map(identity), # Batching
    )

    if epoch_size is not None:
        batch_size_iter = batch_size if batch_size is not None else 1
        datapipe = datapipe.with_epoch(epoch_size // (num_gpus * num_workers * batch_size_iter)) # Pre-define iterator length
    
    if batch_size is not None:
        # Perform multi-threaded dataloading
        return wds.WebLoader(datapipe, num_workers=num_workers, batch_size=None)
    else:
        return datapipe


### Huggingface datasets functions

def text_to_caption(sample):
    """ Rename "text" to "caption". """
    return {'caption': sample['text']}


def build_huggingface_pretraining_dataloader(
        data_path, all_domains, modality_info, modality_transforms, image_augmenter, 
        text_tokenizer, input_tokens_range, target_tokens_range,
        num_gpus, num_workers, batch_size, epoch_size, split,
        streaming=True, rename_text_to_caption=True, shuffle_buffer_load=10_000, shuffle_seed=0):

    # Load huggingface dataset and split samples across workers. Shuffle samples in each worker
    dataset = load_dataset(data_path, split=split, streaming=streaming)
    dataset = split_dataset_by_node(dataset, rank=get_rank(), world_size=get_world_size())
    dataset = dataset.shuffle(seed=shuffle_seed, buffer_size=shuffle_buffer_load)

    modality_info = {mod: modality_info[mod] for mod in modality_info if mod in all_domains}
    
    transform = transforms.Compose([
        UnifiedDataTransform(transforms_dict=modality_transforms, image_augmenter=image_augmenter),
        UnifiedMasking(modality_info=modality_info, text_tokenizer=text_tokenizer,
                       input_tokens_range=input_tokens_range, target_tokens_range=target_tokens_range)
    ])
    
    datapipe = wds.DataPipeline(
        dataset,
        map(text_to_caption) if rename_text_to_caption else map(identity), # Rename "text" to "caption"
        map(filter_metadata), # Remove non-task keys
        map(transform), # Apply data augmentation and masking
        wds.batched(batch_size, collation_fn=default_collate, partial=False)
            if batch_size is not None else map(identity), # Batching
    )

    datapipe.n_shards = dataset.n_shards
    num_workers = min(num_workers, dataset.n_shards)

    if epoch_size is not None:
        batch_size_iter = batch_size if batch_size is not None else 1
        datapipe = datapipe.with_epoch(epoch_size // (num_gpus * num_workers * batch_size_iter)) # Pre-define iterator length

    if batch_size is not None:
        # Perform multi-threaded dataloading
        return wds.WebLoader(datapipe, num_workers=num_workers, batch_size=None)
    else:
        return datapipe


### Multi-dataset loading utils
def make_empty_mod_dict(modality_info):
    empty_mod_dicts = {}

    for mod_name, mod_info in modality_info.items():
        empty_mod = {}

        # Tensor
        if 'num_channels' in mod_info and 'input_size' in mod_info:
            # Handle image-like modalities
            max_tokens = mod_info['max_tokens']
            empty_mod['tensor'] = torch.zeros((mod_info['num_channels'], mod_info['input_size'], mod_info['input_size']), dtype=torch.float32)
        elif mod_name == 't5_caption':
            # Handle T5 embedding
            max_tokens = mod_info['max_tokens']
            orig_emb_dim = mod_info['encoder_embedding']().orig_emb_dim
            empty_mod['tensor'] = torch.zeros((max_tokens, orig_emb_dim), dtype=torch.float32)
        elif mod_info['type'] in ['seq', 'seq_emb', 'seq_token']:
            # Handle all other discrete sequence modalities
            max_tokens = (mod_info['max_tokens'] + 1) * 2
            empty_mod['tensor'] = torch.zeros((max_tokens), dtype=torch.int32)
        else:
            max_tokens = mod_info['max_tokens']
            empty_mod['tensor'] = torch.zeros((max_tokens), dtype=torch.int32)
            
        # Input and target masks
        empty_mod['input_mask'] = torch.ones((max_tokens), dtype=torch.bool)
        empty_mod['target_mask'] = torch.ones((max_tokens), dtype=torch.bool)

        # Decoder attention mask
        empty_mod['decoder_attention_mask'] = torch.zeros((max_tokens), dtype=torch.int32)
        
        empty_mod_dicts[mod_name] = empty_mod
        
    return empty_mod_dicts


class MixtureDataset(IterableDataset):
    def __init__(self, data_iters, weights, modality_info):
        self.orig_data_iters = data_iters
        self.data_iters = [iter(data_iter) for data_iter in data_iters]  # Create initial iterators
        self.sampling_probs = np.array(weights) / sum(weights)
        self.modality_info = modality_info

    def reset_iterator(self, idx):
        """ Reset the iterator when exhausted. """
        self.data_iters[idx] = iter(self.orig_data_iters[idx])

    def __iter__(self):
        while True:
            dataset_idx = np.random.choice(len(self.sampling_probs), p=self.sampling_probs)
            try:
                data = next(self.data_iters[dataset_idx])
            except StopIteration:  # If the iterator is exhausted
                self.reset_iterator(dataset_idx)  # Reset it
                data = next(self.data_iters[dataset_idx])

            mod_dict = make_empty_mod_dict(self.modality_info)
            mod_dict.update(data)
            yield mod_dict


def build_mixture_dataloader(data_iters, weights, modality_info, batch_size, num_workers, epoch_size, num_gpus):
    mixture_pipe = wds.DataPipeline(
        MixtureDataset(data_iters, weights, modality_info),
        wds.batched(batch_size, collation_fn=default_collate, partial=False),
    ).with_epoch(epoch_size // (num_gpus * num_workers * batch_size)) # Pre-define iterator length
    
    mixture_loader = wds.WebLoader(mixture_pipe, num_workers=num_workers, batch_size=None)
    
    return mixture_loader