guidellm 0.4.0a21__py3-none-any.whl → 0.4.0a169__py3-none-any.whl

guidellm/data/__init__.py

@@ -1,4 +1,28 @@
-"""
-Required for python < 3.12
-https://docs.python.org/3/library/importlib.resources.html#importlib.resources.files
-"""
+from .collators import GenerativeRequestCollator
+from .deserializers import (
+    DataNotSupportedError,
+    DatasetDeserializer,
+    DatasetDeserializerFactory,
+)
+from .loaders import DataLoader, DatasetsIterator
+from .preprocessors import (
+    DataDependentPreprocessor,
+    DatasetPreprocessor,
+    PreprocessorRegistry,
+)
+from .processor import ProcessorFactory
+from .schemas import GenerativeDatasetColumnType
+
+__all__ = [
+    "DataDependentPreprocessor",
+    "DataLoader",
+    "DataNotSupportedError",
+    "DatasetDeserializer",
+    "DatasetDeserializerFactory",
+    "DatasetPreprocessor",
+    "DatasetsIterator",
+    "GenerativeDatasetColumnType",
+    "GenerativeRequestCollator",
+    "PreprocessorRegistry",
+    "ProcessorFactory",
+]

guidellm/data/collators.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from guidellm.schemas import GenerationRequest
+
+__all__ = ["GenerativeRequestCollator"]
+
+
+class GenerativeRequestCollator:
+    def __call__(self, batch: list) -> GenerationRequest:
+        if len(batch) != 1:
+            raise NotImplementedError(
+                f"Batch size greater than 1 is not currently supported. "
+                f"Got batch size: {len(batch)}"
+            )
+
+        return batch[0]

guidellm/data/deserializers/__init__.py

@@ -0,0 +1,53 @@
+from .deserializer import (
+    DataNotSupportedError,
+    DatasetDeserializer,
+    DatasetDeserializerFactory,
+)
+from .file import (
+    ArrowFileDatasetDeserializer,
+    CSVFileDatasetDeserializer,
+    DBFileDatasetDeserializer,
+    HDF5FileDatasetDeserializer,
+    JSONFileDatasetDeserializer,
+    ParquetFileDatasetDeserializer,
+    TarFileDatasetDeserializer,
+    TextFileDatasetDeserializer,
+)
+from .huggingface import HuggingFaceDatasetDeserializer
+from .memory import (
+    InMemoryCsvDatasetDeserializer,
+    InMemoryDictDatasetDeserializer,
+    InMemoryDictListDatasetDeserializer,
+    InMemoryItemListDatasetDeserializer,
+    InMemoryJsonStrDatasetDeserializer,
+)
+from .synthetic import (
+    SyntheticTextDatasetConfig,
+    SyntheticTextDatasetDeserializer,
+    SyntheticTextGenerator,
+    SyntheticTextPrefixBucketConfig,
+)
+
+__all__ = [
+    "ArrowFileDatasetDeserializer",
+    "CSVFileDatasetDeserializer",
+    "DBFileDatasetDeserializer",
+    "DataNotSupportedError",
+    "DatasetDeserializer",
+    "DatasetDeserializerFactory",
+    "HDF5FileDatasetDeserializer",
+    "HuggingFaceDatasetDeserializer",
+    "InMemoryCsvDatasetDeserializer",
+    "InMemoryDictDatasetDeserializer",
+    "InMemoryDictListDatasetDeserializer",
+    "InMemoryItemListDatasetDeserializer",
+    "InMemoryJsonStrDatasetDeserializer",
+    "JSONFileDatasetDeserializer",
+    "ParquetFileDatasetDeserializer",
+    "SyntheticTextDatasetConfig",
+    "SyntheticTextDatasetDeserializer",
+    "SyntheticTextGenerator",
+    "SyntheticTextPrefixBucketConfig",
+    "TarFileDatasetDeserializer",
+    "TextFileDatasetDeserializer",
+]

guidellm/data/deserializers/deserializer.py

