guidellm.data

@PreprocessStrategyArgs.register("concatenate")
class ConcatenatePreprocessStrategyArgs(PreprocessStrategyArgs):
    """Concatenate successive short prompts until the target length is reached."""

    kind: Literal["concatenate"] = Field(
        default="concatenate",
        description="Concatenate short prompts until the target length is met.",
    )
    delimiter: str = Field(
        default="",
        description="Delimiter inserted between concatenated prompts.",
        examples=["\\n\\n", " "],
    )

    @field_validator("delimiter", mode="before")
    @classmethod
    def _decode_delimiter_escapes(cls, value: Any) -> Any:
        if not isinstance(value, str):
            return value
        return codecs.decode(value, "unicode_escape")

    def handle_short_prompt(
        self,
        current_prompt: str,
        min_prompt_tokens: int,
        tokenizer: PreTrainedTokenizerBase,
        *,
        dataset_iterator: Iterator[dict[str, Any]] | None = None,
        prompt_column: str | None = None,
    ) -> str | None:
        if dataset_iterator is None or prompt_column is None:
            raise ValueError(
                "concatenate strategy requires dataset_iterator and prompt_column"
            )

        tokens_len = len(tokenizer.encode(current_prompt))
        while tokens_len < min_prompt_tokens:
            try:
                next_row = next(dataset_iterator)
            except StopIteration:
                logger.warning(
                    "Could not concatenate enough prompts to reach minimum "
                    "length, ignoring"
                )
                return None
            current_prompt += self.delimiter + next_row[prompt_column]
            tokens_len = len(tokenizer.encode(current_prompt))
        return current_prompt

class DataArgs(
    PydanticClassRegistryMixin["DataArgs"],
    ABC,
):
    """Base class for data loading and processing argument models.

    This class serves as a base for defining argument models related to data loading
    and processing. It inherits from PydanticClassRegistryMixin to enable automatic
    registration of subclasses, allowing for flexible and extensible data handling
    configurations.

    :cvar schema_discriminator: Field name for polymorphic deserialization
    """

    model_config = standard_model_config()

    schema_discriminator: ClassVar[str] = "kind"

    @classmethod
    def __pydantic_schema_base_type__(cls) -> type[DataArgs]:
        """
        Return base type for polymorphic validation hierarchy.

        :return: Base Profile class for schema validation
        """
        if cls.__name__ == "DataArgs":
            return cls

        return DataArgs

    kind: str = Field(
        description="Type identifier for the data arguments configuration.",
        examples=["text_file", "csv_file"],
    )
    load_kwargs: dict[str, Any] = Field(
        default_factory=dict,
        description=(
            "Additional arguments for data loading. These arguements are "
            "passed to the datasets library when loading the dataset."
        ),
        examples=[{"format": "csv"}],
    )

class DataFinalizerArgs(
    PydanticClassRegistryMixin["DataFinalizerArgs"],
    ABC,
):
    """
    Base class for data finalizer argument models.

    This class serves as a base for defining arguments related to data finalization
    configurations. It inherits from PydanticClassRegistryMixin to enable automatic
    registration of subclasses, allowing for flexible and extensible data finalization
    configurations.

    :cvar schema_discriminator: Field name for polymorphic deserialization
    """

    model_config = standard_model_config()

    schema_discriminator: ClassVar[str] = "kind"

    @classmethod
    def __pydantic_schema_base_type__(cls) -> type[DataFinalizerArgs]:
        """
        Return base type for polymorphic validation hierarchy.

        :return: Base DataFinalizerArgs class for schema validation
        """
        if cls.__name__ == "DataFinalizerArgs":
            return cls

        return DataFinalizerArgs

    kind: str = Field(
        description="Type identifier for the data finalizer arguments.",
        examples=["generative"],
    )

