Split dataset into batches of exactly block_size length
(self, dataset: Dataset, tokens_per_batch: int)
| 46 | } |
| 47 | |
| 48 | def batchify(self, dataset: Dataset, tokens_per_batch: int) -> Iterator[List[torch.Tensor]]: |
| 49 | """Split dataset into batches of exactly block_size length""" |
| 50 | current_tokens = [] # Buffer to store all tokens |
| 51 | |
| 52 | for sentence in dataset: |
| 53 | # Convert input_ids to list and add to buffer |
| 54 | tokens = sentence['input_ids'].tolist() if torch.is_tensor(sentence['input_ids']) else list(sentence['input_ids']) |
| 55 | if not tokens: |
| 56 | continue |
| 57 | current_tokens.extend(tokens) |
| 58 | |
| 59 | # When we have enough tokens, yield batches |
| 60 | while len(current_tokens) >= self.block_size * self.batch_size: |
| 61 | batch = [] |
| 62 | for _ in range(self.batch_size): |
| 63 | # Extract exactly block_size tokens |
| 64 | batch.append(torch.tensor(current_tokens[:self.block_size], dtype=torch.long)) |
| 65 | current_tokens = current_tokens[self.block_size:] |
| 66 | yield batch |
| 67 | |
| 68 | # Handle remaining tokens if they form complete blocks |
| 69 | if len(current_tokens) >= self.block_size: |
| 70 | remaining_batches = len(current_tokens) // self.block_size |
| 71 | remaining_batches = min(remaining_batches, self.batch_size) |
| 72 | if remaining_batches > 0: |
| 73 | batch = [] |
| 74 | for _ in range(remaining_batches): |
| 75 | batch.append(torch.tensor(current_tokens[:self.block_size], dtype=torch.long)) |
| 76 | current_tokens = current_tokens[self.block_size:] |
| 77 | yield batch |
| 78 | |
| 79 | def process_batch(self, batch: List[torch.Tensor]) -> Dict[str, torch.Tensor]: |
| 80 | """Process a single batch of data""" |