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

import pydantic
import typer

from . import advanced_base as base


class LinearBlockArgs(pgml.YamlBaseModel):
    """Linear block configuration"""

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

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



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

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

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



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

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

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



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

    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"""

    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"""


