Function create_dataloaders

bin/finetune_example/posttrain_dataloader.py:288–319

(
    train_datasets: str,
    validation_datasets: str,
    batch_size: int,
    device,  # for debug
    infinite_train: bool = False,
    num_workers: int = 0,
)

def create_dataloaders(
    train_datasets: str,
    validation_datasets: str,
    batch_size: int,
    device,  # for debug
    infinite_train: bool = False,
    num_workers: int = 0,
):

    trainset = TokenizedDataset(dataset_dir=train_datasets, device=device)
    valset = TokenizedDataset(dataset_dir=validation_datasets, device=device)
    trainloader = DataLoader(
        trainset,
        # batch_sampler=trainsampler,
        batch_size=batch_size,
        collate_fn=collate_fn,
        num_workers=12,
        pin_memory=True,
        shuffle=True,
    )

    valloader = DataLoader(
        valset,
        # batch_sampler=valsampler,
        batch_size=batch_size,
        collate_fn=collate_fn,
        num_workers=num_workers,
        pin_memory=True,
        shuffle=False,
    )

    return trainloader, valloader