| 126 | |
| 127 | @torch.no_grad() |
| 128 | def perplexity(self, data_dict: dict) -> dict: |
| 129 | |
| 130 | data_dict = self.tokenizer.tokenize(data_dict) |
| 131 | |
| 132 | input_ids = data_dict['input_ids'] # batch x ntoken |
| 133 | attention_mask = data_dict['attention_mask'] # batch x ntoken |
| 134 | |
| 135 | # set pad_token_id = eos_token_id |
| 136 | input_ids[input_ids == self.tokenizer.pad_id] = self.pad_token_id # <pad> xxx <pad> <pad> |
| 137 | input_ids[:, 0] = self.bos_token_id # <sos> xxx <pad> <pad> |
| 138 | eos_pos_id = attention_mask.sum(1, keepdim=True) - 1 |
| 139 | input_ids = torch.scatter( # <bos> xxx <eos> <pad> |
| 140 | input_ids, |
| 141 | 1, |
| 142 | eos_pos_id.long(), |
| 143 | torch.ones_like(input_ids) * self.eos_token_id |
| 144 | ) |
| 145 | |
| 146 | # llm loss calculation |
| 147 | output = self.transformer( |
| 148 | input_ids = input_ids.long(), |
| 149 | ) |
| 150 | |
| 151 | # compute loss with shift token right |
| 152 | logit = output.logits[:, :-1] # batch x (ntoken - 1) x vocab |
| 153 | label = input_ids[:, 1:] # batch x (ntoken - 1) |
| 154 | masks = attention_mask[:, 1:] # batch x (ntoken - 1) |
| 155 | loss_per_token = nnf.cross_entropy( |
| 156 | logit.permute(0, 2, 1), # batch x (ntoken - 1) x ntoken |
| 157 | label, # batch x (ntoken - 1) |
| 158 | reduction='none' |
| 159 | ) # batch x ntoken |
| 160 | |
| 161 | # compute negative log likelihood for each sequence |
| 162 | neg_log_likelihood = torch.sum(loss_per_token * masks, dim=1) / torch.sum(masks, dim=1) |
| 163 | |
| 164 | data_dict['neg_log_likelihood'] = neg_log_likelihood # batch, |
| 165 | return data_dict |
| 166 | |
| 167 | |
| 168 | |