guidellm 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

guidellm/dataset/in_memory.py

@@ -0,0 +1,132 @@
+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)}

guidellm/dataset/synthetic.py

@@ -0,0 +1,262 @@
+import json
+import random
+from collections.abc import Iterable, Iterator
+from pathlib import Path
+from typing import Any, Literal, Optional, Union
+
+import yaml
+from datasets import (
+    Dataset,
+    DatasetDict,
+    IterableDataset,
+    IterableDatasetDict,
+)
+from pydantic import BaseModel, Field
+from transformers import PreTrainedTokenizerBase  # type: ignore[import]
+
+from guidellm.dataset.creator import ColumnInputTypes, DatasetCreator
+from guidellm.utils import EndlessTextCreator, IntegerRangeSampler, check_load_processor
+
+__all__ = [
+    "SyntheticDatasetCreator",
+    "SyntheticDatasetConfig",
+    "SyntheticTextItemsGenerator",
+]
+
+
+class SyntheticDatasetConfig(BaseModel):
+    prompt_tokens: int = Field(
+        description="The average number of text tokens generated for prompts.",
+        gt=0,
+    )
+    prompt_tokens_stdev: Optional[int] = Field(
+        description="The standard deviation of the tokens generated for prompts.",
+        gt=0,
+        default=None,
+    )
+    prompt_tokens_min: Optional[int] = Field(
+        description="The minimum number of text tokens generated for prompts.",
+        gt=0,
+        default=None,
+    )
+    prompt_tokens_max: Optional[int] = 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: Optional[int] = Field(
+        description="The standard deviation of the tokens generated for outputs.",
+        gt=0,
+        default=None,
+    )
+    output_tokens_min: Optional[int] = Field(
+        description="The minimum number of text tokens generated for outputs.",
+        gt=0,
+        default=None,
+    )
+    output_tokens_max: Optional[int] = Field(
+        description="The maximum number of text tokens generated for outputs.",
+        gt=0,
+        default=None,
+    )
+    samples: int = Field(
+        description="The number of samples to generate for the dataset.",
+        gt=0,
+        default=1000,
+    )
+    source: str = Field(
+        description="The source of the text data to be used for generation.",
+        default="data:prideandprejudice.txt.gz",
+    )
+
+    @staticmethod
+    def parse_str(data: Union[str, Path]) -> "SyntheticDatasetConfig":
+        if (
+            isinstance(data, Path)
+            or data.strip().endswith(".config")
+            or data.strip().endswith(".yaml")
+        ):
+            return SyntheticDatasetConfig.parse_config_file(data)
+
+        if data.strip().startswith("{"):
+            return SyntheticDatasetConfig.parse_json(data)
+
+        if data.count("=") > 1:
+            return SyntheticDatasetConfig.parse_key_value_pairs(data)
+
+        raise ValueError(
+            f"Unsupported data format. Expected JSON or key-value pairs, got {data}"
+        )
+
+    @staticmethod
+    def parse_json(data: str) -> "SyntheticDatasetConfig":
+        config_dict = json.loads(data.strip())
+
+        return SyntheticDatasetConfig(**config_dict)
+
+    @staticmethod
+    def parse_key_value_pairs(data: str) -> "SyntheticDatasetConfig":
+        config_dict = {}
+        items = data.strip().split(",")
+        for item in items:
+            key, value = item.split("=")
+            config_dict[key.strip()] = (
+                int(value.strip()) if value.strip().isnumeric() else value.strip()
+            )
+
+        return SyntheticDatasetConfig(**config_dict)  # type: ignore[arg-type]
+
+    @staticmethod
+    def parse_config_file(data: Union[str, Path]) -> "SyntheticDatasetConfig":
+        with Path(data).open("r") as file:
+            config_dict = yaml.safe_load(file)
+
+        return SyntheticDatasetConfig(**config_dict)
+
+
+class SyntheticTextItemsGenerator(
+    Iterable[
+        dict[
+            Literal["prompt", "prompt_tokens_count", "output_tokens_count"],
+            Union[str, int],
+        ]
+    ]
+):
+    def __init__(
+        self,
+        config: SyntheticDatasetConfig,
+        processor: PreTrainedTokenizerBase,
+        random_seed: int,
+    ):
+        self.config = config
+        self.processor = processor
+        self.random_seed = random_seed
+        self.text_creator = EndlessTextCreator(
+            data=config.source,
+        )
+
+    def __iter__(
+        self,
+    ) -> Iterator[
+        dict[
+            Literal["prompt", "prompt_tokens_count", "output_tokens_count"],
+            Union[str, int],
+        ]
+    ]:
+        prompt_tokens_sampler = 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 = 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
+        )
+        # ensure diff distribution from output tokens
+        rand = random.Random(self.random_seed + 2)  # noqa: S311
+
+        for _, prompt_tokens, output_tokens in zip(
+            range(self.config.samples),
+            prompt_tokens_sampler,
+            output_tokens_sampler,
+        ):
+            start_index = rand.randint(0, len(self.text_creator.words))
+            yield {
+                "prompt": self._create_prompt(prompt_tokens, start_index),
+                "prompt_tokens_count": prompt_tokens,
+                "output_tokens_count": output_tokens,
+            }
+
+    def _create_prompt(self, prompt_tokens: int, start_index: int) -> str:
+        if prompt_tokens <= 0:
+            return ""
+
+        left = start_index
+        right = start_index + 4 * prompt_tokens
+
+        while left < right:
+            mid = (left + right) // 2
+            test_prompt = self.text_creator.create_text(start_index, mid - start_index)
+            test_tokens = len(self.processor.tokenize(test_prompt))
+
+            if test_tokens == prompt_tokens:
+                return test_prompt
+            elif test_tokens < prompt_tokens:
+                left = mid + 1
+            else:
+                right = mid
+
+        return self.text_creator.create_text(start_index, left - start_index)
+
+
+class SyntheticDatasetCreator(DatasetCreator):
+    @classmethod
+    def is_supported(cls, data: Any, data_args: Optional[dict[str, Any]]) -> bool:  # noqa: ARG003
+        if (
+            isinstance(data, Path)
+            and data.exists()
+            and data.suffix in {".config", ".yaml"}
+        ):
+            return True
+
+        if isinstance(data, str):
+            data_str: str = data.strip()
+            if (
+                data_str.startswith("{")
+                or data_str.count("=") > 1
+                or data_str.endswith((".config", ".yaml"))
+            ):
+                return True
+
+        return False
+
+    @classmethod
+    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]:
+        processor = check_load_processor(
+            processor,
+            processor_args,
+            error_msg=(
+                "Processor/tokenizer required for synthetic dataset generation."
+            ),
+        )
+
+        config = SyntheticDatasetConfig.parse_str(data)
+        generator = SyntheticTextItemsGenerator(config, processor, random_seed)
+        items = list(generator)
+
+        return Dataset.from_list(items, **(data_args or {}))
+
+    @classmethod
+    def extract_args_column_mappings(
+        cls,
+        data_args: Optional[dict[str, Any]],
+    ) -> dict[ColumnInputTypes, str]:
+        data_args_columns = super().extract_args_column_mappings(data_args)
+
+        if data_args_columns:
+            raise ValueError(
+                f"Column mappings are not supported for synthetic datasets. "
+                f"Got {data_args_columns}"
+            )
+
+        return {
+            "prompt_column": "prompt",
+            "prompt_tokens_count_column": "prompt_tokens_count",
+            "output_tokens_count_column": "output_tokens_count",
+        }

