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

import pydantic
import typer

from . import advanced_base as base


class TrainingArgs(pgml.YamlBaseModel):
    """Training configuration"""

    num_units: typing.Annotated[
        typing.Optional[int],
        typer.Option(help="Number of units. Maps to 'mlp.layers.num_units'"),
        pydantic.Field(None),
        pgml.ArgRef("mlp.layers.num_units"),
    ]
    """Number of units"""

    activation: typing.Annotated[
        typing.Optional[str],
        typer.Option(help="Activation function. Maps to 'mlp.layers.activation'"),
        pydantic.Field(None),
        pgml.ArgRef("mlp.layers.activation"),
    ]
    """Activation function"""

    type: typing.Annotated[
        typing.Optional[str],
        typer.Option(help="Type of optimizer. Maps to 'optimizer.type'"),
        pydantic.Field(None),
        pgml.ArgRef("optimizer.type"),
    ]
    """Type of optimizer"""

    learning_rate: typing.Annotated[
        typing.Optional[float],
        typer.Option(help="Learning rate. Maps to 'optimizer.learning_rate'"),
        pydantic.Field(None),
        pgml.ArgRef("optimizer.learning_rate"),
    ]
    """Learning rate"""
