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

guidellm/data/schemas.py

@@ -0,0 +1,175 @@
+from __future__ import annotations
+
+from typing import Literal
+
+from pydantic import ConfigDict, Field, model_validator
+
+from guidellm.schemas import StandardBaseModel
+
+__all__ = [
+    "DataConfig",
+    "DataNotSupportedError",
+    "GenerativeDatasetColumnType",
+    "SyntheticTextDatasetConfig",
+    "SyntheticTextPrefixBucketConfig",
+]
+
+
+GenerativeDatasetColumnType = Literal[
+    "prompt_tokens_count_column",
+    "output_tokens_count_column",
+    "prefix_column",
+    "text_column",
+    "image_column",
+    "video_column",
+    "audio_column",
+]
+
+class DataNotSupportedError(Exception):
+    """
+    Exception raised when the data format is not supported by deserializer or config.
+    """
+
+class DataConfig(StandardBaseModel):
+    """
+    A generic parent class for various configs for the data package
+    that can be passed in as key-value pairs or JSON.
+    """
+
+class PreprocessDatasetConfig(DataConfig):
+
+    prompt_tokens: int = Field(
+        description="The average number of text tokens retained or added to prompts.",
+        gt=0,
+    )
+    prompt_tokens_stdev: int | None = Field(
+        description="The standard deviation of the number of tokens retained in or "
+                    "added to prompts.",
+        gt=0,
+        default=None,
+    )
+    prompt_tokens_min: int | None = Field(
+        description="The minimum number of text tokens retained or added to prompts.",
+        gt=0,
+        default=None,
+    )
+    prompt_tokens_max: int | None = Field(
+        description="The maximum number of text tokens retained or added to prompts.",
+        gt=0,
+        default=None,
+    )
+    output_tokens: int = Field(
+        description="The average number of text tokens retained or added to outputs.",
+        gt=0,
+    )
+    output_tokens_stdev: int | None = Field(
+        description="The standard deviation of the number of tokens retained or "
+                    "added to outputs.",
+        gt=0,
+        default=None,
+    )
+    output_tokens_min: int | None = Field(
+        description="The minimum number of text tokens retained or added to outputs.",
+        gt=0,
+        default=None,
+    )
+    output_tokens_max: int | None = Field(
+        description="The maximum number of text tokens retained or added to outputs.",
+        gt=0,
+        default=None,
+    )
+    prefix_tokens_max: int | None = Field(
+        description="The maximum number of text tokens left in the prefixes.",
+        gt=0,
+        default=None,
+    )
+
+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(DataConfig):
+    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,
+    )
+
+    model_config = ConfigDict(
+        extra="allow",
+    )
+
+    prefix_buckets: list[SyntheticTextPrefixBucketConfig] | None = Field(
+        description="Buckets for the prefix tokens distribution.",
+        default=None,
+    )
+
+
+    @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

guidellm/data/utils/__init__.py

@@ -0,0 +1,6 @@
+from .dataset import DEFAULT_SPLITS, resolve_dataset_split
+
+__all__ = [
+    "DEFAULT_SPLITS",
+    "resolve_dataset_split",
+]

