guidellm 0.4.0a18__py3-none-any.whl → 0.4.0a155__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,109 @@
+from __future__ import annotations
+
+import contextlib
+from collections.abc import Callable
+from typing import Any, Protocol, Union, runtime_checkable
+
+from datasets import Dataset, IterableDataset
+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],
+    ) -> dict[str, list]: ...
+
+
+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 = None
+
+        if type_ is None:
+            errors = []
+            # 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:
+                    with contextlib.suppress(DataNotSupportedError):
+                        dataset = deserializer_fn(
+                            data=data,
+                            processor_factory=processor_factory,
+                            random_seed=random_seed,
+                            **data_kwargs,
+                        )
+                except Exception as e:
+                    errors.append(e)
+
+                if dataset is not None:
+                    break # Found one that works. Continuing could overwrite it.
+
+            if dataset is None and len(errors) > 0:
+                raise DataNotSupportedError(f"data deserialization failed; {len(errors)} errors occurred while "
+                                            f"attempting to deserialize data {data}: {errors}")
+
+        elif deserializer := cls.get_registered_object(type_) is not None:
+            deserializer_fn: DatasetDeserializer = (
+                deserializer() if isinstance(deserializer, type) else deserializer
+            )
+
+            dataset = deserializer_fn(
+                data=data,
+                processor_factory=processor_factory,
+                random_seed=random_seed,
+                **data_kwargs,
+            )
+
+        if dataset is None:
+            raise DataNotSupportedError(
+                f"No suitable deserializer found for data {data} "
+                f"with kwargs {data_kwargs} and deserializer type {type_}."
+            )
+
+        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

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],
+    ) -> dict[str, list]:
+        _ = (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],
+    ) -> 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() != ".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],
+    ) -> 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() 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],
+    ) -> 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() != ".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],
+    ) -> 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() != ".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],
+    ) -> 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() 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],
+    ) -> dict[str, list]:
+        _ = (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