guidellm.data.schemas

@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 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"],
    )

@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

@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 PromptTooShortError(Exception):
    """Raised when a prompt is shorter than the required token length."""