guidellm/data/utils/dataset.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+from typing import Literal
+
+from datasets import Dataset, DatasetDict, IterableDataset, IterableDatasetDict
+
+__all__ = ["DEFAULT_SPLITS", "resolve_dataset_split"]
+
+
+DEFAULT_SPLITS: dict[Literal["train", "calib", "val", "test"], list[str]] = {
+    "train": [
+        "train",
+        "training",
+        "train_set",
+        "training_set",
+        "train_dataset",
+        "training_dataset",
+        "train_data",
+        "training_data",
+        "pretrain",
+        "pretrain_set",
+        "pretrain_dataset",
+        "pretrain_data",
+        "pretraining",
+    ],
+    "calib": [
+        "calibration",
+        "calib",
+        "cal",
+        "calibration_set",
+        "calib_set",
+        "cal_set",
+        "calibration_dataset",
+        "calib_dataset",
+        "cal_set",
+        "calibration_data",
+        "calib_data",
+        "cal_data",
+    ],
+    "val": [
+        "validation",
+        "val",
+        "valid",
+        "validation_set",
+        "val_set",
+        "validation_dataset",
+        "val_dataset",
+        "validation_data",
+        "val_data",
+        "dev",
+        "dev_set",
+        "dev_dataset",
+        "dev_data",
+    ],
+    "test": [
+        "test",
+        "testing",
+        "test_set",
+        "testing_set",
+        "test_dataset",
+        "testing_dataset",
+        "test_data",
+        "testing_data",
+        "eval",
+        "eval_set",
+        "eval_dataset",
+        "eval_data",
+    ],
+}
+
+
+def resolve_dataset_split(
+    dataset: Dataset | IterableDataset | DatasetDict | IterableDatasetDict,
+    split: str | None = None,
+) -> Dataset | IterableDataset:
+    if split is not None and isinstance(dataset, DatasetDict | IterableDatasetDict):
+        if split in dataset:
+            return dataset[split]
+
+        raise ValueError(f"Requested split '{split}' not found in dataset: {dataset}.")
+    elif split is not None:
+        raise ValueError(
+            f"Requested split '{split}' but dataset has no splits: {dataset}."
+        )
+
+    if isinstance(dataset, Dataset | IterableDataset):
+        return dataset
+
+    for _, default_splits in DEFAULT_SPLITS.items():
+        for default_split in default_splits:
+            if default_split in dataset:
+                return dataset[default_split]
+
+    return dataset[list(dataset.keys())[0]]

guidellm/extras/__init__.py

@@ -0,0 +1,4 @@
+"""
+Code that depends on optional dependencies.
+Each submodule should be deferred imported.
+"""

guidellm/extras/audio.py

