Skip to content

guidellm.data.deserializers.huggingface

HuggingFaceDatasetDeserializer

Bases: DatasetDeserializer

Deserializer for Hugging Face datasets.

Source code in src/guidellm/data/deserializers/huggingface.py
@DatasetDeserializerFactory.register(["huggingface", "hf"])
class HuggingFaceDatasetDeserializer(DatasetDeserializer):
    """Deserializer for Hugging Face datasets."""

    def __call__(
        self,
        config: HuggingFaceDataArgs,
        processor_factory: Callable[[], PreTrainedTokenizerBase],
        random_seed: int,
    ) -> Dataset | IterableDataset | DatasetDict | IterableDatasetDict:
        _ = (processor_factory, random_seed)
        data = config.source

        if isinstance(
            data, Dataset | IterableDataset | DatasetDict | IterableDatasetDict
        ):
            return data

        load_error = None

        if (
            isinstance(data, str | Path)
            and (path := Path(data)).exists()
            and ((path.is_file() and path.suffix == ".py") or path.is_dir())
        ):
            # Handle python script or nested python script in a directory
            try:
                return load_dataset(str(data), **config.load_kwargs)
            except (
                FileNotFoundDatasetsError,
                DatasetNotFoundError,
                DataFilesNotFoundError,
            ) as err:
                load_error = err
            except Exception:  # noqa: BLE001
                # Try loading as a local dataset directory next
                try:
                    return load_from_disk(str(data), **config.load_kwargs)
                except (
                    FileNotFoundDatasetsError,
                    DatasetNotFoundError,
                    DataFilesNotFoundError,
                ) as err2:
                    load_error = err2

        try:
            # Handle dataset identifier from the Hugging Face Hub
            return load_dataset(str(data), **config.load_kwargs)
        except (
            FileNotFoundDatasetsError,
            DatasetNotFoundError,
            DataFilesNotFoundError,
        ) as err:
            load_error = err

        not_supported = DataNotSupportedError(
            "Unsupported data for HuggingFaceDatasetDeserializer, "
            "expected Dataset, IterableDataset, DatasetDict, IterableDatasetDict, "
            "str or Path to a local dataset directory or a local .py dataset script, "
            f"got {data} and HF load error: {load_error}"
        )

        if load_error is not None:
            raise not_supported from load_error
        else:
            raise not_supported