scripts/eval_compare.py

# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "torch",
#   "numpy",
#   "datasets>=2.19",
#   "huggingface_hub>=0.24",
#   "safetensors>=0.4",
#   "transformers>=4.44",
# ]
# ///
"""Bits-per-byte comparison: Trident vs a real open-source model.

Both models are scored on the *same* held-out UTF-8 byte string, so the
comparison is apples-to-apples:

    BPB = total_negative_log_likelihood_bits / total_UTF8_bytes

For the token baseline the per-token NLL is summed and divided by the number of
UTF-8 bytes of the same text (standard cross-tokenizer BPB, as in byte-LM
papers). Results are written to results/scorecard.json in the code repo.
"""
import json
import math
import os
import sys
import time
from pathlib import Path

import torch


def env(k, d=None):
    v = os.environ.get(k)
    return v if v not in (None, "") else d


def log(m):
    print(f"[eval] {m}", flush=True)


def build_heldout_text(dataset, name, split, text_field, target_bytes):
    from datasets import load_dataset
    ds = load_dataset(dataset, name=name, split=split, streaming=True)
    parts, total = [], 0
    for row in ds:
        t = row.get(text_field)
        if not t:
            continue
        parts.append(t)
        total += len(t.encode("utf-8"))
        if total >= target_bytes:
            break
    text = "\n\n".join(parts)
    return text


def trident_bpb(text, code_repo, ckpt_tag, seq_len, device):
    from huggingface_hub import snapshot_download
    from safetensors.torch import load_file

    work = Path("/tmp/trident_code")
    snapshot_download(repo_id=code_repo, repo_type="model",
                      allow_patterns=["src/**"], local_dir=str(work))
    sys.path.insert(0, str(work / "src"))
    from trident import Trident, TridentConfig

    ckpt = Path("/tmp/ckpt")
    ckpt_dir = env("CKPT_DIR", "checkpoints")
    snapshot_download(repo_id=code_repo, repo_type="model",
                      allow_patterns=[f"{ckpt_dir}/{ckpt_tag}/*"], local_dir=str(ckpt))
    cdir = ckpt / ckpt_dir / ckpt_tag
    cfg = TridentConfig(**json.loads((cdir / "config.json").read_text()))
    model = Trident(cfg).to(device).eval()
    sd = load_file(str(cdir / "model.safetensors"))
    missing, unexpected = model.load_state_dict(sd, strict=False)
    if missing:
        log(f"WARNING missing keys: {len(missing)}")
    data = torch.tensor(list(text.encode("utf-8")), dtype=torch.long)
    n = (data.numel() // seq_len) * seq_len
    data = data[:n].view(-1, seq_len).to(device)
    tot_nll, tot_bytes = 0.0, 0
    amp = torch.bfloat16 if device == "cuda" else torch.float32
    with torch.no_grad():
        for i in range(0, data.shape[0], 8):
            chunk = data[i:i + 8]
            with torch.autocast(device_type="cuda", dtype=amp, enabled=device == "cuda"):
                out = model(chunk, return_logits=True)
            logits = out["logits"].float()
            nll = torch.nn.functional.cross_entropy(
                logits.reshape(-1, logits.shape[-1]), chunk.reshape(-1), reduction="sum")
            tot_nll += nll.item()
            tot_bytes += chunk.numel()
    return tot_nll / tot_bytes / math.log(2), sum(p.numel() for p in model.parameters())


def baseline_bpb(text, model_id, device, ctx=1024):
    from transformers import AutoModelForCausalLM, AutoTokenizer
    tok = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(
        model_id, torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32).to(device).eval()
    ids = tok(text, return_tensors="pt").input_ids[0]
    n_bytes = len(text.encode("utf-8"))
    tot_nll = 0.0
    with torch.no_grad():
        for i in range(0, ids.numel() - 1, ctx):
            window = ids[i:i + ctx + 1].to(device)
            if window.numel() < 2:
                break
            inp = window[:-1].unsqueeze(0)
            tgt = window[1:].unsqueeze(0)
            logits = model(inp).logits.float()
            nll = torch.nn.functional.cross_entropy(
                logits.reshape(-1, logits.shape[-1]), tgt.reshape(-1), reduction="sum")
            tot_nll += nll.item()
    return tot_nll / n_bytes / math.log(2), sum(p.numel() for p in model.parameters())


def main():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    code_repo = env("CODE_REPO", "farguney/trident")
    ckpt_tag = env("CKPT_TAG", "final")
    seq_len = int(env("SEQ_LEN", 2048))
    target_bytes = int(env("TARGET_BYTES", 1_000_000))
    baselines = env("BASELINES", "HuggingFaceTB/SmolLM2-135M,EleutherAI/pythia-160m").split(",")

    log(f"device={device} building held-out ({target_bytes} bytes)")
    text = build_heldout_text(
        env("VAL_DATASET", "Salesforce/wikitext"), env("VAL_NAME", "wikitext-103-raw-v1"),
        env("VAL_SPLIT", "validation"), env("VAL_TEXT_FIELD", "text"), target_bytes)
    real_bytes = len(text.encode("utf-8"))
    log(f"held-out bytes={real_bytes}")

    results = {"held_out": {"dataset": env("VAL_DATASET", "Salesforce/wikitext"),
                            "bytes": real_bytes}, "models": {}}

    t = time.time()
    tri_bpb, tri_params = trident_bpb(text, code_repo, ckpt_tag, seq_len, device)
    results["models"]["trident"] = {"bpb": tri_bpb, "params_M": tri_params / 1e6,
                                    "class": "byte / fixed-state recurrent", "ckpt": ckpt_tag}
    log(f"Trident BPB={tri_bpb:.4f} params={tri_params/1e6:.1f}M ({time.time()-t:.0f}s)")

    for b in baselines:
        b = b.strip()
        if not b:
            continue
        try:
            t = time.time()
            bpb, params = baseline_bpb(text, b, device)
            results["models"][b] = {"bpb": bpb, "params_M": params / 1e6,
                                    "class": "subword transformer (pretrained)"}
            log(f"{b} BPB={bpb:.4f} params={params/1e6:.1f}M ({time.time()-t:.0f}s)")
        except Exception as e:  # noqa
            log(f"baseline {b} failed: {e}")

    # scorecard
    print("\n==== BPB SCORECARD (lower is better) ====", flush=True)
    for name, m in sorted(results["models"].items(), key=lambda kv: kv[1]["bpb"]):
        print(f"  {name:40s} {m['bpb']:.4f}  ({m['params_M']:.1f}M)", flush=True)

    from huggingface_hub import HfApi
    out = Path("/tmp/results"); out.mkdir(exist_ok=True)
    (out / "scorecard.json").write_text(json.dumps(results, indent=2))
    try:
        HfApi().upload_folder(folder_path=str(out), repo_id=code_repo, repo_type="model",
                              path_in_repo="results", commit_message="BPB scorecard")
        log("scorecard pushed to results/")
    except Exception as e:  # noqa
        log(f"scorecard upload failed: {e}")


if __name__ == "__main__":
    main()