class DataLoaderArgs(
    PydanticClassRegistryMixin["DataLoaderArgs"],
    ABC,
):
    """
    Base class for data loader argument models.

    This class serves as a base for defining argument models related to data loading
    configurations. It inherits from PydanticClassRegistryMixin to enable automatic
    registration of subclasses, allowing for flexible and extensible data loading
    configurations.

    :cvar schema_discriminator: Field name for polymorphic deserialization
    """

    model_config = standard_model_config()

    schema_discriminator: ClassVar[str] = "kind"

    @classmethod
    def __pydantic_schema_base_type__(cls) -> type[DataLoaderArgs]:
        """
        Return base type for polymorphic validation hierarchy.

        :return: Base DataLoaderArgs class for schema validation
        """
        if cls.__name__ == "DataLoaderArgs":
            return cls

        return DataLoaderArgs

    kind: str = Field(
        description="Type identifier for the data loader configuration.",
    )
    samples: int = Field(
        default=-1,
        description=(
            "Number of data samples to generate. If -1, the data loader will "
            "generate indefinitely until the dataset is exhausted."
        ),
    )

class DataLoaderRegistry(Generic[DataT_co], RegistryMixin[type[DataLoader]]):
    @classmethod
    def create(
        cls,
        config: DataLoaderArgs,
        datasets: list[DatasetType],
        preprocessors: list[DatasetPreprocessor],
        finalizer: DatasetFinalizer[DataT_co],
        random_seed: int,
        **kwargs: Any,
    ) -> DataLoader[DataT_co]:
        """
        Factory method to create a DataLoader instance based on provided configuration.

        :param config: A DataEntrypointArgs object containing the configuration.
        """
        kind = config.kind
        data_loader_cls = cls.get_registered_object(kind)

        if data_loader_cls is None:
            raise ValueError(
                f"DataLoader type '{kind}' is not registered."
                f"Available types: {list(cls.registry.keys()) if cls.registry else []}"
            )

        return data_loader_cls(
            config=config,
            datasets=datasets,
            preprocessors=preprocessors,
            finalizer=finalizer,
            random_seed=random_seed,
            **kwargs,
        )

class DataNotSupportedError(Exception):
    """
    Exception raised when the data format is not supported by deserializer or config.
    """

class DataPreprocessorArgs(
    PydanticClassRegistryMixin["DataPreprocessorArgs"],
    ABC,
):
    """
    Base class for data preprocessor argument models.

    This class serves as a base for defining arguments related to data preprocessing
    configurations. It inherits from PydanticClassRegistryMixin to enable automatic
    registration of subclasses, allowing for flexible and extensible data preprocessing
    configurations.

    :cvar schema_discriminator: Field name for polymorphic deserialization
    """

    model_config = standard_model_config()

    schema_discriminator: ClassVar[str] = "kind"

    @classmethod
    def __pydantic_schema_base_type__(cls) -> type[DataPreprocessorArgs]:
        """
        Return base type for polymorphic validation hierarchy.

        :return: Base DataPreprocessorArgs class for schema validation
        """
        if cls.__name__ == "DataPreprocessorArgs":
            return cls

        return DataPreprocessorArgs

    kind: str = Field(
        description="Type identifier for the data preprocessor arguments.",
        examples=["generative_column_mapper", "pooling_column_mapper"],
    )

class DataTokenizerArgs(
    PydanticClassRegistryMixin["DataTokenizerArgs"],
    ABC,
):
    """
    Base class for data tokenizer argument models.

    This class serves as a base for defining arguments related to data tokenization
    configurations. It inherits from PydanticClassRegistryMixin to enable automatic
    registration of subclasses, allowing for flexible and extensible data tokenization
    configurations.

    :cvar schema_discriminator: Field name for polymorphic deserialization
    """

    model_config = standard_model_config()

    schema_discriminator: ClassVar[str] = "kind"

    @classmethod
    def __pydantic_schema_base_type__(cls) -> type[DataTokenizerArgs]:
        """
        Return base type for polymorphic validation hierarchy.

        :return: Base DataTokenizerArgs class for schema validation
        """
        if cls.__name__ == "DataTokenizerArgs":
            return cls

        return DataTokenizerArgs

    kind: str = Field(
        description="Type identifier for the data tokenizer arguments.",
        examples=["huggingface"],
    )
    model: str | None = Field(
        default=None,
        description=(
            "Optional model name or path for the tokenizer. This field can be "
            "used by tokenizer implementations that require a model specification, "
            "such as HuggingFace tokenizers."
        ),
        examples=["gpt2"],
    )

