Downloads HuggingFace datasets and concatenates them based on split type.

concatenate_huggingface_datasets.py

"""Downloads HuggingFace datasets and concatenates them based on split type."""

import datasets

from datasets import concatenate_datasets, load_dataset
from datasets.dataset_dict import DatasetDict

# `config_name`s for the `universal_dependencies` dataset
TBIDS = [
    "af_afribooms",
    "ga_idt",
]
SPLITS = ["train", "validation", "test"]

train_files = []
validation_files = []
test_files = []

# Store files based on split type
tmp_files = {split: [] for split in SPLITS}

# Download and sort the files by split type
for tbid in TBIDS:
    for split in SPLITS:
        # first argument is the dataset, second is the config and split determines the split type
        d = load_dataset("universal_dependencies", tbid, split=split)
        tmp_files[split].append(d)

# Create a DatasetDict from concatenated datasets
dd = datasets.DatasetDict(
    {split: concatenate_datasets(files) for split, files in tmp_files.items()}
)

print(f"the concatenated dataset \n {dd}")