#!/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())