@runtime_checkable
class DatasetFinalizer(Protocol[DataT_co]):
    """
    Protocol for finalizing dataset rows into a desired data type.
    """

    def __init__(self, config: DataFinalizerArgs) -> None: ...

    def __call__(self, items: list[dict[str, Any]]) -> DataT_co: ...

@PreprocessStrategyArgs.register("error")
class ErrorPreprocessStrategyArgs(PreprocessStrategyArgs):
    """Raise an error when a prompt is shorter than the target length."""

    kind: Literal["error"] = Field(
        default="error",
        description="Raise an error when a prompt is shorter than the target length.",
    )

    def handle_short_prompt(
        self,
        current_prompt: str,
        min_prompt_tokens: int,
        tokenizer: PreTrainedTokenizerBase,
        *,
        dataset_iterator: Iterator[dict[str, Any]] | None = None,
        prompt_column: str | None = None,
    ) -> str | None:
        _ = (dataset_iterator, prompt_column)
        prompt_len = len(tokenizer.encode(current_prompt))
        if prompt_len < min_prompt_tokens:
            raise PromptTooShortError(
                f"Found too short prompt: {current_prompt}, with length: {prompt_len}. "
                f"Minimum length required: {min_prompt_tokens}.",
            )
        return current_prompt

class FinalizerRegistry(RegistryMixin[type[DatasetFinalizer]]):
    @classmethod
    def create(cls, config: DataFinalizerArgs) -> DatasetFinalizer:
        """
        Factory method to create a DatasetFinalizer instance based on configuration.

        :param config: A DataFinalizerArgs object containing the configuration.
        """
        kind = config.kind
        finalizer_cls = cls.get_registered_object(kind)

        if finalizer_cls is None:
            raise ValueError(
                f"DatasetFinalizer type '{kind}' is not registered."
                f"Available types: {list(cls.registry.keys()) if cls.registry else []}"
            )

        return finalizer_cls(config)

