File size: 4,402 Bytes
1611cfd | 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 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | #!/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()
|