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

guidellm/dataset/creator.py

@@ -1,213 +0,0 @@
-from abc import ABC, abstractmethod
-from pathlib import Path
-from typing import Any, Literal, Optional, Union
-
-from datasets import Dataset, DatasetDict, IterableDataset, IterableDatasetDict
-from transformers import PreTrainedTokenizerBase  # type: ignore[import]
-
-__all__ = ["ColumnInputTypes", "DatasetCreator"]
-
-ColumnInputTypes = Literal[
-    "prompt_column",
-    "text_column",
-    "prompt_tokens_count_column",
-    "output_tokens_count_column",
-]
-
-
-class DatasetCreator(ABC):
-    DEFAULT_SPLITS_TRAIN = [
-        "train",
-        "training",
-        "train_set",
-        "training_set",
-        "train_dataset",
-        "training_dataset",
-        "train_data",
-        "training_data",
-        "pretrain",
-        "pretrain_set",
-        "pretrain_dataset",
-        "pretrain_data",
-        "pretraining",
-    ]
-    DEFAULT_SPLITS_CALIB = [
-        "calibration",
-        "calib",
-        "cal",
-        "calibration_set",
-        "calib_set",
-        "cal_set",
-        "calibration_dataset",
-        "calib_dataset",
-        "cal_set",
-        "calibration_data",
-        "calib_data",
-        "cal_data",
-    ]
-    DEFAULT_SPLITS_VAL = [
-        "validation",
-        "val",
-        "valid",
-        "validation_set",
-        "val_set",
-        "validation_dataset",
-        "val_dataset",
-        "validation_data",
-        "val_data",
-        "dev",
-        "dev_set",
-        "dev_dataset",
-        "dev_data",
-    ]
-    DEFAULT_SPLITS_TEST = [
-        "test",
-        "testing",
-        "test_set",
-        "testing_set",
-        "test_dataset",
-        "testing_dataset",
-        "test_data",
-        "testing_data",
-        "eval",
-        "eval_set",
-        "eval_dataset",
-        "eval_data",
-    ]
-    DEFAULT_SPLITS_DATASET: dict[str, str] = {}
-
-    @classmethod
-    def create(
-        cls,
-        data: Any,
-        data_args: Optional[dict[str, Any]],
-        processor: Optional[Union[str, Path, PreTrainedTokenizerBase]],
-        processor_args: Optional[dict[str, Any]],
-        random_seed: int = 42,
-        split_pref_order: Optional[list[str]] = None,
-    ) -> tuple[Union[Dataset, IterableDataset], dict[ColumnInputTypes, str]]:
-        if not cls.is_supported(data, data_args):
-            raise ValueError(f"Unsupported data type: {type(data)} given for {data}. ")
-
-        split = cls.extract_args_split(data_args)
-        column_mappings = cls.extract_args_column_mappings(data_args)
-        dataset = cls.handle_create(
-            data, data_args, processor, processor_args, random_seed
-        )
-
-        if isinstance(dataset, (DatasetDict, IterableDatasetDict)):
-            dataset = cls.extract_dataset_split(dataset, split, split_pref_order)
-
-        if not isinstance(dataset, (Dataset, IterableDataset)):
-            raise ValueError(
-                f"Unsupported data type: {type(dataset)} given for {dataset}."
-            )
-
-        return dataset, column_mappings
-
-    @classmethod
-    def extract_args_split(cls, data_args: Optional[dict[str, Any]]) -> str:
-        split = "auto"
-
-        if data_args and "split" in data_args:
-            split = data_args["split"]
-            del data_args["split"]
-
-        return split
-
-    @classmethod
-    def extract_args_column_mappings(
-        cls,
-        data_args: Optional[dict[str, Any]],
-    ) -> dict[ColumnInputTypes, str]:
-        columns: dict[ColumnInputTypes, str] = {}
-
-        if data_args:
-            if "prompt_column" in data_args:
-                columns["prompt_column"] = data_args["prompt_column"]
-                del data_args["prompt_column"]
-
-            if "prompt_tokens_count_column" in data_args:
-                columns["prompt_tokens_count_column"] = data_args[
-                    "prompt_tokens_count_column"
-                ]
-                del data_args["prompt_tokens_count_column"]
-
-            if "output_tokens_count_column" in data_args:
-                columns["output_tokens_count_column"] = data_args[
-                    "output_tokens_count_column"
-                ]
-                del data_args["output_tokens_count_column"]
-
-        return columns
-
-    @classmethod
-    def extract_dataset_name(
-        cls, dataset: Union[Dataset, IterableDataset, DatasetDict, IterableDatasetDict]
-    ) -> Optional[str]:
-        if isinstance(dataset, (DatasetDict, IterableDatasetDict)):
-            dataset = dataset[list(dataset.keys())[0]]
-
-        if isinstance(dataset, (Dataset, IterableDataset)):
-            if not hasattr(dataset, "info") or not hasattr(
-                dataset.info, "dataset_name"
-            ):
-                return None
-
-            return dataset.info.dataset_name
-
-        raise ValueError(f"Unsupported data type: {type(dataset)} given for {dataset}.")
-
-    @classmethod
-    def extract_dataset_split(
-        cls,
-        dataset: Union[DatasetDict, IterableDatasetDict],
-        specified_split: Union[Literal["auto"], str] = "auto",
-        split_pref_order: Optional[Union[Literal["auto"], list[str]]] = "auto",
-    ) -> Union[Dataset, IterableDataset]:
-        if not isinstance(dataset, (DatasetDict, IterableDatasetDict)):
-            raise ValueError(
-                f"Unsupported data type: {type(dataset)} given for {dataset}."
-            )
-
-        if specified_split != "auto":
-            if specified_split not in dataset:
-                raise ValueError(
-                    f"Split {specified_split} not found in dataset {dataset}."
-                )
-
-            return dataset[specified_split]
-
-        dataset_name = cls.extract_dataset_name(dataset)
-
-        if dataset_name and dataset_name in cls.DEFAULT_SPLITS_DATASET:
-            return dataset[cls.DEFAULT_SPLITS_DATASET[dataset_name]]
-
-        if split_pref_order == "auto":
-            split_pref_order = [
-                *cls.DEFAULT_SPLITS_TEST,
-                *cls.DEFAULT_SPLITS_VAL,
-                *cls.DEFAULT_SPLITS_CALIB,
-                *cls.DEFAULT_SPLITS_TRAIN,
-            ]
-
-        for test_split in split_pref_order or []:
-            if test_split in dataset:
-                return dataset[test_split]
-
-        return dataset[list(dataset.keys())[0]]
-
-    @classmethod
-    @abstractmethod
-    def is_supported(cls, data: Any, data_args: Optional[dict[str, Any]]) -> bool: ...
-
-    @classmethod
-    @abstractmethod
-    def handle_create(
-        cls,
-        data: Any,
-        data_args: Optional[dict[str, Any]],
-        processor: Optional[Union[str, Path, PreTrainedTokenizerBase]],
-        processor_args: Optional[dict[str, Any]],
-        random_seed: int,
-    ) -> Union[Dataset, DatasetDict, IterableDataset, IterableDatasetDict]: ...