@@ -0,0 +1,144 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any, Protocol, Union, runtime_checkable
+
+from datasets import Dataset, DatasetDict, IterableDataset, IterableDatasetDict
+from transformers import PreTrainedTokenizerBase
+
+from guidellm.data.utils import resolve_dataset_split
+from guidellm.utils import RegistryMixin
+
+__all__ = [
+    "DataNotSupportedError",
+    "DatasetDeserializer",
+    "DatasetDeserializerFactory",
+]
+
+
+class DataNotSupportedError(Exception):
+    """Exception raised when data format is not supported by deserializer."""
+
+
+@runtime_checkable
+class DatasetDeserializer(Protocol):
+    def __call__(
+        self,
+        data: Any,
+        processor_factory: Callable[[], PreTrainedTokenizerBase],
+        random_seed: int,
+        **data_kwargs: dict[str, Any],
+    ) -> Dataset | IterableDataset | DatasetDict | IterableDatasetDict: ...
+
+
+class DatasetDeserializerFactory(
+    RegistryMixin[Union["type[DatasetDeserializer]", DatasetDeserializer]],
+):
+    @classmethod
+    def deserialize(
+        cls,
+        data: Any,
+        processor_factory: Callable[[], PreTrainedTokenizerBase],
+        random_seed: int = 42,
+        type_: str | None = None,
+        resolve_split: bool = True,
+        select_columns: list[str] | None = None,
+        remove_columns: list[str] | None = None,
+        **data_kwargs: dict[str, Any],
+    ) -> Dataset | IterableDataset:
+        dataset: Dataset
+
+        if type_ is None:
+            dataset = cls._deserialize_with_registered_deserializers(
+                data, processor_factory, random_seed, **data_kwargs
+            )
+
+        else:
+            dataset = cls._deserialize_with_specified_deserializer(
+                data, type_, processor_factory, random_seed, **data_kwargs
+            )
+
+        if resolve_split:
+            dataset = resolve_dataset_split(dataset)
+
+        if select_columns is not None or remove_columns is not None:
+            column_names = dataset.column_names or list(next(iter(dataset)).keys())
+            if select_columns is not None:
+                remove_columns = [
+                    col for col in column_names if col not in select_columns
+                ]
+
+            dataset = dataset.remove_columns(remove_columns)
+
+        return dataset
+
+    @classmethod
+    def _deserialize_with_registered_deserializers(
+        cls,
+        data: Any,
+        processor_factory: Callable[[], PreTrainedTokenizerBase],
+        random_seed: int = 42,
+        **data_kwargs: dict[str, Any],
+    ) -> Dataset:
+        if cls.registry is None:
+            raise RuntimeError("registry is None; cannot deserialize dataset")
+        dataset: Dataset | None = None
+
+        errors: dict[str, Exception] = {}
+        # Note: There is no priority order for the deserializers, so all deserializers
+        #  must be mutually exclusive to ensure deterministic behavior.
+        for _name, deserializer in cls.registry.items():
+            deserializer_fn: DatasetDeserializer = (
+                deserializer() if isinstance(deserializer, type) else deserializer
+            )
+
+            try:
+                dataset = deserializer_fn(
+                    data=data,
+                    processor_factory=processor_factory,
+                    random_seed=random_seed,
+                    **data_kwargs,
+                )
+            except Exception as e:  # noqa: BLE001 # The exceptions are saved.
+                errors[_name] = e
+
+            if dataset is not None:
+                return dataset  # Success
+
+        if len(errors) > 0:
+            err_msgs = ""
+            def sort_key(item):
+                return (isinstance(item[1], DataNotSupportedError), item[0])
+            for key, err in sorted(errors.items(), key=sort_key):
+                err_msgs += f"\n  - Deserializer '{key}': ({type(err).__name__}) {err}"
+            raise ValueError(
+                "Data deserialization failed, likely because the input doesn't "
+                f"match any of the input formats. See the {len(errors)} error(s) that "
+                f"occurred while attempting to deserialize the data {data}:{err_msgs}"
+            )
+        return dataset
+
+    @classmethod
+    def _deserialize_with_specified_deserializer(
+        cls,
+        data: Any,
+        type_: str,
+        processor_factory: Callable[[], PreTrainedTokenizerBase],
+        random_seed: int = 42,
+        **data_kwargs: dict[str, Any],
+    ) -> Dataset:
+        deserializer_from_type = cls.get_registered_object(type_)
+        if deserializer_from_type is None:
+            raise ValueError(f"Deserializer type '{type_}' is not registered.")
+        if isinstance(deserializer_from_type, type):
+            deserializer_fn = deserializer_from_type()
+        else:
+            deserializer_fn = deserializer_from_type
+
+        return deserializer_fn(
+            data=data,
+            processor_factory=processor_factory,
+            random_seed=random_seed,
+            **data_kwargs,
+        )
+

guidellm/data/deserializers/file.py

