(
self,
dataset: Dataset,
name: str,
default_batch_size: int,
total_epochs: int = 1,
offset: int = 0,
drop_last: bool = True,
total_steps: Optional[int] = None,
enable_progress_bar: Optional[bool] = True,
)
| 23 | |
| 24 | class _HFBatchReader: |
| 25 | def __init__( |
| 26 | self, |
| 27 | dataset: Dataset, |
| 28 | name: str, |
| 29 | default_batch_size: int, |
| 30 | total_epochs: int = 1, |
| 31 | offset: int = 0, |
| 32 | drop_last: bool = True, |
| 33 | total_steps: Optional[int] = None, |
| 34 | enable_progress_bar: Optional[bool] = True, |
| 35 | ): |
| 36 | self.dataset = dataset |
| 37 | self.dataset_size = len(dataset) |
| 38 | if self.dataset_size == 0: |
| 39 | raise ValueError(f"Dataset [{name}] is empty and cannot be read in batches.") |
| 40 | self.name = name |
| 41 | self.current_batch_size = None |
| 42 | self.drop_last = drop_last |
| 43 | |
| 44 | self.current_offset = offset |
| 45 | |
| 46 | # convert epochs/steps to sample number |
| 47 | if total_steps is not None: |
| 48 | self.total_samples = default_batch_size * total_steps |
| 49 | else: |
| 50 | self.total_samples = self.dataset_size * total_epochs |
| 51 | |
| 52 | if enable_progress_bar: |
| 53 | from ray.experimental.tqdm_ray import tqdm |
| 54 | |
| 55 | self.progress_bar = tqdm( |
| 56 | total=self.total_samples, |
| 57 | desc=f"Dataset [{self.name}] Progressing", |
| 58 | ) |
| 59 | else: |
| 60 | self.progress_bar = DummyProgressBar() |
| 61 | |
| 62 | self.progress_bar.update(self.current_offset) |
| 63 | |
| 64 | def read_batch(self, batch_size: int) -> Tuple[List, List]: |
| 65 | batch, indices = [], [] |
nothing calls this directly
no test coverage detected