File size: 2,831 Bytes
e69b72a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | """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,
)
# output.loss is the mean CE over B*(S-1) predicted tokens.
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,
)
|