| 313 | |
| 314 | |
| 315 | def build_dataloader(dataset, batch_size) -> DataLoader: |
| 316 | # Multiple workers is only supported on linux machines |
| 317 | if "linux" in platform.platform().lower(): |
| 318 | num_workers = min(64, dataset.hf_dataset.n_shards) # type: ignore |
| 319 | else: |
| 320 | num_workers = 0 |
| 321 | |
| 322 | # If using multiple workers, configure each worker to prefetch as many samples as it can, up to |
| 323 | # the aggregate device batch size |
| 324 | # If not using workers, the torch DataLoader expects the default value for prefetch_factor, |
| 325 | # which non-intuitively must be 2. |
| 326 | prefetch_factor = max(1, 2 * batch_size // num_workers) if num_workers > 0 else 2 |
| 327 | |
| 328 | return DataLoader( |
| 329 | dataset=dataset, |
| 330 | sampler=None, |
| 331 | batch_size=batch_size, |
| 332 | num_workers=num_workers, |
| 333 | prefetch_factor=prefetch_factor, |
| 334 | ) |
| 335 | |
| 336 | |
| 337 | def generate_samples(loader: DataLoader, truncate_num_samples: Optional[int] = None) -> Iterable[Dict[str, bytes]]: |