guidellm/objects/__init__.py

@@ -0,0 +1,18 @@
+from .pydantic import StandardBaseModel, StatusBreakdown
+from .statistics import (
+    DistributionSummary,
+    Percentiles,
+    RunningStats,
+    StatusDistributionSummary,
+    TimeRunningStats,
+)
+
+__all__ = [
+    "StandardBaseModel",
+    "StatusBreakdown",
+    "DistributionSummary",
+    "Percentiles",
+    "RunningStats",
+    "StatusDistributionSummary",
+    "TimeRunningStats",
+]

guidellm/objects/pydantic.py

@@ -0,0 +1,60 @@
+from typing import Any, Generic, TypeVar
+
+from loguru import logger
+from pydantic import BaseModel, ConfigDict, Field
+
+__all__ = ["StandardBaseModel", "StatusBreakdown"]
+
+
+class StandardBaseModel(BaseModel):
+    """
+    A base class for Pydantic models throughout GuideLLM enabling standard
+    configuration and logging.
+    """
+
+    model_config = ConfigDict(
+        extra="ignore",
+        use_enum_values=True,
+        validate_assignment=True,
+        from_attributes=True,
+    )
+
+    def __init__(self, /, **data: Any) -> None:
+        super().__init__(**data)
+        logger.debug(
+            "Initialized new instance of {} with data: {}",
+            self.__class__.__name__,
+            data,
+        )
+
+
+SuccessfulT = TypeVar("SuccessfulT")
+ErroredT = TypeVar("ErroredT")
+IncompleteT = TypeVar("IncompleteT")
+TotalT = TypeVar("TotalT")
+
+
+class StatusBreakdown(BaseModel, Generic[SuccessfulT, ErroredT, IncompleteT, TotalT]):
+    """
+    A base class for Pydantic models that are separated by statuses including
+    successful, incomplete, and errored. It additionally enables the inclusion
+    of total, which is intended as the combination of all statuses.
+    Total may or may not be used depending on if it duplicates information.
+    """
+
+    successful: SuccessfulT = Field(
+        description="The results with a successful status.",
+        default=None,  # type: ignore[assignment]
+    )
+    errored: ErroredT = Field(
+        description="The results with an errored status.",
+        default=None,  # type: ignore[assignment]
+    )
+    incomplete: IncompleteT = Field(
+        description="The results with an incomplete status.",
+        default=None,  # type: ignore[assignment]
+    )
+    total: TotalT = Field(
+        description="The combination of all statuses.",
+        default=None,  # type: ignore[assignment]
+    )