| 152 | |
| 153 | # Function to evaluate perplexity (ppl) specifically on the wikitext dataset |
| 154 | def eval_ppl_wikitext(model, testenc, bs=1, device=None): |
| 155 | # Get input IDs |
| 156 | testenc = testenc.input_ids |
| 157 | |
| 158 | # Calculate number of samples |
| 159 | nsamples = testenc.numel() // model.seqlen |
| 160 | |
| 161 | # List to store negative log likelihoods |
| 162 | nlls = [] |
| 163 | print(f"nsamples {nsamples}") |
| 164 | |
| 165 | # Loop through each batch |
| 166 | for i in range(0, nsamples, bs): |
| 167 | if i % 50 == 0: |
| 168 | print(f"sample {i}") |
| 169 | |
| 170 | # Calculate end index |
| 171 | j = min(i + bs, nsamples) |
| 172 | |
| 173 | # Prepare inputs and move to device |
| 174 | inputs = testenc[:, (i * model.seqlen) : (j * model.seqlen)].to(device) |
| 175 | inputs = inputs.reshape(j - i, model.seqlen) |
| 176 | |
| 177 | # Forward pass through the model |
| 178 | lm_logits = model(inputs).logits |
| 179 | |
| 180 | # Shift logits and labels for next token prediction |
| 181 | shift_logits = lm_logits[:, :-1, :].contiguous() |
| 182 | shift_labels = inputs[:, 1:] |
| 183 | |
| 184 | # Compute loss |
| 185 | loss_fct = nn.CrossEntropyLoss() |
| 186 | loss = loss_fct( |
| 187 | shift_logits.reshape(-1, shift_logits.size(-1)), shift_labels.reshape(-1) |
| 188 | ) |
| 189 | |
| 190 | # Calculate negative log likelihood |
| 191 | neg_log_likelihood = loss.float() * model.seqlen * (j - i) |
| 192 | |
| 193 | # Append to list of negative log likelihoods |
| 194 | nlls.append(neg_log_likelihood) |
| 195 | |
| 196 | # Compute perplexity |
| 197 | ppl = torch.exp(torch.stack(nlls).sum() / (nsamples * model.seqlen)) |
| 198 | |
| 199 | # Empty CUDA cache to save memory |
| 200 | torch.cuda.empty_cache() |
| 201 | |
| 202 | return ppl.item() |
| 203 | |
| 204 | |
| 205 | def eval_zero_shot( |