Upload src/xscript/eval/bpb.py with huggingface_hub

src/xscript/eval/bpb.py

@@ -0,0 +1,111 @@
+"""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
+        # sliding, non-overlapping windows of seq_len+1 (predict positions 1..)
+        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