# Autogenerated code. DO NOT EDIT.
import py_gen_ml as pgml

from . import config_base as base


class OptimizerPatch(pgml.YamlBaseModel):
    """Optimizer configuration"""

    learning_rate: float | None = None
    """Learning rate"""

    beta1: float | None = None
    """Decay rate"""


class DataPatch(pgml.YamlBaseModel):
    """Data configuration"""

    batch_size: int | None = None
    """Batch size for a single GPU"""

    num_epochs: int | None = None
    """Number of epochs to train"""


class ConvBlockPatch(pgml.YamlBaseModel):
    """Convolutional layer configuration"""

    out_channels: int | None = None
    """Number of output channels"""

    kernel_size: int | None = None
    """Square kernel size"""

    pool_size: int | None = None
    """Square pool size"""

    activation: base.Activation | None = None
    """Activation function"""


class LinearBlockPatch(pgml.YamlBaseModel):
    """Linear layer configuration"""

    out_features: int | None = None
    """Number of output features"""

    activation: base.Activation | None = None
    """Activation function"""


class ConvNetPatch(pgml.YamlBaseModel):
    """Convolutional neural network configuration"""

    block: ConvBlockPatch | None = None
    """Conv layer configuration"""

    num_layers: int | None = None
    """Number of layers"""


class MLPPatch(pgml.YamlBaseModel):
    """Multi-layer perceptron configuration"""

    block: LinearBlockPatch | None = None
    """Linear layer configuration"""

    num_layers: int | None = None
    """Number of layers"""


class ModelPatch(pgml.YamlBaseModel):
    """Model configuration"""

    conv_net: ConvNetPatch | None = None
    """Conv blocks"""

    head: MLPPatch | None = None
    """MLP head"""


class ProjectPatch(pgml.YamlBaseModel):
    """Global configuration"""

    net: ModelPatch | None = None
    """Model configuration"""

    optimizer: OptimizerPatch | None = None
    """Optimizer configuration"""

    data: DataPatch | None = None
    """Data configuration"""