guidellm/dataset/entrypoints.py

@@ -1,42 +0,0 @@
-from pathlib import Path
-from typing import Any, Optional, Union
-
-from datasets import Dataset, IterableDataset
-from transformers import PreTrainedTokenizerBase  # type: ignore[import]
-
-from guidellm.dataset.creator import ColumnInputTypes
-from guidellm.dataset.file import FileDatasetCreator
-from guidellm.dataset.hf_datasets import HFDatasetsCreator
-from guidellm.dataset.in_memory import InMemoryDatasetCreator
-from guidellm.dataset.synthetic import SyntheticDatasetCreator
-
-__all__ = ["load_dataset"]
-
-
-def load_dataset(
-    data: Any,
-    data_args: Optional[dict[str, Any]],
-    processor: Optional[Union[str, Path, PreTrainedTokenizerBase]],
-    processor_args: Optional[dict[str, Any]],
-    random_seed: int = 42,
-    split_pref_order: Optional[list[str]] = None,
-) -> tuple[Union[Dataset, IterableDataset], dict[ColumnInputTypes, str]]:
-    creators = [
-        InMemoryDatasetCreator,
-        SyntheticDatasetCreator,
-        FileDatasetCreator,
-        HFDatasetsCreator,
-    ]
-
-    for creator in creators:
-        if creator.is_supported(data, data_args):
-            return creator.create(
-                data,
-                data_args,
-                processor,
-                processor_args,
-                random_seed,
-                split_pref_order,
-            )
-
-    raise ValueError(f"Unsupported data type: {type(data)} given for {data}. ")

