Function create_dataloaders

pretrain/main.py:741–775

This calls `create_dataloader` twice, one for train, another for validation.

(
    break_into_chunks: int,
    batch_size: int,
    block_size: int,
    fabric,
    data_dir: Path,
    seed: int = 12345,
)

def create_dataloaders(
    break_into_chunks: int,
    batch_size: int,
    block_size: int,
    fabric,
    data_dir: Path,
    seed: int = 12345,
) -> Tuple[DataLoader, DataLoader]:
    """This calls `create_dataloader` twice, one for train, another for validation."""

    # Increase by one because we need the next word as well
    effective_block_size = block_size + 1
    train_dataloader = create_dataloader(
        break_into_chunks=break_into_chunks,
        batch_size=batch_size,
        block_size=effective_block_size,
        fabric=fabric,
        data_dir=data_dir,
        shuffle=True,
        seed=seed,
        split='train',
    )

    val_dataloader = create_dataloader(
        break_into_chunks=break_into_chunks,
        batch_size=batch_size,
        block_size=effective_block_size,
        fabric=fabric,
        data_dir=data_dir,
        shuffle=False,
        seed=seed,
        split='validation',
    )

    return train_dataloader, val_dataloader


# learning rate decay scheduler (cosine with warmup)