@@ -0,0 +1,222 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any
+
+import pandas as pd
+from datasets import Dataset, load_dataset
+from transformers import PreTrainedTokenizerBase
+
+from guidellm.data.deserializers.deserializer import (
+    DataNotSupportedError,
+    DatasetDeserializer,
+    DatasetDeserializerFactory,
+)
+
+__all__ = [
+    "ArrowFileDatasetDeserializer",
+    "CSVFileDatasetDeserializer",
+    "DBFileDatasetDeserializer",
+    "HDF5FileDatasetDeserializer",
+    "JSONFileDatasetDeserializer",
+    "ParquetFileDatasetDeserializer",
+    "TarFileDatasetDeserializer",
+    "TextFileDatasetDeserializer",
+]
+
+
+@DatasetDeserializerFactory.register("text_file")
+class TextFileDatasetDeserializer(DatasetDeserializer):
+    def __call__(
+        self,
+        data: Any,
+        processor_factory: Callable[[], PreTrainedTokenizerBase],
+        random_seed: int,
+        **data_kwargs: dict[str, Any],
+    ) -> Dataset:
+        _ = (processor_factory, random_seed)  # Ignore unused args format errors
+
+        if (
+            not isinstance(data, str | Path)
+            or not (path := Path(data)).exists()
+            or not path.is_file()
+            or path.suffix.lower() not in {".txt", ".text"}
+        ):
+            raise DataNotSupportedError(
+                "Unsupported data for TextFileDatasetDeserializer, "
+                f"expected str or Path to a local .txt or .text file, got {data}"
+            )
+
+        with path.open() as file:
+            lines = file.readlines()
+
+        return Dataset.from_dict({"text": lines}, **data_kwargs)
+
+
+@DatasetDeserializerFactory.register("csv_file")
+class CSVFileDatasetDeserializer(DatasetDeserializer):
+    def __call__(
+        self,
+        data: Any,
+        processor_factory: Callable[[], PreTrainedTokenizerBase],
+        random_seed: int,
+        **data_kwargs: dict[str, Any],
+    ) -> Dataset:
+        _ = (processor_factory, random_seed)
+        if (
+            not isinstance(data, str | Path)
+            or not (path := Path(data)).exists()
+            or not path.is_file()
+            or path.suffix.lower() != ".csv"
+        ):
+            raise DataNotSupportedError(
+                "Unsupported data for CSVFileDatasetDeserializer, "
+                f"expected str or Path to a local .csv file, got {data}"
+            )
+
+        return load_dataset("csv", data_files=str(path), **data_kwargs)
+
+
+@DatasetDeserializerFactory.register("json_file")
+class JSONFileDatasetDeserializer(DatasetDeserializer):
+    def __call__(
+        self,
+        data: Any,
+        processor_factory: Callable[[], PreTrainedTokenizerBase],
+        random_seed: int,
+        **data_kwargs: dict[str, Any],
+    ) -> Dataset:
+        _ = (processor_factory, random_seed)
+        if (
+            not isinstance(data, str | Path)
+            or not (path := Path(data)).exists()
+            or not path.is_file()
+            or path.suffix.lower() not in {".json", ".jsonl"}
+        ):
+            raise DataNotSupportedError(
+                f"Unsupported data for JSONFileDatasetDeserializer, "
+                f"expected str or Path to a local .json or .jsonl file, got {data}"
+            )
+
+        return load_dataset("json", data_files=str(path), **data_kwargs)
+
+
+@DatasetDeserializerFactory.register("parquet_file")
+class ParquetFileDatasetDeserializer(DatasetDeserializer):
+    def __call__(
+        self,
+        data: Any,
+        processor_factory: Callable[[], PreTrainedTokenizerBase],
+        random_seed: int,
+        **data_kwargs: dict[str, Any],
+    ) -> Dataset:
+        _ = (processor_factory, random_seed)
+        if (
+            not isinstance(data, str | Path)
+            or not (path := Path(data)).exists()
+            or not path.is_file()
+            or path.suffix.lower() != ".parquet"
+        ):
+            raise DataNotSupportedError(
+                f"Unsupported data for ParquetFileDatasetDeserializer, "
+                f"expected str or Path to a local .parquet file, got {data}"
+            )
+
+        return load_dataset("parquet", data_files=str(path), **data_kwargs)
+
+
+@DatasetDeserializerFactory.register("arrow_file")
+class ArrowFileDatasetDeserializer(DatasetDeserializer):
+    def __call__(
+        self,
+        data: Any,
+        processor_factory: Callable[[], PreTrainedTokenizerBase],
+        random_seed: int,
+        **data_kwargs: dict[str, Any],
+    ) -> Dataset:
+        _ = (processor_factory, random_seed)
+        if (
+            not isinstance(data, str | Path)
+            or not (path := Path(data)).exists()
+            or not path.is_file()
+            or path.suffix.lower() != ".arrow"
+        ):
+            raise DataNotSupportedError(
+                f"Unsupported data for ArrowFileDatasetDeserializer, "
+                f"expected str or Path to a local .arrow file, got {data}"
+            )
+
+        return load_dataset("arrow", data_files=str(path), **data_kwargs)
+
+
+@DatasetDeserializerFactory.register("hdf5_file")
+class HDF5FileDatasetDeserializer(DatasetDeserializer):
+    def __call__(
+        self,
+        data: Any,
+        processor_factory: Callable[[], PreTrainedTokenizerBase],
+        random_seed: int,
+        **data_kwargs: dict[str, Any],
+    ) -> Dataset:
+        _ = (processor_factory, random_seed)
+        if (
+            not isinstance(data, str | Path)
+            or not (path := Path(data)).exists()
+            or not path.is_file()
+            or path.suffix.lower() not in {".hdf5", ".h5"}
+        ):
+            raise DataNotSupportedError(
+                f"Unsupported data for HDF5FileDatasetDeserializer, "
+                f"expected str or Path to a local .hdf5 or .h5 file, got {data}"
+            )
+
+        return Dataset.from_pandas(pd.read_hdf(str(path)), **data_kwargs)
+
+
+@DatasetDeserializerFactory.register("db_file")
+class DBFileDatasetDeserializer(DatasetDeserializer):
+    def __call__(
+        self,
+        data: Any,
+        processor_factory: Callable[[], PreTrainedTokenizerBase],
+        random_seed: int,
+        **data_kwargs: dict[str, Any],
+    ) -> dict[str, list]:
+        _ = (processor_factory, random_seed)
+        if (
+            not isinstance(data, str | Path)
+            or not (path := Path(data)).exists()
+            or not path.is_file()
+            or path.suffix.lower() != ".db"
+        ):
+            raise DataNotSupportedError(
+                f"Unsupported data for DBFileDatasetDeserializer, "
+                f"expected str or Path to a local .db file, got {data}"
+            )
+
+        return Dataset.from_sql(con=str(path), **data_kwargs)
+
+
+@DatasetDeserializerFactory.register("tar_file")
+class TarFileDatasetDeserializer(DatasetDeserializer):
+    def __call__(
+        self,
+        data: Any,
+        processor_factory: Callable[[], PreTrainedTokenizerBase],
+        random_seed: int,
+        **data_kwargs: dict[str, Any],
+    ) -> dict[str, list]:
+        _ = (processor_factory, random_seed)
+        if (
+            not isinstance(data, str | Path)
+            or not (path := Path(data)).exists()
+            or not path.is_file()
+            or path.suffix.lower() != ".tar"
+        ):
+            raise DataNotSupportedError(
+                f"Unsupported data for TarFileDatasetDeserializer, "
+                f"expected str or Path to a local .tar file, got {data}"
+            )
+
+        return load_dataset("webdataset", data_files=str(path), **data_kwargs)

