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

guidellm/data/deserializers/synthetic.py

@@ -0,0 +1,246 @@
+from __future__ import annotations
+
+import math
+from collections.abc import Callable, Iterator
+from random import Random
+from typing import Any
+
+import numpy as np
+from datasets import DatasetInfo, Features, IterableDataset, Value
+from datasets.iterable_dataset import _BaseExamplesIterable
+from faker import Faker
+from transformers import PreTrainedTokenizerBase
+
+from guidellm.data.config import load_config
+from guidellm.data.deserializers.deserializer import (
+    DataNotSupportedError,
+    DatasetDeserializer,
+    DatasetDeserializerFactory,
+)
+from guidellm.data.schemas import SyntheticTextDatasetConfig
+from guidellm.utils import IntegerRangeSampler
+
+__all__ = [
+    "SyntheticTextDataset",
+    "SyntheticTextDatasetDeserializer",
+]
+
+
+class _SyntheticTextExamplesIterable(_BaseExamplesIterable):
+    """Custom examples iterable for synthetic text generation."""
+
+    def __init__(
+        self,
+        config: SyntheticTextDatasetConfig,
+        processor: PreTrainedTokenizerBase,
+        random_seed: int,
+    ):
+        super().__init__()
+        self.config = config
+        self.processor = processor
+        self.random_seed = random_seed
+        self.iteration_count = 0
+
+    def __iter__(self) -> Iterator[tuple[int, dict[str, Any]]]:
+        iter_random_seed = self.random_seed + self.iteration_count
+        self.iteration_count += 1
+
+        faker = Faker()
+        faker.seed_instance(iter_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=iter_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=iter_random_seed + 1,  # ensure diff dist from prompts
+            )
+        )
+
+        # Create a shared prefix if specified
+        rand = Random(iter_random_seed + 3)
+        prefix_iter = self._create_prefix_iter(faker, rand)
+        samples_count = 0
+
+        while True:
+            prompt_tokens_count = next(prompt_tokens_sampler)
+            output_tokens_count = next(output_tokens_sampler)
+
+            yield (
+                samples_count,
+                {
+                    "prefix": next(prefix_iter),
+                    "prompt": self._create_prompt(
+                        prompt_tokens_count,
+                        faker,
+                        f"{self.iteration_count} {samples_count} ",
+                    ),
+                    "prompt_tokens_count": prompt_tokens_count,
+                    "output_tokens_count": output_tokens_count,
+                },
+            )
+            samples_count += 1
+
+    @property
+    def is_typed(self) -> bool:
+        return True
+
+    @property
+    def features(self) -> Features:
+        return Features(
+            {
+                "prefix": Value("string"),
+                "prompt": Value("string"),
+                "prompt_tokens_count": Value("int32"),
+                "output_tokens_count": Value("int32"),
+            }
+        )
+
+    @property
+    def num_shards(self) -> int:
+        return 1
+
+    def shuffle_data_sources(
+        self,
+        generator: np.random.Generator,  # noqa: ARG002
+    ) -> _SyntheticTextExamplesIterable:
+        """Return self since synthetic data doesn't have fixed sources to shuffle."""
+        return self
+
+    def shard_data_sources(
+        self,
+        num_shards: int,  # noqa: ARG002
+        index: int,  # noqa: ARG002
+        contiguous: bool = True,  # noqa: ARG002
+    ) -> _SyntheticTextExamplesIterable:
+        """Return self since synthetic data generation is infinite and stateless."""
+        return self
+
+    def load_state_dict(self, state_dict: dict) -> None:
+        """Load the state from a state dict."""
+        self.iteration_count = state_dict.get("iteration_count", 0)
+
+    def _init_state_dict(self) -> dict:
+        """Initialize the state dict for the iterable."""
+        self._state_dict = {"iteration_count": self.iteration_count}
+        return self._state_dict
+
+    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)
+
+
+class SyntheticTextDataset(IterableDataset):
+    def __init__(
+        self,
+        config: SyntheticTextDatasetConfig,
+        processor: PreTrainedTokenizerBase,
+        random_seed: int = 42,
+    ):
+        self.config = config
+        self.processor = processor
+        self.random_seed = random_seed
+
+        # Create the examples iterable
+        ex_iterable = _SyntheticTextExamplesIterable(
+            config=config,
+            processor=processor,
+            random_seed=random_seed,
+        )
+
+        # Initialize parent with proper ex_iterable
+        super().__init__(
+            ex_iterable=ex_iterable,
+            info=DatasetInfo(
+                description="Synthetic text dataset generator",
+                features=ex_iterable.features,
+            ),
+        )
+
+    def set_epoch(self, epoch: int):
+        """Set the epoch for the dataset iteration."""
+        if isinstance(self._ex_iterable, _SyntheticTextExamplesIterable):
+            self._ex_iterable.iteration_count = epoch
+
+
+@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 and string pathways; deserialize and call self again
+        if (config := load_config(data, SyntheticTextDatasetConfig)) 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 SyntheticTextDataset(
+            config=data,
+            processor=processor_factory(),
+            random_seed=random_seed,
+        )

guidellm/data/entrypoints.py