guidellm/dataset/file.py

@@ -1,92 +0,0 @@
-from pathlib import Path
-from typing import Any, Optional, Union
-
-import pandas as pd  # type: ignore[import]
-from datasets import (
-    Dataset,
-    DatasetDict,
-    IterableDataset,
-    IterableDatasetDict,
-    load_dataset,
-)
-from transformers import PreTrainedTokenizerBase  # type: ignore[import]
-
-from guidellm.dataset.creator import DatasetCreator
-
-__all__ = ["FileDatasetCreator"]
-
-
-class FileDatasetCreator(DatasetCreator):
-    SUPPORTED_TYPES = {
-        ".txt",
-        ".text",
-        ".csv",
-        ".json",
-        ".jsonl",
-        ".parquet",
-        ".arrow",
-        ".hdf5",
-        ".tar",
-    }
-
-    @classmethod
-    def is_supported(cls, data: Any, data_args: Optional[dict[str, Any]]) -> bool:  # noqa: ARG003
-        if isinstance(data, (str, Path)) and (path := Path(data)).exists():
-            # local folder or py file, assume supported
-            return path.suffix.lower() in cls.SUPPORTED_TYPES
-
-        return False
-
-    @classmethod
-    def handle_create(
-        cls,
-        data: Any,
-        data_args: Optional[dict[str, Any]],
-        processor: Optional[Union[str, Path, PreTrainedTokenizerBase]],  # noqa: ARG003
-        processor_args: Optional[dict[str, Any]],  # noqa: ARG003
-        random_seed: int,  # noqa: ARG003
-    ) -> Union[Dataset, DatasetDict, IterableDataset, IterableDatasetDict]:
-        if not isinstance(data, (str, Path)):
-            raise ValueError(f"Unsupported data type: {type(data)} given for {data}. ")
-
-        path = Path(data)
-        if not path.exists():
-            raise FileNotFoundError(f"File not found: {path}")
-
-        if not path.is_file():
-            raise ValueError(f"Unsupported data type: {path} given for {path}. ")
-
-        if path.suffix.lower() not in cls.SUPPORTED_TYPES:
-            raise ValueError(f"Unsupported file type: {path.suffix} given for {path}. ")
-
-        return cls.load_dataset(path, data_args)
-
-    @classmethod
-    def load_dataset(
-        cls, path: Path, data_args: Optional[dict[str, Any]]
-    ) -> Union[Dataset, IterableDataset]:
-        if path.suffix.lower() in {".txt", ".text"}:
-            with path.open("r") as file:
-                items = file.readlines()
-
-            dataset = Dataset.from_dict({"text": items}, **(data_args or {}))
-        elif path.suffix.lower() == ".csv":
-            dataset = load_dataset("csv", data_files=str(path), **(data_args or {}))
-        elif path.suffix.lower() in {".json", ".jsonl"}:
-            dataset = load_dataset("json", data_files=str(path), **(data_args or {}))
-        elif path.suffix.lower() == ".parquet":
-            dataset = load_dataset("parquet", data_files=str(path), **(data_args or {}))
-        elif path.suffix.lower() == ".arrow":
-            dataset = load_dataset("arrow", data_files=str(path), **(data_args or {}))
-        elif path.suffix.lower() == ".hdf5":
-            dataset = Dataset.from_pandas(pd.read_hdf(str(path)), **(data_args or {}))
-        elif path.suffix.lower() == ".db":
-            dataset = Dataset.from_sql(con=str(path), **(data_args or {}))
-        elif path.suffix.lower() == ".tar":
-            dataset = load_dataset(
-                "webdataset", data_files=str(path), **(data_args or {})
-            )
-        else:
-            raise ValueError(f"Unsupported file type: {path.suffix} given for {path}. ")
-
-        return dataset

