| |
| """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() |
|
|