@FinalizerRegistry.register("generative")
class GenerativeRequestFinalizer(
    DatasetFinalizer[Iterable[tuple[GenerationRequest, RequestSettings]]]
):
    """
    Finalizer that converts dataset rows into GenerationRequest objects,
    aggregating usage metrics from the provided columns.
    """

    def __init__(self, config: GenerativeRequestFinalizerArgs) -> None:
        self.config = config

    def __call__(
        self, items: list[dict[str, Any]]
    ) -> list[tuple[GenerationRequest, RequestSettings]]:
        results: list[tuple[GenerationRequest, RequestSettings]] = []
        for item in items:
            request, settings = self.finalize_turn(item)
            if request.turn_type == "client_tool_call":
                # Split tool-calling turns: the tool_response_column moves
                # to a separate injection turn that follows the tool call.
                tool_response_data = request.columns.pop("tool_response_column", [])
                injection_columns: dict[str, list[Any]] = {}
                if tool_response_data:
                    injection_columns["tool_response_column"] = tool_response_data
                # Move output metrics to next turn
                metrics_config = request.output_metrics
                request.output_metrics = UsageMetrics()
                results.append((request, RequestSettings()))
                results.append(
                    (
                        GenerationRequest(
                            columns=injection_columns,
                            turn_type="tool_response_injection",
                            output_metrics=metrics_config,
                        ),
                        settings,
                    )
                )
            else:
                results.append((request, settings))
        return results

    def finalize_turn(  # noqa: C901 PLR0912
        self, columns: dict[str, Any]
    ) -> tuple[GenerationRequest, RequestSettings]:
        input_metrics = UsageMetrics()
        output_metrics = UsageMetrics()

        # Sum prompt token column
        if prompt_tokens := sum(
            count for count in columns.get("prompt_tokens_count_column", []) if count
        ):
            input_metrics.text_tokens = prompt_tokens

        # Sum output token column
        if output_tokens := sum(
            count for count in columns.get("output_tokens_count_column", []) if count
        ):
            output_metrics.text_tokens = output_tokens

        # Count words in prefixes
        for prefix in columns.get("prefix_column", []):
            if not prefix:
                continue

            input_metrics.add_text_metrics(prefix)

        # Count words in text prompts
        for text in columns.get("text_column", []):
            if not text:
                continue

            input_metrics.add_text_metrics(text)

        # Count pixels and bytes in images
        for image in columns.get("image_column", []):
            if not image:
                continue

            if (image_pixels := image.get("image_pixels")) is not None:
                input_metrics.image_pixels = (
                    input_metrics.image_pixels or 0
                ) + image_pixels
            if (image_bytes := image.get("image_bytes")) is not None:
                input_metrics.image_bytes = (
                    input_metrics.image_bytes or 0
                ) + image_bytes

        # Count frames, seconds, and bytes in videos
        for video in columns.get("video_column", []):
            if not video:
                continue

            if (video_frames := video.get("video_frames")) is not None:
                input_metrics.video_frames = (
                    input_metrics.video_frames or 0
                ) + video_frames
            if (video_seconds := video.get("video_seconds")) is not None:
                input_metrics.video_seconds = (
                    input_metrics.video_seconds or 0.0
                ) + video_seconds
            if (video_bytes := video.get("video_bytes")) is not None:
                input_metrics.video_bytes = (
                    input_metrics.video_bytes or 0
                ) + video_bytes

        # Count samples, seconds, and bytes in audio
        for audio in columns.get("audio_column", []):
            if not audio:
                continue

            if (audio_samples := audio.get("audio_samples")) is not None:
                input_metrics.audio_samples = (
                    input_metrics.audio_samples or 0
                ) + audio_samples
            if (audio_seconds := audio.get("audio_seconds")) is not None:
                input_metrics.audio_seconds = (
                    input_metrics.audio_seconds or 0.0
                ) + audio_seconds
            if (audio_bytes := audio.get("audio_bytes")) is not None:
                input_metrics.audio_bytes = (
                    input_metrics.audio_bytes or 0
                ) + audio_bytes

        # Resolve turn type with priority:
        # 1. Explicit turn_type_column (from synthetic server_tool_call_turns
        #    or hand-crafted datasets)
        # 2. tools_column presence + tool_call_mode config
        # 3. Default to "standard"
        turn_type_values = columns.get("turn_type_column", [])
        turn_type: TurnType
        if turn_type_values and turn_type_values[0]:
            turn_type = turn_type_values[0]
        elif columns.get("tools_column"):
            if self.config.tool_call_mode == "server":
                turn_type = "server_tool_call"
                # Tools are server-managed; strip data-provided definitions
                # so the request handler doesn't inject them.
                columns.pop("tools_column", None)
                columns.pop("tool_response_column", None)
            else:
                turn_type = "client_tool_call"
        else:
            turn_type = "standard"

        return GenerationRequest(
            columns=columns,
            turn_type=turn_type,
            input_metrics=input_metrics,
            output_metrics=output_metrics,
        ), RequestSettings(
            relative_timestamp=self._get_optional_column_value(
                columns, "relative_timestamp_column"
            ),
            requeue_delay=self._get_optional_column_value(
                columns, "requeue_delay_column"
            ),
        )

    @staticmethod
    def _get_optional_column_value(
        columns: dict[str, Any],
        column_name: str,
    ) -> Any | None:
        """
        Retrieve the first value from a specified column in the columns dictionary.

        :param columns: A dictionary containing columns.
        :param column_name: The name of the column to retrieve the value from.
        :return: The first value from the specified column
        """
        values = columns.get(column_name, [])
        return values[0] if values else None