guidellm/dataset/hf_datasets.py

@@ -1,62 +0,0 @@
-from pathlib import Path
-from typing import Any, Optional, Union
-
-from datasets import (
-    Dataset,
-    DatasetDict,
-    IterableDataset,
-    IterableDatasetDict,
-    get_dataset_config_info,
-    load_dataset,
-)
-from transformers import PreTrainedTokenizerBase  # type: ignore[import]
-
-from guidellm.dataset.creator import DatasetCreator
-
-__all__ = ["HFDatasetsCreator"]
-
-
-class HFDatasetsCreator(DatasetCreator):
-    @classmethod
-    def is_supported(cls, data: Any, data_args: Optional[dict[str, Any]]) -> bool:  # noqa: ARG003
-        if isinstance(
-            data, (Dataset, DatasetDict, IterableDataset, IterableDatasetDict)
-        ):
-            # base type is supported
-            return True
-
-        if isinstance(data, (str, Path)) and (path := Path(data)).exists():
-            # local folder or py file, assume supported
-            return path.is_dir() or path.suffix == ".py"
-
-        if isinstance(data, (str, Path)):
-            try:
-                # try to load dataset
-                return get_dataset_config_info(data) is not None
-            except Exception:  # noqa: BLE001, S110
-                pass
-
-        return False
-
-    @classmethod
-    def handle_create(
-        cls,
-        data: Any,
-        data_args: Optional[dict[str, Any]],
-        processor: Optional[Union[str, Path, PreTrainedTokenizerBase]],  # noqa: ARG003
-        processor_args: Optional[dict[str, Any]],  # noqa: ARG003
-        random_seed: int,  # noqa: ARG003
-    ) -> Union[Dataset, DatasetDict, IterableDataset, IterableDatasetDict]:
-        if isinstance(data, (str, Path)):
-            data = load_dataset(data, **(data_args or {}))
-        elif data_args:
-            raise ValueError(
-                f"data_args should not be provided when data is a {type(data)}"
-            )
-
-        if isinstance(
-            data, (Dataset, DatasetDict, IterableDataset, IterableDatasetDict)
-        ):
-            return data
-
-        raise ValueError(f"Unsupported data type: {type(data)} given for {data}. ")

guidellm/dataset/in_memory.py

