guidellm.data.loaders

class DataLoaderRegistry(Generic[DataT_co], RegistryMixin[type[DataLoader]]):
    @classmethod
    def create(
        cls,
        config: DataLoaderArgs,
        datasets: list[DatasetType],
        preprocessors: list[DatasetPreprocessor],
        finalizer: DatasetFinalizer[DataT_co],
        random_seed: int,
        **kwargs: Any,
    ) -> DataLoader[DataT_co]:
        """
        Factory method to create a DataLoader instance based on provided configuration.

        :param config: A DataEntrypointArgs object containing the configuration.
        """
        kind = config.kind
        data_loader_cls = cls.get_registered_object(kind)

        if data_loader_cls is None:
            raise ValueError(
                f"DataLoader type '{kind}' is not registered."
                f"Available types: {list(cls.registry.keys()) if cls.registry else []}"
            )

        return data_loader_cls(
            config=config,
            datasets=datasets,
            preprocessors=preprocessors,
            finalizer=finalizer,
            random_seed=random_seed,
            **kwargs,
        )

@DataLoaderArgs.register("pytorch")
class TorchDataLoaderArgs(DataLoaderArgs):
    """Model for PyTorch data loader arguments."""

    kind: Literal["pytorch"] = Field(  # type: ignore[assignment]
        default="pytorch",
        description="Type identifier for the generative data loader.",
    )
    shuffle: bool = Field(
        default=False,
        description="Shuffle data rows at every epoch.",
    )
    num_workers: int = Field(
        default=1,
        description=(
            "Number of worker processes for data loading. If 0, data loading "
            "will be performed in the main process."
        ),
    )