Process a single batch of data
(self, batch: List[torch.Tensor])
| 77 | yield batch |
| 78 | |
| 79 | def process_batch(self, batch: List[torch.Tensor]) -> Dict[str, torch.Tensor]: |
| 80 | """Process a single batch of data""" |
| 81 | # Stack the tensors - no need for padding since all sequences are block_size |
| 82 | input_ids = torch.stack(batch).to(self.device) |
| 83 | |
| 84 | # Calculate number of blocks for each sequence |
| 85 | blocks = [ |
| 86 | (self.block_size-1)//self.bucket_size |
| 87 | for _ in range(input_ids.shape[0]) |
| 88 | ] |
| 89 | |
| 90 | # Prepare labels |
| 91 | labels = input_ids.clone() |
| 92 | |
| 93 | # Forward pass |
| 94 | outputs = self.model(input_ids, labels=labels) |
| 95 | |
| 96 | # Calculate next token prediction labels |
| 97 | next_token_labels = torch.cat(( |
| 98 | input_ids[..., 1:], |
| 99 | torch.full_like(input_ids[:, :1], self.tokenizer.eos_token_id) |
| 100 | ), -1) |
| 101 | |
| 102 | # Calculate negative log likelihood |
| 103 | nlls = (-outputs['logits'].log_softmax(-1)).gather(-1, next_token_labels.unsqueeze(-1)).squeeze(-1) |
| 104 | |
| 105 | return { |
| 106 | 'input_ids': input_ids, |
| 107 | 'loss': outputs['loss'], |
| 108 | 'nlls': nlls, |
| 109 | 'labels': next_token_labels, |
| 110 | 'blocks': blocks |
| 111 | } |
| 112 | |
| 113 | |
| 114 | def evaluate(self, dataset: Dataset) -> Dict[str, Any]: |