@@ -0,0 +1,220 @@
+from __future__ import annotations
+
+import base64
+from pathlib import Path
+from typing import Any, Literal
+
+import httpx
+import numpy as np
+import torch
+
+try:
+    from torchcodec import AudioSamples
+    from torchcodec.decoders import AudioDecoder
+    from torchcodec.encoders import AudioEncoder
+except ImportError as e:
+    raise ImportError("Please install guidellm[audio] to use audio features") from e
+
+__all__ = [
+    "encode_audio",
+    "is_url",
+]
+
+
+def is_url(text: Any) -> bool:
+    return isinstance(text, str) and text.startswith(("http://", "https://"))
+
+
+def encode_audio(
+    audio: AudioDecoder
+    | bytes
+    | str
+    | Path
+    | np.ndarray
+    | torch.Tensor
+    | dict[str, Any],
+    b64encode: bool = False,
+    sample_rate: int | None = None,
+    file_name: str = "audio.wav",
+    encode_sample_rate: int = 16000,
+    max_duration: float | None = None,
+    mono: bool = True,
+    audio_format: str = "mp3",
+    bitrate: str = "64k",
+) -> dict[
+    Literal[
+        "type",
+        "audio",
+        "format",
+        "mimetype",
+        "audio_samples",
+        "audio_seconds",
+        "audio_bytes",
+        "file_name",
+    ],
+    str | int | float | bytes | None,
+]:
+    """Decode audio (if necessary) and re-encode to specified format."""
+    samples = _decode_audio(audio, sample_rate=sample_rate, max_duration=max_duration)
+
+    bitrate_val = (
+        int(bitrate.rstrip("k")) * 1000 if bitrate.endswith("k") else int(bitrate)
+    )
+    format_val = audio_format.lower()
+
+    encoded_audio = _encode_audio(
+        samples=samples,
+        resample_rate=encode_sample_rate,
+        bitrate=bitrate_val,
+        audio_format=format_val,
+        mono=mono,
+    )
+
+    return {
+        "type": "audio_base64" if b64encode else "audio_file",
+        "audio": (
+            base64.b64encode(encoded_audio).decode("utf-8")
+            if b64encode
+            else encoded_audio
+        ),
+        "file_name": get_file_name(audio)
+        if isinstance(audio, str | Path)
+        else file_name,
+        "format": audio_format,
+        "mimetype": f"audio/{format_val}",
+        "audio_samples": samples.sample_rate,
+        "audio_seconds": samples.duration_seconds,
+        "audio_bytes": len(encoded_audio),
+    }
+
+
+def _decode_audio(  # noqa: C901, PLR0912
+    audio: AudioDecoder
+    | bytes
+    | str
+    | Path
+    | np.ndarray
+    | torch.Tensor
+    | dict[str, Any],
+    sample_rate: int | None = None,
+    max_duration: float | None = None,
+) -> AudioSamples:
+    """Decode audio from various input types into AudioSamples."""
+    # If input is a dict, unwrap it into a function call
+    if isinstance(audio, dict):
+        sample_rate = audio.get("sample_rate", audio.get("sampling_rate", sample_rate))
+        if "data" not in audio and "url" not in audio:
+            raise ValueError(
+                f"Audio dict must contain either 'data' or 'url' keys, got {audio}"
+            )
+        audio_data = audio["data"] if "data" in audio else audio.get("url")
+        if audio_data is None:
+            raise ValueError(
+                f"Audio dict must contain either 'data' or 'url' keys, got {audio}"
+            )
+        return _decode_audio(
+            audio=audio_data,
+            sample_rate=sample_rate,
+            max_duration=max_duration,
+        )
+
+    # Convert numpy array to torch tensor and re-call
+    if isinstance(audio, np.ndarray):
+        return _decode_audio(
+            audio=torch.from_numpy(audio),
+            sample_rate=sample_rate,
+            max_duration=max_duration,
+        )
+
+    samples: AudioSamples
+
+    data: torch.Tensor | bytes
+    # HF datasets return AudioDecoder for audio column
+    if isinstance(audio, AudioDecoder):
+        samples = audio.get_samples_played_in_range(stop_seconds=max_duration)
+    elif isinstance(audio, torch.Tensor):
+        # If float stream assume decoded audio
+        if torch.is_floating_point(audio):
+            if sample_rate is None:
+                raise ValueError("Sample rate must be set for decoded audio")
+
+            full_duration = audio.shape[1] / sample_rate
+            # If max_duration is set, trim the audio to that duration
+            if max_duration is not None:
+                num_samples = int(max_duration * sample_rate)
+                duration = min(max_duration, full_duration)
+                data = audio[:, :num_samples]
+            else:
+                duration = full_duration
+                data = audio
+
+            samples = AudioSamples(
+                data=data,
+                pts_seconds=0.0,
+                duration_seconds=duration,
+                sample_rate=sample_rate,
+            )
+        # If bytes tensor assume encoded audio
+        elif audio.dtype == torch.uint8:
+            decoder = AudioDecoder(
+                source=audio,
+                sample_rate=sample_rate,
+            )
+            samples = decoder.get_samples_played_in_range(stop_seconds=max_duration)
+
+        else:
+            raise ValueError(f"Unsupported audio type: {type(audio)}")
+
+    # If bytes, assume encoded audio
+    elif isinstance(audio, bytes):
+        decoder = AudioDecoder(
+            source=audio,
+            sample_rate=sample_rate,
+        )
+        samples = decoder.get_samples_played_in_range(stop_seconds=max_duration)
+
+    # If str or Path, assume file path or URL to encoded audio
+    elif isinstance(audio, str | Path):
+        if isinstance(audio, str) and is_url(audio):
+            response = httpx.get(audio)
+            response.raise_for_status()
+            data = response.content
+        else:
+            if not Path(audio).exists():
+                raise ValueError(f"Audio file does not exist: {audio}")
+            data = Path(audio).read_bytes()
+        decoder = AudioDecoder(
+            source=data,
+        )
+        samples = decoder.get_samples_played_in_range(stop_seconds=max_duration)
+    else:
+        raise ValueError(f"Unsupported audio type: {type(audio)}")
+
+    return samples
+
+
+def _encode_audio(
+    samples: AudioSamples,
+    resample_rate: int | None = None,
+    bitrate: int = 64000,
+    audio_format: str = "mp3",
+    mono: bool = True,
+) -> bytes:
+    encoder = AudioEncoder(
+        samples=samples.data,
+        sample_rate=samples.sample_rate,
+    )
+
+    audio_tensor = encoder.to_tensor(
+        format=audio_format,
+        bit_rate=bitrate if audio_format == "mp3" else None,
+        num_channels=1 if mono else None,
+        sample_rate=resample_rate,
+    )
+
+    return audio_tensor.numpy().tobytes()
+
+
+def get_file_name(path: Path | str) -> str:
+    """Get file name from path."""
+    return Path(path).name