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

import pydantic
import typer

from . import dataloader_base as base


class DataLoaderConfigArgs(pgml.YamlBaseModel):
    """DataLoader configuration"""

    batch_size: typing.Annotated[
        typing.Optional[int],
        typer.Option(help="Batch size. Maps to 'batch_size'"),
        pydantic.Field(None),
        pgml.ArgRef("batch_size"),
    ]
    """Batch size"""

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

    pin_memory: typing.Annotated[
        typing.Optional[bool],
        typer.Option(help="Pin memory. Maps to 'pin_memory'"),
        pydantic.Field(None),
        pgml.ArgRef("pin_memory"),
    ]
    """Pin memory"""

    persistent_workers: typing.Annotated[
        typing.Optional[bool],
        typer.Option(help="Persistent workers. Maps to 'persistent_workers'"),
        pydantic.Field(None),
        pgml.ArgRef("persistent_workers"),
    ]
    """Persistent workers"""

    prefetch_factor: typing.Annotated[
        typing.Optional[int],
        typer.Option(help="Prefetch factor. Maps to 'prefetch_factor'"),
        pydantic.Field(None),
        pgml.ArgRef("prefetch_factor"),
    ]
    """Prefetch factor"""
