| """Bits-per-byte evaluation. |
| |
| BPB = (sum of next-token NLL in nats) / (ln2 * total UTF-8 bytes of the text). |
| |
| The byte denominator makes BPB comparable across tokenizers with different |
| fertilities -- the whole reason the plan reports BPB rather than per-token loss. |
| Each document is scored with a leading <bos>; the <bos> itself is never a target |
| and <eos> is included as a target (it is real signal the model must predict). |
| Long documents are scored in sliding windows of the model's context length. |
| """ |
| import json |
| from pathlib import Path |
|
|
| import torch |
|
|
| from ..tok.wrapper import BOS_ID, EOS_ID |
|
|
|
|
| @torch.no_grad() |
| def score_texts(model, tok, texts, device, seq_len: int, batch_tokens: int = 8192): |
| """Return (nll_nats, n_bytes, n_target_tokens) summed over `texts`.""" |
| model.eval() |
| total_nll = 0.0 |
| total_bytes = 0 |
| total_tokens = 0 |
| for text in texts: |
| b = len(text.encode("utf-8")) |
| if b == 0: |
| continue |
| ids = tok.encode(text, bos=True, eos=True) |
| total_bytes += b |
| |
| for st in range(0, len(ids) - 1, seq_len): |
| chunk = ids[st:st + seq_len + 1] |
| if len(chunk) < 2: |
| continue |
| x = torch.tensor(chunk[:-1], device=device).unsqueeze(0) |
| y = torch.tensor(chunk[1:], device=device).unsqueeze(0) |
| logits, _ = model(x, y) |
| nll = torch.nn.functional.cross_entropy( |
| logits.view(-1, logits.size(-1)), y.view(-1), reduction="sum") |
| total_nll += float(nll) |
| total_tokens += y.numel() |
| return total_nll, total_bytes, total_tokens |
|
|
|
|
| def bpb(nll_nats: float, n_bytes: int) -> float: |
| import math |
| return nll_nats / (math.log(2) * max(n_bytes, 1)) |
|
|
|
|
| def eval_sources(model, tok, sources: dict[str, list[str]], device, seq_len: int) -> dict: |
| """sources: name -> list[str]. Returns {name: {bpb, ppl_tok, bytes, tokens}}.""" |
| import math |
| out = {} |
| for name, texts in sources.items(): |
| nll, nbytes, ntok = score_texts(model, tok, texts, device, seq_len) |
| out[name] = { |
| "bpb": bpb(nll, nbytes), |
| "ppl_token": math.exp(nll / max(ntok, 1)), |
| "bytes": nbytes, |
| "tokens": ntok, |
| } |
| return out |
|
|
|
|
| def run(run_name: str, tok_name: str, tag: str = "final", |
| langs=None, out_dir=None) -> dict: |
| """Re-evaluate a saved checkpoint's BPB on FLORES+ dev and holdout.""" |
| from pathlib import Path |
| from ..model import ModelConfig, Transformer |
| from ..tok.wrapper import Tok |
| from ..paths import RUNS, RESULTS, tokenizer_dir, ensure |
| from .. import flores |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| ck = torch.load(RUNS / run_name / "checkpoints" / f"{tag}.pt", |
| map_location="cpu", weights_only=False) |
| model = Transformer(ModelConfig(**ck["cfg"]["model"])) |
| model.load_state_dict(ck["model"]) |
| model = model.to(device).eval() |
| tok = Tok(tokenizer_dir(tok_name)) |
| langs = langs or list(ck["cfg"]["langs"]) |
| srcs = {} |
| for l in langs: |
| h = load_holdout(l) |
| if h: |
| srcs[f"holdout_{l}"] = h |
| for l, sents in flores.load_parallel(langs, "dev").items(): |
| srcs[f"flores_{l}"] = sents |
| res = eval_sources(model, tok, srcs, device, model.cfg.max_seq_len) |
| out_dir = ensure(Path(out_dir) if out_dir else RESULTS / "bpb") |
| (out_dir / f"{run_name}_{tag}.json").write_text(json.dumps(res, indent=2)) |
| print(f"[bpb] {run_name} ({tag}): " + |
| ", ".join(f"{k}={v['bpb']:.4f}" for k, v in res.items())) |
| return res |
|
|
|
|
| def load_holdout(lang: str, max_docs: int = 2000) -> list[str]: |
| """In-domain eval text from the reserved FineWeb holdout shard.""" |
| from ..paths import HOLDOUT |
| import zstandard, io |
| texts = [] |
| for p in sorted(HOLDOUT.glob(f"{lang}_*.jsonl.zst")): |
| with open(p, "rb") as raw: |
| r = zstandard.ZstdDecompressor().stream_reader(raw) |
| for line in io.TextIOWrapper(r, encoding="utf-8"): |
| if line.strip(): |
| texts.append(json.loads(line)["text"]) |
| if len(texts) >= max_docs: |
| return texts |
| return texts |
|
|