guidellm/data/deserializers/huggingface.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any
+
+from datasets import (
+    Dataset,
+    DatasetDict,
+    IterableDataset,
+    IterableDatasetDict,
+    load_dataset,
+    load_from_disk,
+)
+from datasets.exceptions import (
+    DataFilesNotFoundError,
+    DatasetNotFoundError,
+    FileNotFoundDatasetsError,
+)
+from transformers import PreTrainedTokenizerBase
+
+from guidellm.data.deserializers.deserializer import (
+    DataNotSupportedError,
+    DatasetDeserializer,
+    DatasetDeserializerFactory,
+)
+
+__all__ = ["HuggingFaceDatasetDeserializer"]
+
+
+@DatasetDeserializerFactory.register("huggingface")
+class HuggingFaceDatasetDeserializer(DatasetDeserializer):
+    def __call__(
+        self,
+        data: Any,
+        processor_factory: Callable[[], PreTrainedTokenizerBase],
+        random_seed: int,
+        **data_kwargs: dict[str, Any],
+    ) -> Dataset | IterableDataset | DatasetDict | IterableDatasetDict:
+        _ = (processor_factory, random_seed)
+
+        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), **data_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), **data_kwargs)
+                except (
+                    FileNotFoundDatasetsError,
+                    DatasetNotFoundError,
+                    DataFilesNotFoundError,
+                ) as err2:
+                    load_error = err2
+
+        try:
+            # Handle dataset identifier from the Hugging Face Hub
+            return load_dataset(str(data), **data_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