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

import pydantic
import typer

from . import cli_demo_base as base


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

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



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

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



class DatasetArgs(pgml.YamlBaseModel):
    """Dataset configuration"""

    path: typing.Annotated[typing.Optional[str], typer.Option(help="Path to the dataset. Maps to 'path'"), pydantic.Field(None), pgml.ArgRef("path")]
    """Path to the dataset"""



class DataArgs(pgml.YamlBaseModel):
    """Data config"""

    num_workers: typing.Annotated[typing.Optional[int], typer.Option(help="Number of workers for loading the dataset. Maps to 'num_workers'"), pydantic.Field(None), pgml.ArgRef("num_workers")]
    """Number of workers for loading the dataset"""

    path: typing.Annotated[typing.Optional[str], typer.Option(help="Path to the dataset. Maps to 'dataset.path'"), pydantic.Field(None), pgml.ArgRef("dataset.path")]
    """Path to the dataset"""



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

    num_workers: typing.Annotated[typing.Optional[int], typer.Option(help="Number of workers for loading the dataset. Maps to 'data.num_workers'"), pydantic.Field(None), pgml.ArgRef("data.num_workers")]
    """Number of workers for loading the dataset"""

    path: typing.Annotated[typing.Optional[str], typer.Option(help="Path to the dataset. Maps to 'data.dataset.path'"), pydantic.Field(None), pgml.ArgRef("data.dataset.path")]
    """Path to the dataset"""

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

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