@@ -1,132 +0,0 @@
-from collections.abc import Iterable
-from pathlib import Path
-from typing import Any, Optional, Union
-
-from datasets import (
-    Dataset,
-    DatasetDict,
-    IterableDataset,
-    IterableDatasetDict,
-)
-from transformers import PreTrainedTokenizerBase  # type: ignore[import]
-
-from guidellm.dataset.creator import DatasetCreator
-
-__all__ = ["InMemoryDatasetCreator"]
-
-
-class InMemoryDatasetCreator(DatasetCreator):
-    @classmethod
-    def is_supported(cls, data: Any, data_args: Optional[dict[str, Any]]) -> bool:  # noqa: ARG003
-        return isinstance(data, Iterable) and not isinstance(data, str)
-
-    @classmethod
-    def handle_create(
-        cls,
-        data: Any,
-        data_args: Optional[dict[str, Any]],
-        processor: Optional[Union[str, Path, PreTrainedTokenizerBase]],  # noqa: ARG003
-        processor_args: Optional[dict[str, Any]],  # noqa: ARG003
-        random_seed: int,  # noqa: ARG003
-    ) -> Union[Dataset, DatasetDict, IterableDataset, IterableDatasetDict]:
-        if not isinstance(data, Iterable):
-            raise TypeError(
-                f"Unsupported data format. Expected Iterable[Any], got {type(data)}"
-            )
-
-        if not data:
-            raise ValueError("Data is empty")
-
-        if isinstance(data, dict):
-            # assume data is a dictionary of columns and values: {"c1": ["i1", "i2"]}
-            data_dict = cls.format_data_dict(data)
-        elif isinstance(data[0], dict):  # type: ignore[index]
-            # assume data is a list of dictionaries: [{"c1": "i1"}, {"c1": "i2"}]
-            data_dict = cls.format_data_iterable_dicts(data)
-        else:
-            # assume data is a list of items with no columns: ["i1", "i2"]
-            data_dict = cls.format_data_iterable_values(data)
-
-        return Dataset.from_dict(data_dict, **(data_args or {}))
-
-    @classmethod
-    def format_data_dict(cls, data: dict[Any, Any]) -> dict[str, Any]:
-        if not isinstance(data, dict):
-            raise TypeError(
-                f"Unsupported data format. Expected Dict[str, Iterable[Any]], "
-                f"got {type(data)}"
-            )
-
-        if not all(
-            isinstance(key, str) and isinstance(val, Iterable)
-            for key, val in data.items()
-        ):
-            raise TypeError(
-                "Unsupported data format. Expected Dict[str, Iterable[Any]], "
-                f"got {type(data)}"
-            )
-
-        samples = len(list(data.values())[0])
-        if not all(len(val) == samples for val in data.values()):
-            raise ValueError(
-                "Unsupported data format. Not all columns have the same number samples "
-                f"for {data}"
-            )
-
-        return data
-
-    @classmethod
-    def format_data_iterable_dicts(
-        cls, data: Iterable[dict[Any, Any]]
-    ) -> dict[str, Any]:
-        if not isinstance(data, Iterable):
-            raise TypeError(
-                f"Unsupported data format. Expected Iterable[Dict[str, Any]], "
-                f"got {type(data)}"
-            )
-
-        if not all(isinstance(item, dict) for item in data):
-            raise TypeError(
-                f"Unsupported data format. Expected Iterable[Dict[str, Any]], "
-                f"got {type(data)}"
-            )
-
-        if not all(isinstance(key, str) for key in data[0]):  # type: ignore[index]
-            raise TypeError(
-                "Unsupported data format. Expected Dict[str, Any], "
-                f"but one of the items had a non string column for {data}"
-            )
-
-        columns = list(data[0].keys())  # type: ignore[index]
-        if not all(
-            len(item) == len(columns) and all(key in item for key in columns)
-            for item in data
-        ):
-            raise ValueError(
-                "Unsupported data format. Not all items have the same columns "
-                f"for {data}"
-            )
-
-        data_dict: dict[str, Any] = {key: [] for key in columns}
-        for item in data:
-            for key, value in item.items():
-                data_dict[key].append(value)
-
-        return data_dict
-
-    @classmethod
-    def format_data_iterable_values(cls, data: Iterable[Any]) -> dict[str, Any]:
-        if not isinstance(data, Iterable):
-            raise TypeError(
-                f"Unsupported data format. Expected Iterable[Iterable[Any]], "
-                f"got {type(data)}"
-            )
-
-        first_item = next(iter(data), None)
-        first_type = type(first_item)
-        if not all(isinstance(item, first_type) for item in data):
-            raise TypeError(
-                f"Unsupported data format. Not all types are the same for {data}"
-            )
-
-        return {"data": list(data)}