#!/usr/bin/env python3
"""Bits-per-char / perplexity / top-k против базовой модели. CPU, без GPU.

    python bench/eval_bpc.py --model Rahspide/MNN-Pocket-0.3B-RU \
        --base taobao-mnn/MNN-Pocket-0.3B-Base --dataset fineweb-2 --config rus_Cyrl

Метрика bits/char сравнима между разными токенайзерами, поэтому она здесь главная.
Код выхода 1, если цель проиграла базе больше, чем --max-bpc-regression.
"""
import argparse
import json
import math
import sys

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


def read_text(args, tok):
    if args.text:
        return open(args.text, encoding="utf-8").read()
    from datasets import load_dataset
    ds = load_dataset(args.dataset, args.config, split="train", streaming=True)
    parts, seen = [], 0
    for row in ds:
        chunk = row.get("text") or ""
        if not chunk:
            continue
        parts.append(chunk)
        seen += len(tok(chunk, add_special_tokens=False)["input_ids"])
        if seen >= args.tokens * 2:
            break
    return "\n".join(parts)


def evaluate(model_id, text, tok, args):
    dtype = torch.bfloat16 if args.bf16 else torch.float32
    model = AutoModelForCausalLM.from_pretrained(
        model_id, dtype=dtype, attn_implementation="sdpa", token=args.hf_token)
    model.eval()
    ids = torch.tensor(tok(text, add_special_tokens=False)["input_ids"][: args.tokens + 1])
    n = ids.numel() - 1
    if n < 8:
        sys.exit("not enough tokens in the text sample")
    nll, c1, c5 = 0.0, 0, 0
    with torch.inference_mode():
        for i in range(0, n, args.seq):
            j = min(i + args.seq, n)
            logits = model(input_ids=ids[i:j][None, :]).logits[0].float()
            y = ids[i + 1:j + 1]
            nll += torch.nn.functional.cross_entropy(logits, y).item() * (j - i)
            top5 = logits.topk(5, dim=-1).indices
            c1 += (top5[:, 0] == y).sum().item()
            c5 += (top5 == y[:, None]).any(-1).sum().item()
    chars = len(tok.decode(ids.tolist()))
    del model
    return {"model": model_id, "bpc": nll / math.log(2) / chars, "ppl": math.exp(nll / n),
            "top1": c1 / n, "top5": c5 / n, "tokens": int(n), "chars": chars}


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--model", default="Rahspide/MNN-Pocket-0.3B-RU")
    p.add_argument("--base", default="taobao-mnn/MNN-Pocket-0.3B-Base")
    p.add_argument("--text", help="файл с текстом; если не задан, берётся из --dataset")
    p.add_argument("--dataset", default="HuggingFaceFW/fineweb-2")
    p.add_argument("--config", default="rus_Cyrl")
    p.add_argument("--tokens", type=int, default=20000)
    p.add_argument("--seq", type=int, default=1024)
    p.add_argument("--threads", type=int, default=16)
    p.add_argument("--bf16", action="store_true", help="быстрее, но числа слегка расходятся с fp32")
    p.add_argument("--max-bpc-regression", type=float, default=0.0,
                   help="допустимый рост bits/char против базы; 0 = только улучшение")
    p.add_argument("--json", action="store_true")
    p.add_argument("--hf-token", default=None)
    args = p.parse_args()

    if args.base == "base":
        args.base = "taobao-mnn/MNN-Pocket-0.3B-Base"
    torch.set_num_threads(args.threads)
    tok = AutoTokenizer.from_pretrained(args.base, token=args.hf_token)
    text = read_text(args, tok)

    rows = [evaluate(m, text, tok, args) for m in (args.base, args.model)]
    delta = (rows[1]["bpc"] - rows[0]["bpc"]) / rows[0]["bpc"]
    if args.json:
        print(json.dumps({"rows": rows, "bpc_delta": delta}, indent=1))
    else:
        print(f"{'model':52} {'bpc':>7} {'ppl':>9} {'top1':>7} {'top5':>7}")
        for r in rows:
            print(f"{r['model'][:52]:52} {r['bpc']:7.3f} {r['ppl']:9.2f} "
                  f"{r['top1']*100:6.1f}% {r['top5']*100:6.1f}%")
        print(f"\nbits/char {rows[1]['bpc']:.3f} against base {rows[0]['bpc']:.3f} "
              f"({delta:+.1%}) on {rows[1]['tokens']} tokens")
    if delta > args.max_bpc_regression:
        print("FAIL: regression above the allowed margin", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()