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

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)}"
+        )

guidellm/data/deserializers/synthetic.py

@@ -0,0 +1,348 @@
+from __future__ import annotations
+
+import math
+from collections.abc import Callable, Iterator
+from pathlib import Path
+from random import Random
+from typing import Any
+
+import yaml
+from datasets import Features, IterableDataset, Value
+from faker import Faker
+from pydantic import ConfigDict, Field, model_validator
+from transformers import PreTrainedTokenizerBase
+
+from guidellm.data.deserializers.deserializer import (
+    DataNotSupportedError,
+    DatasetDeserializer,
+    DatasetDeserializerFactory,
+)
+from guidellm.utils import IntegerRangeSampler, StandardBaseModel
+
+__all__ = [
+    "SyntheticTextDatasetConfig",
+    "SyntheticTextDatasetDeserializer",
+    "SyntheticTextGenerator",
+    "SyntheticTextPrefixBucketConfig",
+]
+
+
+class SyntheticTextPrefixBucketConfig(StandardBaseModel):
+    bucket_weight: int = Field(
+        description="Weight of this bucket in the overall distribution.",
+        gt=0,
+        default=100,
+    )
+    prefix_count: int = Field(
+        description="The number of unique prefixes to generate for this bucket.",
+        ge=1,
+        default=1,
+    )
+    prefix_tokens: int = Field(
+        description="The number of prefix tokens per-prompt for this bucket.",
+        ge=0,
+        default=0,
+    )
+
+
+class SyntheticTextDatasetConfig(StandardBaseModel):
+    model_config = ConfigDict(
+        extra="allow",
+    )
+
+    prefix_buckets: list[SyntheticTextPrefixBucketConfig] | None = Field(
+        description="Buckets for the prefix tokens distribution.",
+        default=None,
+    )
+    prompt_tokens: int = Field(
+        description="The average number of text tokens generated for prompts.",
+        gt=0,
+    )
+    prompt_tokens_stdev: int | None = Field(
+        description="The standard deviation of the tokens generated for prompts.",
+        gt=0,
+        default=None,
+    )
+    prompt_tokens_min: int | None = Field(
+        description="The minimum number of text tokens generated for prompts.",
+        gt=0,
+        default=None,
+    )
+    prompt_tokens_max: int | None = Field(
+        description="The maximum number of text tokens generated for prompts.",
+        gt=0,
+        default=None,
+    )
+    output_tokens: int = Field(
+        description="The average number of text tokens generated for outputs.",
+        gt=0,
+    )
+    output_tokens_stdev: int | None = Field(
+        description="The standard deviation of the tokens generated for outputs.",
+        gt=0,
+        default=None,
+    )
+    output_tokens_min: int | None = Field(
+        description="The minimum number of text tokens generated for outputs.",
+        gt=0,
+        default=None,
+    )
+    output_tokens_max: int | None = Field(
+        description="The maximum number of text tokens generated for outputs.",
+        gt=0,
+        default=None,
+    )
+    source: str = Field(
+        description="The source of the text data to be used for generation.",
+        default="data:prideandprejudice.txt.gz",
+    )
+
+    @model_validator(mode="after")
+    def check_prefix_options(self) -> SyntheticTextDatasetConfig:
+        if self.__pydantic_extra__ is not None:
+            prefix_count = self.__pydantic_extra__.get("prefix_count", None)  # type: ignore[attr-defined]
+            prefix_tokens = self.__pydantic_extra__.get("prefix_tokens", None)  # type: ignore[attr-defined]
+
+            if prefix_count is not None or prefix_tokens is not None:
+                if self.prefix_buckets:
+                    raise ValueError(
+                        "prefix_buckets is mutually exclusive"
+                        " with prefix_count and prefix_tokens"
+                    )
+
+                self.prefix_buckets = [
+                    SyntheticTextPrefixBucketConfig(
+                        prefix_count=prefix_count or 1,
+                        prefix_tokens=prefix_tokens or 0,
+                    )
+                ]
+
+        return self
+
+
+class SyntheticTextGenerator:
+    def __init__(
+        self,
+        config: SyntheticTextDatasetConfig,
+        processor: PreTrainedTokenizerBase,
+        random_seed: int = 42,
+    ):
+        self.config = config
+        self.processor = processor
+        self.random_seed = random_seed
+
+    def __iter__(self) -> Iterator[dict[str, Any]]:
+        samples_generated = 0
+
+        faker = Faker()
+        faker.seed_instance(self.random_seed)
+        prompt_tokens_sampler = iter(
+            IntegerRangeSampler(
+                average=self.config.prompt_tokens,
+                variance=self.config.prompt_tokens_stdev,
+                min_value=self.config.prompt_tokens_min,
+                max_value=self.config.prompt_tokens_max,
+                random_seed=self.random_seed,
+            )
+        )
+        output_tokens_sampler = iter(
+            IntegerRangeSampler(
+                average=self.config.output_tokens,
+                variance=self.config.output_tokens_stdev,
+                min_value=self.config.output_tokens_min,
+                max_value=self.config.output_tokens_max,
+                random_seed=self.random_seed + 1,  # ensure diff dist from prompts
+            )
+        )
+
+        # Create a shared prefix if specified
+        rand = Random(self.random_seed + 3)
+        prefix_iter = self._create_prefix_iter(faker, rand)
+
+        while True:
+            prompt_tokens_count = next(prompt_tokens_sampler)
+            output_tokens_count = next(output_tokens_sampler)
+
+            yield {
+                "prefix": next(prefix_iter),
+                "prompt": self._create_prompt(
+                    prompt_tokens_count, faker, f"{samples_generated} "
+                ),
+                "prompt_tokens_count": prompt_tokens_count,
+                "output_tokens_count": output_tokens_count,
+            }
+            samples_generated += 1
+
+    def _create_prompt(
+        self, prompt_tokens_count: int, faker: Faker, unique: str = ""
+    ) -> str:
+        prompt_token_ids: list[int] = []
+        avg_chars_per_token = 5
+        margin_of_safety = 1.5
+        attempts = 0
+
+        while len(prompt_token_ids) < prompt_tokens_count:
+            attempts += 1
+            num_chars = int(
+                prompt_tokens_count * avg_chars_per_token * margin_of_safety * attempts
+            )
+            text = unique + faker.text(max_nb_chars=num_chars)
+            prompt_token_ids = self.processor.encode(text)
+
+        return self.processor.decode(
+            prompt_token_ids[:prompt_tokens_count], skip_special_tokens=True
+        )
+
+    def _create_prefix_iter(self, faker: Faker, rand: Random) -> Iterator[str]:
+        if not self.config.prefix_buckets:
+            while True:
+                yield ""
+
+        # Increase weights to ensure an integer number of samples per per-prefix
+        least_common_prefix_count = math.lcm(
+            *(bucket.prefix_count for bucket in self.config.prefix_buckets)
+        )
+        unnorm_weights = [
+            least_common_prefix_count * bucket.bucket_weight // bucket.prefix_count
+            for bucket in self.config.prefix_buckets
+        ]
+        # Use GCD to reduce the weights to smallest integer ratio
+        common_divisor = math.gcd(*unnorm_weights)
+
+        # Create prefix list maintaining the correct distribution
+        prefixes = []
+        for bucket, weight in zip(
+            self.config.prefix_buckets, unnorm_weights, strict=False
+        ):
+            bucket_prefixes = [
+                self._create_prompt(bucket.prefix_tokens, faker)
+                for _ in range(bucket.prefix_count)
+            ]
+            sample_count = weight // common_divisor
+            prefixes.extend(bucket_prefixes * sample_count)
+
+        while True:
+            yield rand.choice(prefixes)
+
+
+@DatasetDeserializerFactory.register("synthetic_text")
+class SyntheticTextDatasetDeserializer(DatasetDeserializer):
+    def __call__(
+        self,
+        data: Any,
+        processor_factory: Callable[[], PreTrainedTokenizerBase],
+        random_seed: int,
+        **data_kwargs: dict[str, Any],
+    ) -> IterableDataset:
+        # Config file pathways, deserialize and call self again
+        if (config := self._load_config_file(data)) is not None:
+            return self(config, processor_factory, random_seed, **data_kwargs)
+
+        # Config str pathways, deserialize and call self again
+        if (config := self._load_config_str(data)) is not None:
+            return self(config, processor_factory, random_seed, **data_kwargs)
+
+        if not isinstance(data, SyntheticTextDatasetConfig):
+            raise DataNotSupportedError(
+                "Unsupported data for SyntheticTextDatasetDeserializer, "
+                "expected SyntheticTextDatasetConfig, str or Path to a config file, "
+                f"got {data}"
+            )
+
+        return IterableDataset.from_generator(
+            SyntheticTextGenerator,
+            gen_kwargs={
+                "config": data,
+                "processor": processor_factory(),
+                "random_seed": random_seed,
+            },
+            features=Features(
+                {
+                    "prefix": Value("string"),
+                    "prompt": Value("string"),
+                    "prompt_tokens_count": Value("int32"),
+                    "output_tokens_count": Value("int32"),
+                }
+            ),
+        )
+
+    def _load_config_file(self, data: Any) -> SyntheticTextDatasetConfig | None:
+        if (not isinstance(data, str) and not isinstance(data, Path)) or (
+            not Path(data).is_file()
+        ):
+            return None
+
+        data_path = Path(data) if isinstance(data, str) else data
+        error = None
+
+        if Path(data).is_file() and data_path.suffix.lower() == ".json":
+            try:
+                return SyntheticTextDatasetConfig.model_validate_json(
+                    data_path.read_text()
+                )
+            except Exception as err:  # noqa: BLE001
+                error = err
+
+        if Path(data).is_file() and data_path.suffix.lower() in {
+            ".yaml",
+            ".yml",
+            ".config",
+        }:
+            try:
+                return SyntheticTextDatasetConfig.model_validate(
+                    yaml.safe_load(data_path.read_text())
+                )
+            except Exception as err:  # noqa: BLE001
+                error = err
+
+        err_message = (
+            f"Unsupported file {data_path} for "
+            f"SyntheticTextDatasetDeserializer, expected .json, "
+            f".yaml, .yml, or .config"
+        )
+
+        if error is not None:
+            err_message += f" with error: {error}"
+            raise DataNotSupportedError(err_message) from error
+        raise DataNotSupportedError(err_message)
+
+    def _load_config_str(self, data: str) -> SyntheticTextDatasetConfig | None:
+        if not isinstance(data, str):
+            return None
+
+        data_str = data.strip()
+        error = None
+
+        if (data_str.startswith("{") and data_str.endswith("}")) or (
+            data_str.startswith("[") and data_str.endswith("]")
+        ):
+            try:
+                return SyntheticTextDatasetConfig.model_validate_json(data_str)
+            except Exception as err:  # noqa: BLE001
+                error = err
+
+        if data_str.count("=") > 1:
+            # key=value pairs separated by commas
+            try:
+                config_dict = {}
+                items = data_str.split(",")
+                for item in items:
+                    key, value = item.split("=")
+                    config_dict[key.strip()] = (
+                        int(value.strip())
+                        if value.strip().isnumeric()
+                        else value.strip()
+                    )
+
+                return SyntheticTextDatasetConfig.model_validate(config_dict)
+            except Exception as err:  # noqa: BLE001
+                error = err
+
+        err_message = (
+            "Unsupported string data for SyntheticTextDatasetDeserializer, "
+            f"expected JSON or key-value pairs, got {data}"
+        )
+        if error is not None:
+            err_message += f" with error: {error}"
+            raise DataNotSupportedError(err_message) from error
+        raise DataNotSupportedError(err_message)

