Skip to content

guidellm.data.utils

resolve_dataset_split(dataset, split=None)

Filter dataset to a specific split or return the first available split.

Source code in src/guidellm/data/utils/dataset.py 
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
def resolve_dataset_split(
    dataset: DatasetDictType,
    split: str | None = None,
) -> DatasetType:
    """
    Filter dataset to a specific split or return the first available split.
    """
    if split is not None and isinstance(dataset, DatasetDict | IterableDatasetDict):
        if split in dataset:
            return dataset[split]

        raise ValueError(f"Requested split '{split}' not found in dataset: {dataset}.")
    elif split is not None:
        raise ValueError(
            f"Requested split '{split}' but dataset has no splits: {dataset}."
        )

    if isinstance(dataset, Dataset | IterableDataset):
        return dataset

    for _, default_splits in DEFAULT_SPLITS.items():
        for default_split in default_splits:
            if default_split in dataset:
                return dataset[default_split]

    return dataset[list(dataset.keys())[0]]