# /// 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()