@@ -0,0 +1,52 @@
+from pathlib import Path
+from typing import Any
+
+from transformers import PreTrainedTokenizerBase
+
+from guidellm.data import builders
+from guidellm.data.builders import ShortPromptStrategy
+
+
+def process_dataset(
+    data: str | Path,
+    output_path: str | Path,
+    processor: str | Path | PreTrainedTokenizerBase,
+    config: str | Path,
+    processor_args: dict[str, Any] | None = None,
+    data_args: dict[str, Any] | None = None,
+    data_column_mapper: dict[str, str] | None = None,
+    short_prompt_strategy: ShortPromptStrategy = ShortPromptStrategy.IGNORE,
+    pad_char: str | None = None,
+    concat_delimiter: str | None = None,
+    include_prefix_in_token_count: bool = False,
+    push_to_hub: bool = False,
+    hub_dataset_id: str | None = None,
+    random_seed: int = 42,
+) -> None:
+    """
+    Main method to process and save a dataset with sampled prompt/output token counts.
+
+    :param data: Path or identifier for dataset input.
+    :param output_path: File path to save the processed dataset.
+    :param processor: Tokenizer object or its config.
+    :param config: PreprocessDatasetConfig string or file path.
+    :param processor_args: Optional processor arguments.
+    :param data_args: Optional data loading arguments.
+    :param data_column_mapper: Optional column mapping dictionary.
+    :param short_prompt_strategy: Strategy for handling short prompts.
+    :param pad_char: Character used when padding short prompts.
+    :param concat_delimiter: Delimiter for concatenation strategy.
+    :param include_prefix_in_token_count:
+        Whether to include prefix in prompt token count, simplifying the token counts.
+        When True, prefix trimming is disabled and the prefix is kept as-is. The prefix
+        token count is subtracted from the prompt token budget instead.
+    :param push_to_hub: Whether to push to Hugging Face Hub.
+    :param hub_dataset_id: Dataset ID on Hugging Face Hub.
+    :param random_seed: Seed for random sampling.
+    :raises ValueError: If the output path is invalid or pushing conditions unmet.
+    """
+    builders.process_dataset(
+        data, output_path, processor, config, processor_args, data_args,
+        data_column_mapper, short_prompt_strategy, pad_char, concat_delimiter,
+        include_prefix_in_token_count, push_to_hub, hub_dataset_id, random_seed,
+    )

guidellm/data/loaders.py

@@ -0,0 +1,190 @@
+from __future__ import annotations
+
+import contextlib
+from collections.abc import Callable, Iterator
+from typing import Any, Literal, TypeVar
+
+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
+from guidellm.utils import InfoMixin
+
+__all__ = ["DataLoader", "DatasetsIterator"]
+
+
+DataT = TypeVar("DataT")
+
+
+class DatasetsIterator(TorchIterableDataset[DataT]):
+    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
+        )
+        self.epoch = 0
+
+    def __iter__(self) -> Iterator[DataT]:
+        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, epoch=self.epoch
+            )
+
+    def set_epoch(self, epoch: int):
+        self.epoch = epoch
+
+    def generator(
+        self,
+        max_items: int | None = None,
+        modulus: int | None = None,
+        offset: int | None = None,
+        epoch: int = 0,
+    ) -> Iterator[DataT]:
+        gen_count = 0
+
+        with contextlib.suppress(StopIteration):
+            dataset_iters = []
+            for dataset in self.datasets:
+                if hasattr(dataset, "set_epoch"):
+                    with contextlib.suppress(Exception):
+                        dataset.set_epoch(epoch)
+                dataset_iters.append(iter(dataset))
+
+            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  # type: ignore[misc]
+                except StopIteration:
+                    raise  # Stop iteration when any dataset is exhausted
+                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[DataT], InfoMixin):
+    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[DataT] = DatasetsIterator(
+            data=data,
+            data_args=data_args,
+            data_samples=data_samples,
+            processor_factory=processor_factory,
+            preprocessors=preprocessors,
+            random_seed=random_seed,
+        )
+        self._info: dict[str, Any] = {
+            "data": str(data),
+            "data_args": str(data_args),
+            "data_samples": data_samples,
+            "preprocessors": [
+                preprocessor.__class__.__name__ for preprocessor in preprocessors
+            ],
+            "collator": collator.__class__.__name__,
+            "sampler": str(sampler),
+            "num_workers": num_workers,
+            "random_seed": random_seed,
+        }
+        self.epoch = 0
+
+        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,
+        )
+
+    def __iter__(self):
+        if isinstance(self.dataset, DatasetsIterator):
+            self.dataset.set_epoch(self.epoch)
+        self.epoch += 1
+
+        return super().__iter__()
+
+    @property
+    def info(self) -> dict[str, Any]:
+        return self._info

guidellm/data/preprocessors/__init__.py

@@ -0,0 +1,27 @@
+from .formatters import (
+    GenerativeAudioTranscriptionRequestFormatter,
+    GenerativeAudioTranslationRequestFormatter,
+    GenerativeChatCompletionsRequestFormatter,
+    GenerativeTextCompletionsRequestFormatter,
+    RequestFormatter,
+)
+from .mappers import GenerativeColumnMapper
+from .preprocessor import (
+    DataDependentPreprocessor,
+    DatasetPreprocessor,
+    PreprocessorRegistry,
+)
+
+__all__ = [
+    "ColumnMapper",
+    "ColumnMapperRegistry",
+    "DataDependentPreprocessor",
+    "DatasetPreprocessor",
+    "GenerativeAudioTranscriptionRequestFormatter",
+    "GenerativeAudioTranslationRequestFormatter",
+    "GenerativeChatCompletionsRequestFormatter",
+    "GenerativeColumnMapper",
+    "GenerativeTextCompletionsRequestFormatter",
+    "PreprocessorRegistry",
+    "RequestFormatter",
+]