# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/hf.datasets.utils.ipynb.

# %% auto 0
__all__ = ['log', 'load_datasets', 'concatenate_dataset_dicts', 'chunk_random_dataset']

# %% ../../../nbs/hf.datasets.utils.ipynb 3
from datasets import concatenate_datasets, DatasetDict, load_dataset, Dataset

from ...logging import get_logger
from ...utils import chunk_random

log = get_logger(__name__)

# %% ../../../nbs/hf.datasets.utils.ipynb 4
def load_datasets(dataset_kwargs_list: list[dict]) -> Dataset:
    datasets = []
    for dataset_kwargs in dataset_kwargs_list:
        dataset = load_dataset(**dataset_kwargs)
        datasets.append(dataset)
    return concatenate_datasets(datasets)

# %% ../../../nbs/hf.datasets.utils.ipynb 5
def concatenate_dataset_dicts(dataset_dicts: list[DatasetDict]):
    """
    Concatenate multiple `DatasetDict` objects into a single `DatasetDict`.
    """
    splits = {key for dd in dataset_dicts for key in dd.keys()}
    output_dsd = {}
    for split in splits:
        dataset_list = []
        for dd in dataset_dicts:
            if split in dd.keys():
                dataset_list.append(dd[split])
        output_dsd[split] = concatenate_datasets(dataset_list)
    return DatasetDict(output_dsd)


# %% ../../../nbs/hf.datasets.utils.ipynb 7
def chunk_random_dataset(ds, min_chunk=1, max_chunk=3):
    for indices in chunk_random(range(len(ds)), min_chunk, max_chunk):
        yield ds.select(indices).to_list()
