Spaces:
Sleeping
Sleeping
File size: 3,743 Bytes
4e316d6 | 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 92 93 94 95 96 97 98 99 100 101 102 103 104 | """
Perplexity evaluation on WikiText-2.
Perplexity measures how well a language model predicts a held-out
corpus — lower is better. It is defined as exp(average NLL) over
all tokens. For sequences longer than the model's context window
a sliding-window approach is used: the window advances by a stride
and only the newly exposed tokens contribute to the loss, avoiding
double-counting while still conditioning on full context.
This module downloads WikiText-2 from HuggingFace Datasets and
evaluates any model + tokenizer pair with configurable stride.
"""
import math
import torch
import torch.nn.functional as F
from datasets import load_dataset
def evaluate_perplexity(model, tokenizer, max_length=1024, stride=512, device=None, dataset_text=None):
"""
Compute perplexity on WikiText-2 test set using a sliding window.
The window is `max_length` tokens wide and advances by `stride` tokens.
Only the last `stride` tokens in each window are scored (except the
first window, which scores all tokens). Earlier tokens in the window
provide conditioning context but don't contribute to the loss, so
every token is scored exactly once with maximal left context.
Args:
model: Language model with forward(input_ids) -> logits
tokenizer: Tokenizer with encode(text) -> list[int]
max_length: Context window size for each forward pass
stride: How far the window advances each step (stride <= max_length)
device: Torch device
dataset_text: Override text instead of WikiText-2 (for testing)
Returns:
Perplexity as a float
"""
model.eval()
if device is None:
device = next(model.parameters()).device
# Load and tokenize
if dataset_text is None:
dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")
text = "\n\n".join(dataset["text"])
else:
text = dataset_text
encodings = tokenizer.encode(text)
token_ids = torch.tensor(encodings, dtype=torch.long)
seq_len = token_ids.size(0)
if seq_len < 2:
raise ValueError(f"Need at least 2 tokens to compute perplexity, got {seq_len}")
total_loss = 0.0
total_tokens = 0
prev_end = 0
with torch.no_grad():
for begin in range(0, seq_len - 1, stride):
end = min(begin + max_length, seq_len)
input_ids = token_ids[begin:end].unsqueeze(0).to(device)
logits = model(input_ids) # (1, end-begin, vocab_size)
# logits[:, i, :] predicts token at absolute position begin+i+1.
# Targets: token_ids[begin+1 .. min(end+1, seq_len)-1]
target_end = min(end + 1, seq_len)
num_preds = target_end - begin - 1
targets = token_ids[begin + 1:target_end].unsqueeze(0).to(device)
logits = logits[:, :num_preds, :]
# Only score newly exposed tokens to avoid double-counting.
# Positions [begin+1, prev_end] were already scored by the
# previous window; skip them (they served as context here).
score_from = max(prev_end - begin, 0)
scored_logits = logits[:, score_from:, :]
scored_targets = targets[:, score_from:]
if scored_targets.numel() == 0:
continue
loss = F.cross_entropy(
scored_logits.reshape(-1, scored_logits.size(-1)),
scored_targets.reshape(-1),
reduction="sum",
)
total_loss += loss.item()
total_tokens += scored_targets.numel()
prev_end = end
if end == seq_len:
break
avg_loss = total_loss / total_tokens
return math.exp(avg_loss)
|