| """Held-out perplexity over a packed token corpus.""" |
|
|
| from __future__ import annotations |
|
|
| import math |
| from dataclasses import dataclass |
| from pathlib import Path |
|
|
| import torch |
|
|
| from strata.training.lm_data import PackedLMDataset |
|
|
|
|
| @dataclass(slots=True) |
| class PerplexityResult: |
| corpus: str |
| perplexity: float |
| loss: float |
| tokens: int |
| windows: int |
|
|
| def to_dict(self) -> dict[str, object]: |
| return { |
| "corpus": self.corpus, |
| "perplexity": self.perplexity, |
| "loss": self.loss, |
| "tokens": self.tokens, |
| "windows": self.windows, |
| } |
|
|
|
|
| @torch.no_grad() |
| def evaluate_perplexity( |
| model, |
| corpus_dir: Path, |
| *, |
| seq_len: int, |
| device: torch.device, |
| batch_size: int = 8, |
| max_windows: int | None = None, |
| precision: str = "bf16", |
| predicate_memory_intervention: str = "none", |
| predicate_memory_residual_scale: float | None = None, |
| graph_object_residual_scale: float | None = None, |
| ) -> PerplexityResult: |
| """Token-weighted mean NLL / perplexity over sequential corpus windows. |
| |
| Evaluation is deterministic (no shuffle). Windows are scored in order; pass |
| ``max_windows`` to cap the number evaluated for a quick estimate. |
| """ |
|
|
| dataset = PackedLMDataset(corpus_dir, seq_len=seq_len) |
| n_windows = len(dataset) if max_windows is None else min(len(dataset), max_windows) |
| if n_windows == 0: |
| raise ValueError(f"corpus {corpus_dir} has no windows at seq_len {seq_len}") |
|
|
| device_type = device.type |
| autocast = ( |
| torch.autocast(device_type=device_type, enabled=False) |
| if (precision == "fp32" or device_type == "cpu") |
| else torch.autocast(device_type=device_type, dtype=torch.bfloat16 if precision == "bf16" else torch.float16) |
| ) |
|
|
| model.eval() |
| loss_sum = 0.0 |
| token_count = 0 |
| for start in range(0, n_windows, batch_size): |
| indices = range(start, min(start + batch_size, n_windows)) |
| batch = torch.stack([dataset[i] for i in indices]).to(device) |
| with autocast: |
| output = model( |
| batch, |
| labels=batch, |
| predicate_memory_intervention=predicate_memory_intervention, |
| predicate_memory_residual_scale=predicate_memory_residual_scale, |
| graph_object_residual_scale=graph_object_residual_scale, |
| ) |
| |
| predicted = batch.shape[0] * (batch.shape[1] - 1) |
| loss_sum += float(output.loss) * predicted |
| token_count += predicted |
|
|
| mean_loss = loss_sum / token_count |
| return PerplexityResult( |
| corpus=corpus_dir.name, |
| perplexity=math.exp(mean_loss), |
| loss=mean_loss, |
| tokens=token_count, |
| windows=n_windows, |
| ) |
|
|