@DataFinalizerArgs.register("generative")
class GenerativeRequestFinalizerArgs(DataFinalizerArgs):
    """Model for generative request finalizer arguments."""

    kind: Literal["generative"] = Field(
        default="generative",
        description="Type identifier for the generative request finalizer.",
    )
    tool_call_mode: Literal["client", "server"] = Field(
        default="client",
        description="How to handle turns with tool definitions. "
        "'client' (default) creates client_tool_call + injection turns. "
        "'server' creates server_tool_call turns (no injection, "
        "tools are server-managed).",
    )

@PreprocessStrategyArgs.register("ignore")
class IgnorePreprocessStrategyArgs(PreprocessStrategyArgs):
    """Skip prompts that are shorter than the target token length."""

    kind: Literal["ignore"] = Field(
        default="ignore",
        description="Skip short prompts instead of extending them.",
    )

    def handle_short_prompt(
        self,
        current_prompt: str,
        min_prompt_tokens: int,
        tokenizer: PreTrainedTokenizerBase,
        *,
        dataset_iterator: Iterator[dict[str, Any]] | None = None,
        prompt_column: str | None = None,
    ) -> str | None:
        _ = (dataset_iterator, prompt_column)
        if len(tokenizer.encode(current_prompt)) < min_prompt_tokens:
            logger.warning("Prompt too short, ignoring")
            return None
        return current_prompt

@PreprocessStrategyArgs.register("pad")
class PadPreprocessStrategyArgs(PreprocessStrategyArgs):
    """Pad short prompts with a character until the target length is reached."""

    kind: Literal["pad"] = Field(
        default="pad",
        description="Pad short prompts with a character until the target length.",
    )
    pad: str = Field(
        default="",
        description="Character or string used to pad short prompts.",
        examples=[" ", "X"],
    )

    @field_validator("pad", mode="before")
    @classmethod
    def _decode_pad_escapes(cls, value: Any) -> Any:
        if not isinstance(value, str):
            return value
        return codecs.decode(value, "unicode_escape")

    def handle_short_prompt(
        self,
        current_prompt: str,
        min_prompt_tokens: int,
        tokenizer: PreTrainedTokenizerBase,
        *,
        dataset_iterator: Iterator[dict[str, Any]] | None = None,
        prompt_column: str | None = None,
    ) -> str | None:
        _ = (dataset_iterator, prompt_column)
        tokens = tokenizer.encode(current_prompt)
        pad_count = 1
        pad_multiplier = 2
        prompt = current_prompt
        while len(tokens) < min_prompt_tokens:
            prompt += self.pad * pad_count
            tokens = tokenizer.encode(prompt)
            pad_count *= pad_multiplier
        return prompt