guidellm/data/loaders.py

@@ -0,0 +1,149 @@
+from __future__ import annotations
+
+import contextlib
+from collections.abc import Callable, Iterator
+from typing import Any, Literal
+
+import torch
+from torch.utils.data import Sampler
+from torch.utils.data.dataloader import DataLoader as PyTorchDataLoader
+from torch.utils.data.dataset import IterableDataset as TorchIterableDataset
+from transformers import PreTrainedTokenizerBase
+
+from guidellm.data.deserializers import DatasetDeserializerFactory
+from guidellm.data.preprocessors import DataDependentPreprocessor, DatasetPreprocessor
+from guidellm.logger import logger
+
+__all__ = ["DataLoader", "DatasetsIterator"]
+
+
+
+class DatasetsIterator(TorchIterableDataset):
+    def __init__(
+        self,
+        data: list[Any],
+        data_args: list[dict[str, Any]] | None,
+        data_samples: int,
+        processor_factory: Callable[[], PreTrainedTokenizerBase],
+        preprocessors: list[DatasetPreprocessor | DataDependentPreprocessor],
+        random_seed: int,
+    ):
+        if not data or not isinstance(data, list):
+            raise ValueError(f"Data must be a non-empty list, got {data}.")
+
+        if not data_args:
+            data_args = [{} for _ in data]
+
+        if len(data) != len(data_args):
+            raise ValueError(
+                f"Length of data ({len(data)}) must match length of data_args "
+                f"({len(data_args)})."
+            )
+
+        self.datasets = []
+        for datum, data_kwargs in zip(data, data_args, strict=False):
+            self.datasets.append(
+                DatasetDeserializerFactory.deserialize(
+                    data=datum,
+                    processor_factory=processor_factory,
+                    random_seed=random_seed,
+                    **data_kwargs,
+                )
+            )
+        self.preprocessors = preprocessors
+        for preprocessor in self.preprocessors:
+            if isinstance(preprocessor, DataDependentPreprocessor):
+                preprocessor.setup_data(
+                    datasets=self.datasets,
+                    data_args=data_args,
+                )
+        self.precache: list[Any] | None = (
+            list(self.generator(data_samples)) if data_samples else None
+        )
+
+    def __iter__(self):
+        worker_info = torch.utils.data.get_worker_info()
+        worker_modulus = worker_info.num_workers if worker_info is not None else 1
+        worker_index = worker_info.id if worker_info is not None else 0
+
+        if self.precache:
+            for index, item in enumerate(self.precache):
+                if (index + worker_index) % worker_modulus == 0:
+                    yield item
+        else:
+            yield from self.generator(modulus=worker_modulus, offset=worker_index)
+
+    def generator(
+        self,
+        max_items: int | None = None,
+        modulus: int | None = None,
+        offset: int | None = None,
+    ) -> Iterator[Any]:
+        gen_count = 0
+
+        with contextlib.suppress(StopIteration):
+            dataset_iters = [iter(dataset) for dataset in self.datasets]
+
+            while max_items is None or gen_count < max_items:
+                try:
+                    row: dict[str, Any] = {
+                        "items": [next(dataset_iter) for dataset_iter in dataset_iters]
+                    }
+                    gen_count += 1
+
+                    if (
+                        modulus is not None
+                        and offset is not None
+                        and (gen_count % modulus) != offset
+                    ):
+                        continue
+
+                    for preprocessor in self.preprocessors:
+                        # This can assign a GenerationRequest, which would then be
+                        # passed into the preprocessor, which is a type violation.
+                        # This should be fixed at some point.
+                        row = preprocessor(row)  # type: ignore[assignment]
+                    yield row
+                except Exception as err:  # noqa: BLE001 # Exception logged
+                    logger.error(f"Skipping data row due to error: {err}")
+                    gen_count -= 1
+
+        if max_items is not None and gen_count < max_items:
+            raise ValueError(
+                f"Requested {max_items} samples, but only {gen_count} "
+                "available from the provided datasets."
+            )
+
+
+class DataLoader(PyTorchDataLoader):
+    def __init__(
+        self,
+        data: list[Any],
+        data_args: list[dict[str, Any]] | None,
+        data_samples: int,
+        processor_factory: Callable[[], PreTrainedTokenizerBase],
+        preprocessors: list[DatasetPreprocessor | DataDependentPreprocessor],
+        collator: Callable,
+        sampler: Sampler[int] | Literal["shuffle"] | None = None,
+        num_workers: int | None = 1,
+        random_seed: int = 42,
+        **kwargs: Any,
+    ):
+        iterator = DatasetsIterator(
+            data=data,
+            data_args=data_args,
+            data_samples=data_samples,
+            processor_factory=processor_factory,
+            preprocessors=preprocessors,
+            random_seed=random_seed,
+        )
+
+        super().__init__(
+            dataset=iterator,
+            batch_size=1,
+            shuffle=sampler == "shuffle",
+            sampler=sampler if sampler != "shuffle" else None,
+            collate_fn=collator,
+            num_workers=num_workers or 0,
+            **kwargs,
+        )