bench/run_bpe_vs_char.py

"""Benchmark: char-level vs global BPE on a real corpus (TinyStories).

Measures the quality lift from switching to a shared BPE tokenizer, on real
prose instead of synthetic data. For each tokenizer mode, on the same corpus
extract and same D:

  - build time + throughput
  - next-token accuracy (memorization of seen text)
  - .exp compression ratio
  - query latency
  - a fluency probe (generated continuation quality)

Usage:
    python bench/run_bpe_vs_char.py --corpus ../../helios/data/tinystories_100k.txt \
        --chars 500000 --D 5000 --vocab 2000 --out results_bpe_vs_char.json
"""
from __future__ import annotations

import argparse
import json
import sys
import time
from dataclasses import asdict, dataclass
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from ensemble import Expert
from ensemble.io import expert_size_bytes
import tempfile, os

# reuse the incremental accuracy measure from run_benchmark
sys.path.insert(0, str(Path(__file__).resolve().parent))
import run_benchmark as bench  # noqa: E402


@dataclass
class ComparisonResult:
    tokenizer: str          # "char" | "bpe"
    D: int
    vocab_size: int
    build_seconds: float
    throughput_tok_per_s: float
    n_tokens: int
    ram_mb: float
    next_token_accuracy: float
    query_latency_ms: float
    exp_size_bytes: int
    source_size_bytes: int
    compression_ratio: float
    fluency_sample: str


def _fluency_probe(expert: Expert, prompt: str, max_new_tokens: int = 40) -> str:
    """Generate a short continuation to eyeball fluency."""
    try:
        return expert.answer(prompt, max_new_tokens=max_new_tokens, temperature=0.0)
    except Exception as e:
        return f"<error: {e}>"


def run_one(tokenizer_name: str, corpus: str, D: int, vocab: int,
            max_eval: int, bpe=None) -> ComparisonResult:
    print(f"\n=== {tokenizer_name.upper()} (D={D}) ===", flush=True)
    t0 = time.perf_counter()
    expert = Expert.from_text(corpus, domain="tinystories", D=D,
                              tokenizer=bpe if tokenizer_name == "bpe" else None)
    build_s = time.perf_counter() - t0
    n_tok = expert.n_traces
    print(f"  built {n_tok} traces in {build_s:.1f}s ({n_tok/build_s:.0f} tok/s)",
          flush=True)

    ram = bench._ram_mb(expert)
    print(f"  measuring next-token accuracy ({max_eval} pts)...", flush=True)
    nta = bench._next_token_accuracy(expert, corpus, max_eval=max_eval)
    print(f"  nta = {nta:.1%}", flush=True)

    # latency
    times = []
    probe = "once upon a time"
    for _ in range(3):
        t1 = time.perf_counter()
        expert.answer(probe, max_new_tokens=20)
        times.append((time.perf_counter() - t1) * 1000)
    lat = float(np.median(times))

    # compression
    with tempfile.TemporaryDirectory() as td:
        res = expert.save(os.path.join(td, "ts.exp"))
        exp_b = res.expert_size_bytes
        src_b = res.source_size_bytes
        ratio = res.compression_ratio

    fluency = _fluency_probe(expert, "once upon a time")

    return ComparisonResult(
        tokenizer=tokenizer_name, D=D,
        vocab_size=(bpe.vocab_size_actual if bpe else expert.vocab_size),
        build_seconds=build_s, throughput_tok_per_s=n_tok / build_s,
        n_tokens=n_tok, ram_mb=ram, next_token_accuracy=nta,
        query_latency_ms=lat, exp_size_bytes=exp_b,
        source_size_bytes=src_b, compression_ratio=ratio,
        fluency_sample=fluency,
    )


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--corpus", required=True, help="path to a text corpus")
    ap.add_argument("--chars", type=int, default=500_000,
                    help="max chars to read from the corpus")
    ap.add_argument("--D", type=int, default=5000)
    ap.add_argument("--vocab", type=int, default=2000, help="BPE target vocab")
    ap.add_argument("--bpe-file", default=None,
                    help="pre-trained BPE vocab json (skip training)")
    ap.add_argument("--max-eval", type=int, default=200)
    ap.add_argument("--out", default="results_bpe_vs_char.json")
    ap.add_argument("--skip-char", action="store_true")
    ap.add_argument("--skip-bpe", action="store_true")
    args = ap.parse_args()

    corpus_path = Path(args.corpus)
    print(f"reading {args.chars:,} chars from {corpus_path.name}...", flush=True)
    corpus = corpus_path.read_text(encoding="utf-8", errors="ignore")[:args.chars]
    print(f"  corpus: {len(corpus):,} chars, {len(corpus.encode('utf-8')):,} bytes",
          flush=True)

    # BPE: load pre-trained or train fresh on the corpus
    bpe = None
    if not args.skip_bpe:
        from palimseste.bpe import BPETokenizer
        from palimseste.lm import PalimpsesteForCausalLM, PalimpsesteConfig
        if args.bpe_file and Path(args.bpe_file).exists():
            print(f"loading pre-trained BPE from {args.bpe_file}...", flush=True)
            cfg = PalimpsesteConfig(D=args.D)
            m = PalimpsesteForCausalLM(config=cfg, rng=np.random.default_rng(0))
            bpe = BPETokenizer.load_vocabulary(args.bpe_file, encoder=m.encoder)
        else:
            print(f"training BPE (vocab={args.vocab}) on corpus...", flush=True)
            t0 = time.perf_counter()
            bpe = Expert.build_bpe(corpus, vocab_size=args.vocab, D=args.D)
            print(f"  BPE trained in {time.perf_counter()-t0:.1f}s, "
                  f"vocab={bpe.vocab_size_actual}", flush=True)

    results = []
    if not args.skip_char:
        results.append(run_one("char", corpus, args.D, args.vocab, args.max_eval))
    if not args.skip_bpe and bpe is not None:
        results.append(run_one("bpe", corpus, args.D, args.vocab, args.max_eval, bpe=bpe))

    with open(args.out, "w", encoding="utf-8") as f:
        json.dump([asdict(r) for r in results], f, indent=2, ensure_ascii=False)

    # summary table
    print("\n" + "=" * 70)
    print(f"{'metric':24} {'char':>16} {'bpe':>16}")
    print("-" * 70)
    by_t = {r.tokenizer: r for r in results}
    for metric in ["build_seconds", "throughput_tok_per_s", "ram_mb",
                   "next_token_accuracy", "query_latency_ms",
                   "compression_ratio", "vocab_size"]:
        c = by_t.get("char")
        b = by_t.get("bpe")
        cv = getattr(c, metric, None) if c else None
        bv = getattr(b, metric, None) if b else None
        def fmt(v):
            if v is None: return "—"
            if isinstance(v, float) and metric == "next_token_accuracy":
                return f"{v:.1%}"
            if isinstance(v, float):
                return f"{v:.2f}"
            return str(v)
        print(f"{metric:24} {fmt(cv):>16} {fmt(bv):>16}")
    print("=" * 70)
    for r in results:
        print(f"\n[{r.tokenizer}] fluency probe ('once upon a time' ->):")
        print(f"    {r.fluency_sample!r}")
    print(f"\nwrote {args.out}", flush=True)
    return 0


if __name__ == "__main__":
    sys.exit(main())