class PreprocessStrategyArgs(
    PydanticClassRegistryMixin["PreprocessStrategyArgs"],
    ABC,
):
    """
    Base class for dataset preprocessing strategy configurations.

    Combines target token-count settings with a short-prompt handling strategy
    selected by ``kind``. Subclasses encapsulate strategy-specific options and
    short-prompt handling logic.
    """

    model_config = standard_model_config()

    schema_discriminator: ClassVar[str] = "kind"

    @classmethod
    def __pydantic_schema_base_type__(cls) -> type[PreprocessStrategyArgs]:
        """
        Return base type for polymorphic validation hierarchy.

        :return: Base PreprocessStrategyArgs class for schema validation
        """
        if cls.__name__ == "PreprocessStrategyArgs":
            return cls

        return PreprocessStrategyArgs

    kind: str = Field(
        description="Short-prompt strategy kind for dataset preprocessing.",
        examples=["ignore", "concatenate", "pad", "error"],
    )
    prompt_tokens: int = Field(
        description="The average number of text tokens retained or added to prompts.",
        examples=[100],
        gt=0,
    )
    prompt_tokens_stdev: int | None = Field(
        description=(
            "The standard deviation of the number of tokens retained in or "
            "added to prompts."
        ),
        examples=[10],
        gt=0,
        default=None,
    )
    prompt_tokens_min: int | None = Field(
        description="The minimum number of text tokens retained or added to prompts.",
        examples=[100],
        gt=0,
        default=None,
    )
    prompt_tokens_max: int | None = Field(
        description="The maximum number of text tokens retained or added to prompts.",
        examples=[100],
        gt=0,
        default=None,
    )
    output_tokens: int = Field(
        description="The average number of text tokens retained or added to outputs.",
        examples=[100],
        gt=0,
    )
    output_tokens_stdev: int | None = Field(
        description=(
            "The standard deviation of the number of tokens retained or "
            "added to outputs."
        ),
        examples=[10],
        gt=0,
        default=None,
    )
    output_tokens_min: int | None = Field(
        description="The minimum number of text tokens retained or added to outputs.",
        examples=[100],
        gt=0,
        default=None,
    )
    output_tokens_max: int | None = Field(
        description="The maximum number of text tokens retained or added to outputs.",
        examples=[100],
        gt=0,
        default=None,
    )
    prefix_tokens_max: int | None = Field(
        description="The maximum number of text tokens left in the prefixes.",
        gt=0,
        examples=[100],
        default=None,
    )
    count_prefix: bool = Field(
        default=False,
        description=(
            "When True, include prefix tokens in the prompt token budget. "
            "Prefix trimming via prefix_tokens_max still applies when set."
        ),
    )

    @abstractmethod
    def handle_short_prompt(
        self,
        current_prompt: str,
        min_prompt_tokens: int,
        tokenizer: PreTrainedTokenizerBase,
        *,
        dataset_iterator: Iterator[dict[str, Any]] | None = None,
        prompt_column: str | None = None,
    ) -> str | None:
        """
        Apply this strategy when a prompt is shorter than the target length.

        :param current_prompt: The input prompt string.
        :param min_prompt_tokens: Minimum required token count.
        :param tokenizer: Tokenizer used to count tokens.
        :param dataset_iterator: Optional iterator for strategies that consume
            additional rows (e.g. concatenate).
        :param prompt_column: Optional prompt column name for row extraction.
        :return: Adjusted prompt text, or None to skip the row.
        """

class PreprocessorRegistry(
    RegistryMixin[type[DatasetPreprocessor] | type[DataDependentPreprocessor]]
):
    @classmethod
    def create(cls, config: DataPreprocessorArgs) -> DatasetPreprocessor:
        """
        Factory method to create a DatasetPreprocessor instance based on configuration.

        :param config: A DataPreprocessorArgs object containing the configuration.
        """
        kind = config.kind
        preprocessor_cls = cls.get_registered_object(kind)

        if preprocessor_cls is None:
            raise ValueError(
                f"DatasetPreprocessor type '{kind}' is not registered."
                f"Available types: {list(cls.registry.keys()) if cls.registry else []}"
            )

        return preprocessor_cls(config)

class PromptTooShortError(Exception):
    """Raised when a prompt is shorter than the required token length."""

@DataLoaderArgs.register("pytorch")
class TorchDataLoaderArgs(DataLoaderArgs):
    """Model for PyTorch data loader arguments."""

    kind: Literal["pytorch"] = Field(  # type: ignore[assignment]
        default="pytorch",
        description="Type identifier for the generative data loader.",
    )
    shuffle: bool = Field(
        default=False,
        description="Shuffle data rows at every epoch.",
    )
    num_workers: int = Field(
        default=1,
        description=(
            "Number of worker processes for data loading. If 0, data loading "
            "will be performed in the main process."
        ),
    )