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


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

    num_layers: int
    """Number of layers"""


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

    num_epochs: int
    """Number of epochs"""


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

    path: str
    """Path to the dataset"""


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

    train_dataset: Dataset
    """Path to the dataset"""

    test_dataset: Dataset
    """Path to the dataset"""

    num_workers: int
    """Number of workers for loading the dataset"""


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

    data: Data
    """Dataset configuration"""

    model: Model
    """Model configuration"""

    training: Training
    """Training configuration"""
