Create a dataloader
(
break_into_chunks: int,
batch_size: int,
block_size: int,
data_dir: Path,
fabric,
shuffle: bool = True,
seed: int = 12345,
split='train',
)
| 670 | |
| 671 | |
| 672 | def create_dataloader( |
| 673 | break_into_chunks: int, |
| 674 | batch_size: int, |
| 675 | block_size: int, |
| 676 | data_dir: Path, |
| 677 | fabric, |
| 678 | shuffle: bool = True, |
| 679 | seed: int = 12345, |
| 680 | split='train', |
| 681 | ) -> DataLoader: |
| 682 | """Create a dataloader""" |
| 683 | |
| 684 | data_config = TRAIN_DATA_CONFIG if split == 'train' else VAL_DATA_CONFIG |
| 685 | |
| 686 | # Support training without validation split. |
| 687 | if not VAL_DATA_CONFIG or len(VAL_DATA_CONFIG) < 1: |
| 688 | return None |
| 689 | |
| 690 | consolidated_filenames = [] |
| 691 | |
| 692 | for dataset_name in data_config: |
| 693 | dataset_dir = data_dir / dataset_name / split |
| 694 | |
| 695 | if not isdir(dataset_dir): |
| 696 | if split == 'train': |
| 697 | error_msg = f'Directory [{dataset_dir}] for dataset [{dataset_name}] does not exist.' |
| 698 | error_msg += f' Please check. Configuration:\n\n{TRAIN_DATA_CONFIG}' |
| 699 | raise RuntimeError(error_msg) |
| 700 | else: |
| 701 | error_msg = f'Validation folder [{dataset_dir}] missing for dataset[{dataset_name}].' |
| 702 | error_msg += 'Did you misconfigure? Or forget to include? Configuration:\n\n{VAL_DATA_CONFIG}' |
| 703 | fabric.print(error_msg) |
| 704 | |
| 705 | # In our data preparation, files are prefixed as such |
| 706 | prefix = f'{split}_{dataset_name}' |
| 707 | filepath_pattern = f'{dataset_dir}/{prefix}*' |
| 708 | |
| 709 | filenames = sorted(glob.glob(filepath_pattern)) |
| 710 | number_files = len(filenames) |
| 711 | if number_files < 1: |
| 712 | raise RuntimeError( |
| 713 | f'No data found at "{filepath_pattern}". Did you specify the right directory?' |
| 714 | ) |
| 715 | fabric.print(f'Found [{number_files}] files in "{filepath_pattern}".') |
| 716 | |
| 717 | consolidated_filenames.extend(filenames) |
| 718 | |
| 719 | random.seed(seed) |
| 720 | random.shuffle(consolidated_filenames) |
| 721 | dataset = PackedDataset( |
| 722 | consolidated_filenames, |
| 723 | # n_chunks control the buffer size. |
| 724 | # Note that the buffer size also impacts the random shuffle |
| 725 | # (PackedDataset is an IterableDataset. So the shuffle is done by prefetch a buffer and shuffle the buffer) |
| 726 | n_chunks=break_into_chunks, |
| 727 | block_size=block_size, |
| 728 | shuffle=shuffle, |
| 729 | seed=seed + fabric.global_rank, |
no test coverage detected