Mean per-sequence next-token NLL, exp'd; logits chunked along seq axis for memory.
(model, dataloader, loss_chunk: int = 2048)
| 636 | |
| 637 | @torch.no_grad() |
| 638 | def evaluate_ppl(model, dataloader, loss_chunk: int = 2048) -> float: |
| 639 | """Mean per-sequence next-token NLL, exp'd; logits chunked along seq axis for memory.""" |
| 640 | model.eval() |
| 641 | device = next(model.parameters()).device |
| 642 | loss_fn = nn.CrossEntropyLoss(reduction="none") |
| 643 | nll_means = [] |
| 644 | for batch in tqdm(dataloader, desc="PPL", unit="batch"): |
| 645 | ids = batch["input_ids"].to(device) |
| 646 | mask = batch["attention_mask"].to(device) |
| 647 | logits = model(input_ids=ids, attention_mask=mask, use_cache=False).logits |
| 648 | shift_logits = logits[:, :-1, :] |
| 649 | shift_targets = ids[:, 1:] |
| 650 | S = shift_logits.shape[1] |
| 651 | parts = [] |
| 652 | for s in range(0, S, loss_chunk): |
| 653 | e = min(s + loss_chunk, S) |
| 654 | parts.append(loss_fn( |
| 655 | shift_logits[:, s:e, :].float().permute(0, 2, 1), |
| 656 | shift_targets[:, s:e], |
| 657 | )) |
| 658 | nll = torch.cat(parts, dim=1) |
| 659 | del logits, shift_logits |
| 660 | m = mask[:, 1:].float() |
| 661 | nll_means.append(((nll * m).sum(dim=1) / m.sum(dim=1).clamp_min(1)).cpu()) |
| 662 | return torch.exp(torch.cat(nll_means).mean()).item() |
| 663 | |
| 664 | |
| 665 | __all__ = [ |