File size: 2,789 Bytes
40e76c0 | 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 | #!/usr/bin/env python3
"""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())
|