| |
| """Sliding-window WT2 + C4 perplexity for one causal-LM checkpoint. |
| |
| Uses the checkpoint's own tokenizer, so it works for a plain HF model id or a |
| local fake-quant checkpoint. Prints the two numbers and (optionally) writes them |
| to a JSON file. |
| |
| Usage: |
| ppl_eval.py MODEL [--out ppl.json] [--seq 2048] [--stride 512] [--device cuda:0] |
| """ |
| import argparse |
| import json |
|
|
| import numpy as np |
| import torch |
| from datasets import load_dataset |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
|
|
| def wt2_text(): |
| return "\n\n".join(load_dataset("wikitext", "wikitext-2-raw-v1", split="test")["text"]) |
|
|
|
|
| def c4_text(min_chars): |
| ds = load_dataset("allenai/c4", "en", split="validation", streaming=True) |
| parts, tot = [], 0 |
| for ex in ds: |
| parts.append(ex["text"]); tot += len(ex["text"]) |
| if tot >= min_chars: |
| break |
| return "\n\n".join(parts) |
|
|
|
|
| def perplexity(model, tok, text, seq, stride, device): |
| ids = tok(text, return_tensors="pt").input_ids[0] |
| n = len(ids); nlls = []; prev_end = 0 |
| for begin in range(0, n - 1, stride): |
| end = min(begin + seq, n) |
| trg_len = end - prev_end |
| chunk = ids[begin:end].unsqueeze(0).to(device) |
| with torch.no_grad(): |
| logits = model(chunk, labels=chunk).logits |
| sl = logits[:, prev_end - begin:-1, :].contiguous() |
| lbl = chunk[:, prev_end - begin + 1:].contiguous() |
| loss = torch.nn.functional.cross_entropy(sl.view(-1, sl.size(-1)), lbl.view(-1)) |
| nlls.append(loss.item() * trg_len) |
| prev_end = end |
| if end == n: |
| break |
| return float(np.exp(sum(nlls) / prev_end)) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("model") |
| ap.add_argument("--out", default=None, help="write {wikitext2, c4} JSON here") |
| ap.add_argument("--seq", type=int, default=2048) |
| ap.add_argument("--stride", type=int, default=512) |
| ap.add_argument("--device", default="cuda:0") |
| ap.add_argument("--c4-chars", type=int, default=2_621_440) |
| args = ap.parse_args() |
|
|
| tok = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True) |
| model = AutoModelForCausalLM.from_pretrained( |
| args.model, torch_dtype=torch.bfloat16, device_map=args.device, |
| trust_remote_code=True).eval() |
|
|
| res = {} |
| res["wikitext2"] = perplexity(model, tok, wt2_text(), args.seq, args.stride, args.device) |
| print(f"WT2 {res['wikitext2']:.4f}", flush=True) |
| res["c4"] = perplexity(model, tok, c4_text(args.c4_chars), args.seq, args.stride, args.device) |
| print(f"C4 {res['c4']:.4f}", flush=True) |
| if args.out: |
| with open(args.out, "w") as f: |
| json.dump(res, f, indent=2) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
