Function concatenate_datasets

src/datasets/combine.py:168–232

Concatenate several datasets (sources) into a single dataset. Use axis=0 to concatenate vertically (default), or axis=1 to concatenate horizontally. Note for iterable datasets: * if axis=0, the resulting dataset's `num_shards` is the sum of each dataset's `num_shards`. * if a

(
    dsets: list[DatasetType],
    info: Optional[DatasetInfo] = None,
    split: Optional[NamedSplit] = None,
    axis: int = 0,
)

def concatenate_datasets(
    dsets: list[DatasetType],
    info: Optional[DatasetInfo] = None,
    split: Optional[NamedSplit] = None,
    axis: int = 0,
) -> DatasetType:
    """
    Concatenate several datasets (sources) into a single dataset.

    Use axis=0 to concatenate vertically (default), or axis=1 to concatenate horizontally.

    Note for iterable datasets:

    * if axis=0, the resulting dataset's `num_shards` is the sum of each dataset's `num_shards`.
    * if axis=1, the resulting dataset has one (1) shard to not misalign data.

    Args:
        dsets (`List[datasets.Dataset]` or `List[datasets.IterableDataset]`):
            List of Datasets to concatenate.
        info (`DatasetInfo`, *optional*):
            Dataset information, like description, citation, etc.
        split (`NamedSplit`, *optional*):
            Name of the dataset split.
        axis (`{0, 1}`, defaults to `0`):
            Axis to concatenate over, where `0` means over rows (vertically) and `1` means over columns
            (horizontally).

            <Added version="1.6.0"/>

    Example:

    ```py
    >>> ds3 = concatenate_datasets([ds1, ds2])
    ```
    """

    if not dsets:
        raise ValueError("Unable to concatenate an empty list of datasets.")
    for i, dataset in enumerate(dsets):
        if not isinstance(dataset, (Dataset, IterableDataset)):
            if isinstance(dataset, (DatasetDict, IterableDatasetDict)):
                if not dataset:
                    raise ValueError(
                        f"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} "
                        "is an empty dataset dictionary."
                    )
                raise ValueError(
                    f"Dataset at position {i} has at least one split: {list(dataset)}\n"
                    f"Please pick one to interleave with the other datasets, for example: dataset['{next(iter(dataset))}']"
                )
            raise ValueError(
                f"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} is a {type(dataset).__name__}."
            )
        if i == 0:
            dataset_type, other_type = (
                (Dataset, IterableDataset) if isinstance(dataset, Dataset) else (IterableDataset, Dataset)
            )
        elif not isinstance(dataset, dataset_type):