guidellm 0.3.1__py3-none-any.whl → 0.6.0a5__py3-none-any.whl

guidellm/data/deserializers/__init__.py

@@ -0,0 +1,49 @@
+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 (
+    SyntheticTextDataset,
+    SyntheticTextDatasetDeserializer,
+)
+
+__all__ = [
+    "ArrowFileDatasetDeserializer",
+    "CSVFileDatasetDeserializer",
+    "DBFileDatasetDeserializer",
+    "DataNotSupportedError",
+    "DatasetDeserializer",
+    "DatasetDeserializerFactory",
+    "HDF5FileDatasetDeserializer",
+    "HuggingFaceDatasetDeserializer",
+    "InMemoryCsvDatasetDeserializer",
+    "InMemoryDictDatasetDeserializer",
+    "InMemoryDictListDatasetDeserializer",
+    "InMemoryItemListDatasetDeserializer",
+    "InMemoryJsonStrDatasetDeserializer",
+    "JSONFileDatasetDeserializer",
+    "ParquetFileDatasetDeserializer",
+    "SyntheticTextDataset",
+    "SyntheticTextDatasetDeserializer",
+    "TarFileDatasetDeserializer",
+    "TextFileDatasetDeserializer",
+]

guidellm/data/deserializers/deserializer.py

@@ -0,0 +1,141 @@
+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.schemas import DataNotSupportedError
+from guidellm.data.utils import resolve_dataset_split
+from guidellm.utils import RegistryMixin
+
+__all__ = [
+    "DatasetDeserializer",
+    "DatasetDeserializerFactory",
+]
+
+
+@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,223 @@
+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 valid 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

guidellm/data/deserializers/memory.py

@@ -0,0 +1,194 @@
+from __future__ import annotations
+
+import contextlib
+import csv
+import json
+from collections.abc import Callable
+from io import StringIO
+from typing import Any, cast
+
+from datasets import Dataset
+from transformers import PreTrainedTokenizerBase
+
+from guidellm.data.deserializers.deserializer import (
+    DataNotSupportedError,
+    DatasetDeserializer,
+    DatasetDeserializerFactory,
+)
+
+__all__ = [
+    "InMemoryCsvDatasetDeserializer",
+    "InMemoryDictDatasetDeserializer",
+    "InMemoryDictListDatasetDeserializer",
+    "InMemoryItemListDatasetDeserializer",
+    "InMemoryJsonStrDatasetDeserializer",
+]
+
+
+@DatasetDeserializerFactory.register("in_memory_dict")
+class InMemoryDictDatasetDeserializer(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 data
+            or not isinstance(data, dict)
+            or not all(
+                isinstance(key, str) and isinstance(val, list)
+                for key, val in data.items()
+            )
+        ):
+            raise DataNotSupportedError(
+                f"Unsupported data for InMemoryDictDatasetDeserializer, "
+                f"expected dict[str, list], got {data}"
+            )
+
+        rows = len(list(data.values())[0])
+        if not all(len(val) == rows for val in data.values()):
+            raise DataNotSupportedError(
+                "All lists in the data dictionary must have the same length, "
+                f"expected {rows} for all keys {list(data.keys())}"
+            )
+
+        return Dataset.from_dict(data, **data_kwargs)
+
+
+@DatasetDeserializerFactory.register("in_memory_dict_list")
+class InMemoryDictListDatasetDeserializer(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 data
+            or not isinstance(data, list)
+            or not all(isinstance(item, dict) for item in data)
+            or not all(isinstance(key, str) for item in data for key in item)
+        ):
+            raise DataNotSupportedError(
+                f"Unsupported data for InMemoryDictListDatasetDeserializer, "
+                f"expected list of dicts, got {data}"
+            )
+
+        typed_data: list[dict[str, Any]] = cast("list[dict[str, Any]]", data)
+        first_keys = set(typed_data[0].keys())
+        for index, item in enumerate(typed_data):
+            if set(item.keys()) != first_keys:
+                raise DataNotSupportedError(
+                    f"All dictionaries must have the same keys. "
+                    f"Expected keys: {first_keys}, "
+                    f"got keys at index {index}: {set(item.keys())}"
+                )
+
+        # Convert list of dicts to dict of lists
+        result_dict: dict = {key: [] for key in first_keys}
+        for item in typed_data:
+            for key, value in item.items():
+                result_dict[key].append(value)
+
+        return Dataset.from_dict(result_dict, **data_kwargs)
+
+
+@DatasetDeserializerFactory.register("in_memory_item_list")
+class InMemoryItemListDatasetDeserializer(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
+
+        primitive_types = (str, int, float, bool, type(None))
+        if (
+            not data
+            or not isinstance(data, list)
+            or not all(isinstance(item, primitive_types) for item in data)
+        ):
+            raise DataNotSupportedError(
+                f"Unsupported data for InMemoryItemListDatasetDeserializer, "
+                f"expected list of primitive items, got {data}"
+            )
+
+        column_name = data_kwargs.pop("column_name", "data")
+
+        return Dataset.from_dict({column_name: data}, **data_kwargs)
+
+
+@DatasetDeserializerFactory.register("in_memory_json_str")
+class InMemoryJsonStrDatasetDeserializer(DatasetDeserializer):
+    def __call__(
+        self,
+        data: Any,
+        processor_factory: Callable[[], PreTrainedTokenizerBase],
+        random_seed: int,
+        **data_kwargs: dict[str, Any],
+    ) -> Dataset:
+        if (
+            isinstance(data, str)
+            and (json_str := data.strip())
+            and (
+                (json_str.startswith("{") and json_str.endswith("}"))
+                or (json_str.startswith("[") and json_str.endswith("]"))
+            )
+        ):
+            with contextlib.suppress(Exception):
+                parsed_data = json.loads(data)
+
+            deserializers = [
+                InMemoryDictDatasetDeserializer(),
+                InMemoryDictListDatasetDeserializer(),
+                InMemoryItemListDatasetDeserializer(),
+            ]
+
+            for deserializer in deserializers:
+                with contextlib.suppress(DataNotSupportedError):
+                    return deserializer(
+                        parsed_data, processor_factory, random_seed, **data_kwargs
+                    )
+
+        raise DataNotSupportedError(
+            f"Unsupported data for InMemoryJsonStrDatasetDeserializer, "
+            f"expected JSON string with a list or dict of items, got {data}"
+        )
+
+
+@DatasetDeserializerFactory.register("in_memory_csv_str")
+class InMemoryCsvDatasetDeserializer(DatasetDeserializer):
+    def __call__(
+        self,
+        data: Any,
+        processor_factory: Callable[[], PreTrainedTokenizerBase],
+        random_seed: int,
+        **data_kwargs: dict[str, Any],
+    ) -> Dataset:
+        if (
+            isinstance(data, str)
+            and (csv_str := data.strip())
+            and len(csv_str.split("\n")) > 0
+        ):
+            with contextlib.suppress(Exception):
+                csv_buffer = StringIO(data)
+                reader = csv.DictReader(csv_buffer)
+                rows = list(reader)
+
+                return InMemoryDictListDatasetDeserializer()(
+                    rows, processor_factory, random_seed, **data_kwargs
+                )
+
+        raise DataNotSupportedError(
+            f"Unsupported data for InMemoryCsvDatasetDeserializer, "
+            f"expected CSV string, got {type(data)}"
+        )