guidellm.data

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: ...

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[GenerationRequest]]):
    """
    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[GenerationRequest]:
        return [self.finalize_turn(item) for item in items]

    def finalize_turn(  # noqa: C901 PLR0912
        self, columns: dict[str, Any]
    ) -> GenerationRequest:
        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

        # A turn expects a tool call if it has tool definitions.
        # Which turns carry tools_column is controlled by the data pipeline
        # (synthetic generator or dataset columns).
        expects_tool_call = bool(columns.get("tools_column"))

        return GenerationRequest(
            columns=columns,
            expects_tool_call=expects_tool_call,
            input_metrics=input_metrics,
            output_metrics=output_metrics,
            settings=self._request_settings_from_columns(columns),
        )

    def _request_settings_from_columns(
        self, columns: dict[str, Any]
    ) -> RequestSettings:
        relative_values = columns.get("relative_timestamp_column", [])
        if relative_values and relative_values[0] is not None:
            return RequestSettings(relative_timestamp=float(relative_values[0]))
        return RequestSettings()

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

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)

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