source/eval/perplexity.py

"""
Compute sliding-window perplexity of a trained LLM on a binary token dataset.

The sliding-window approach avoids the boundary effect of chunking: a window
of ``seq_len`` tokens is evaluated every ``stride`` tokens.  Positions in
the first (stride) tokens of each window are considered "fresh" context and
their NLL contributions are accumulated; positions in the overlap region are
not double-counted because only the *new* stride tokens are aggregated at
each step.

Reference: Press et al., 2022 "Train Short, Test Long" (sliding-window PPL).

Usage:
    python eval/perplexity.py \
        --checkpoint checkpoints/checkpoint-0100000 \
        --data data/val.bin \
        --seq_len 2048 \
        --batch_size 4 \
        --device cuda:0 \
        --stride 512
"""

from __future__ import annotations

import argparse
import math
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm
from model.transformer import LLM


# ---------------------------------------------------------------------------
# Sliding-window dataset
# ---------------------------------------------------------------------------

class SlidingWindowDataset(Dataset):
    """
    Yields (input_ids, targets, loss_mask) tuples for sliding-window PPL.

    ``loss_mask`` is 1 for positions that contribute to the perplexity
    estimate (i.e. the *new* stride tokens at the right end of the window)
    and 0 for the context-only positions.

    Args:
        tokens:   Flat 1-D numpy array of token IDs (uint16).
        seq_len:  Context window size.
        stride:   Step size between consecutive windows.
    """

    def __init__(self, tokens: np.ndarray, seq_len: int, stride: int) -> None:
        self.tokens  = tokens
        self.seq_len = seq_len
        self.stride  = stride
        # Number of windows that fit inside the token array.
        self.n_windows = max(0, (len(tokens) - seq_len + stride - 1) // stride)

    def __len__(self) -> int:
        return self.n_windows

    def __getitem__(self, idx: int):
        start = idx * self.stride
        end   = start + self.seq_len

        # Clamp end to array length; pad if needed.
        actual_end = min(end, len(self.tokens))
        chunk_len  = actual_end - start  # may be < seq_len for last window

        input_ids = torch.zeros(self.seq_len, dtype=torch.long)
        targets   = torch.full((self.seq_len,), fill_value=-100, dtype=torch.long)
        loss_mask = torch.zeros(self.seq_len, dtype=torch.bool)

        if chunk_len > 1:
            toks = torch.from_numpy(
                self.tokens[start : actual_end].astype(np.int64)
            )
            input_ids[: chunk_len]     = toks
            targets  [: chunk_len - 1] = toks[1:]  # shifted labels

        # The "new" tokens start at stride positions from the beginning of the
        # window (they haven't been seen as targets in any previous window).
        # For the very first window (idx == 0) all positions are new.
        new_start_in_window = 0 if idx == 0 else self.stride
        # Loss mask covers [new_start_in_window, chunk_len - 1) because we
        # predict token[t+1] from token[t], so the last input position has no
        # target within this window.
        if chunk_len > 1:
            mask_end = chunk_len - 1  # positions 0 … chunk_len-2 have valid targets
            for pos in range(new_start_in_window, mask_end):
                loss_mask[pos] = True

        return input_ids, targets, loss_mask


# ---------------------------------------------------------------------------
# PPL computation
# ---------------------------------------------------------------------------

@torch.inference_mode()
def compute_perplexity(
    model: torch.nn.Module,
    data_path: str,
    seq_len: int,
    batch_size: int,
    device: str,
    stride: int,
) -> float:
    """
    Compute sliding-window perplexity on the token file at ``data_path``.

    Returns:
        Perplexity (float).
    """
    path = Path(data_path)
    if not path.exists():
        raise FileNotFoundError(f"Data file not found: {path}")

    tokens = np.memmap(path, dtype="uint16", mode="r")
    total_tokens = len(tokens)
    print(f"Loaded {total_tokens:,} tokens from {path}")

    dataset = SlidingWindowDataset(tokens, seq_len=seq_len, stride=stride)
    if len(dataset) == 0:
        raise ValueError(
            f"No windows fit in {total_tokens} tokens with seq_len={seq_len}."
        )

    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=0,
        pin_memory=True,
    )

    model.eval()

    # Accumulate log-probabilities (sum of NLL) and the count of evaluated tokens.
    total_nll   = 0.0
    total_count = 0

    for batch_input_ids, batch_targets, batch_loss_mask in tqdm(
        loader, desc="Evaluating perplexity", unit="batch"
    ):
        batch_input_ids = batch_input_ids.to(device)     # [B, seq_len]
        batch_targets   = batch_targets.to(device)       # [B, seq_len]
        batch_loss_mask = batch_loss_mask.to(device)     # [B, seq_len]

        logits, _ = model(batch_input_ids)              # [B, seq_len, vocab]

        # Cross-entropy loss per position (reduction='none').
        B, S, V = logits.shape
        ce = F.cross_entropy(
            logits.reshape(B * S, V),
            batch_targets.reshape(B * S),
            ignore_index=-100,
            reduction="none",
        ).reshape(B, S)  # [B, seq_len]

        # Apply sliding-window loss mask.
        # Positions where targets == -100 are already zeroed by ignore_index;
        # we additionally zero positions outside the stride window.
        masked_ce = ce * batch_loss_mask.float()
        total_nll   += masked_ce.sum().item()
        total_count += batch_loss_mask.sum().item()

    if total_count == 0:
        raise RuntimeError("No valid token positions were evaluated.")

    avg_nll    = total_nll / total_count
    perplexity = math.exp(avg_nll)
    return perplexity


# ---------------------------------------------------------------------------
# Argument parsing
# ---------------------------------------------------------------------------

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Compute sliding-window perplexity of a trained LLM."
    )
    parser.add_argument(
        "--checkpoint",
        required=True,
        help="Path to the checkpoint directory.",
    )
    parser.add_argument(
        "--data",
        required=True,
        help="Path to the .bin token data file.",
    )
    parser.add_argument(
        "--seq_len",
        type=int,
        default=2048,
        help="Context window length (default: 2048).",
    )
    parser.add_argument(
        "--batch_size",
        type=int,
        default=4,
        help="Evaluation batch size (default: 4).",
    )
    parser.add_argument(
        "--device",
        default="cuda:0",
        help="Torch device string (default: cuda:0).",
    )
    parser.add_argument(
        "--stride",
        type=int,
        default=512,
        help=(
            "Stride for sliding window PPL; smaller = more accurate, "
            "slower (default: 512)."
        ),
    )
    return parser.parse_args()


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

def main() -> None:
    args = parse_args()

    ckpt_path = Path(args.checkpoint)
    if not ckpt_path.exists():
        raise FileNotFoundError(f"Checkpoint directory not found: {ckpt_path}")

    print(f"Loading model from: {ckpt_path}")
    model = LLM.from_pretrained(str(ckpt_path)).to(device=args.device, dtype=torch.float16)
    model.eval()
    print(f"Model parameters: {model.num_params / 1e6:.1f}M")

    print(
        f"\nPerplexity config: seq_len={args.seq_len}, "
        f"stride={args.stride}, batch_size={args.batch_size}"
    )

    ppl = compute_perplexity(
        model=model,
        data_path=args.data,
        seq_len=args.seq_len,
        batch_size=args.batch_size,
        device=args.device,
        stride=args.stride,
    )

    print("\n" + "=" * 50)
    print(f"  Perplexity: {ppl:.4f}")
    print(f"  Bits/token: {math.log2(math.e) * math.log(ppl):.4f}")
    print("=" * 50)


if __name__ == "__main__":
    main()