Spaces:
Runtime error
Runtime error
| """ | |
| Phase L - standalone val-set perplexity for Abyss-1. | |
| Walks sequentially through val.bin with non-overlapping windows (no | |
| randomness), ao the reported perplexity is a true full-dataset figure. | |
| Usage (from the_final_product/ directory): | |
| python -m ml.eval | |
| python -m ml.eval --checkpoint ml/checkpoints/ckpt.pt --batch_size 32 | |
| """ | |
| import argparse | |
| import math | |
| import os | |
| import numpy as np | |
| import torch | |
| from contextlib import nullcontext | |
| from ml.model import TheAbyssModel | |
| VAL_BIN = os.path.join(os.path.dirname(os.path.abspath(__file__)), | |
| "data", "tinystories", "val.bin") | |
| def compute_perplexity(model, device, block_size, batch_size, ctx): | |
| model.eval() | |
| data = np.memmap(VAL_BIN, dtype=np.uint16, mode="r") | |
| n = len(data) | |
| total_loss = 0.0 | |
| total_tokens = 0 | |
| batches_done = 0 | |
| pos = 0 | |
| while pos + block_size + 1 <= n: | |
| starts = [] | |
| while len(starts) < batch_size and pos + block_size + 1 <= n: | |
| starts.append(pos) | |
| pos += block_size | |
| xs = torch.stack([ | |
| torch.from_numpy(data[i : i + block_size].astype(np.int64)) | |
| for i in starts | |
| ]) | |
| ys = torch.stack([ | |
| torch.from_numpy(data[i + 1 : i + 1 + block_size].astype(np.int64)) | |
| for i in starts | |
| ]) | |
| if device.startswith("cuda"): | |
| xs = xs.pin_memory().to(device, non_blocking=True) | |
| ys = ys.pin_memory().to(device, non_blocking=True) | |
| else: | |
| xs, ys = xs.to(device), ys.to(device) | |
| with ctx: | |
| _, loss = model(xs, ys) | |
| n_tok = xs.numel() | |
| total_loss += loss.item() * n_tok | |
| total_tokens += n_tok | |
| batches_done += 1 | |
| if batches_done % 50 == 0: | |
| running_ppl = math.exp(total_loss / total_tokens) | |
| pct = 100.0 * pos / n | |
| print(f" {pct:5.1f}% | {total_tokens:>10,} tok | running ppl {running_ppl:.2f}", | |
| flush=True) | |
| avg_loss = total_loss / total_tokens | |
| return avg_loss, math.exp(avg_loss), total_tokens | |
| def main(): | |
| parser = argparse.ArgumentParser(description="Val-set perplexity for Abyss-1.") | |
| parser.add_argument("--checkpoint", default="ml/checkpoints/ckpt.pt") | |
| parser.add_argument("--batch_size", type=int, default=16) | |
| parser.add_argument("--block_size", type=int, default=512) | |
| parser.add_argument("--device", | |
| default="cuda" if torch.cuda.is_available() else "cpu") | |
| args = parser.parse_args() | |
| # Load Checkpoint | |
| print(f"loading {args.checkpoint} ...", flush=True) | |
| ck = torch.load(args.checkpoint, map_location=args.device, weights_only=False) | |
| config = ck["config"] | |
| model = TheAbyssModel(config).to(args.device) | |
| model.load_state_dict(ck["model"]) | |
| print(f" iter {ck['iter']} | best_val {ck['best_val']:.4f} " | |
| f" | params {model.num_params():,}", flush=True) | |
| ptdtype = (torch.bfloat16 if args.device == "cuda" and torch.cuda.is_bf16_supported() | |
| else torch.float16 if args.device == "cuda" else torch.float32) | |
| ctx = (nullcontext() if args.device == "cpu" | |
| else torch.autocast(device_type="cuda", dtype=ptdtype)) | |
| print(f"evaluating full val set " | |
| f"(block={args.block_size}, batch={args.batch_size}, {args.device}) ...\n", | |
| flush=True) | |
| loss, ppl, n_tok = compute_perplexity( | |
| model, args.device, args.block_size, args.batch_size, ctx | |
| ) | |
| bar = "=" * 42 | |
| print(f"\n{bar}") | |
| print(f" checkpoint : {args.checkpoint}") | |
| print(f" iter : {ck['iter']}") | |
| print(f" val loss : {loss:.4f}") | |
| print(f" perplexity : {ppl:.2f}") | |
| print(f" tokens : {n_tok:,}") | |
| print(bar) | |
| if __name__ == "__main__": | |
| main() | |