| """Validation perplexity evaluation.""" |
|
|
| from __future__ import annotations |
|
|
| import math |
|
|
| import torch |
|
|
| from train.loss import masked_cross_entropy |
|
|
|
|
| @torch.no_grad() |
| def evaluate_perplexity( |
| model: torch.nn.Module, |
| dataloader, |
| device: torch.device, |
| dtype: torch.dtype | None = None, |
| max_batches: int = 16, |
| ) -> dict[str, float]: |
| """Evaluate average loss and perplexity on a validation loader.""" |
| model.eval() |
| losses: list[float] = [] |
| for index, batch in enumerate(dataloader): |
| if index >= max_batches: |
| break |
| input_ids = batch["input_ids"].to(device) |
| labels = batch["labels"].to(device) |
| loss_mask = batch["loss_mask"].to(device) |
| if dtype is not None and device.type != "cpu": |
| with torch.autocast(device_type=device.type, dtype=dtype): |
| logits, _ = model(input_ids) |
| loss = masked_cross_entropy(logits, labels, loss_mask) |
| else: |
| logits, _ = model(input_ids) |
| loss = masked_cross_entropy(logits, labels, loss_mask) |
| losses.append(float(loss)) |
| model.train() |
| mean_loss = sum(losses) / max(len(losses), 1) |
| return {"loss": mean_loss, "perplexity": math.exp(min(mean_loss, 20.0))} |
|
|
