spb/forge Public MIT

1# Author: Simon-Pierre Boucher — contact@spboucher.ai2"""Cross-check the Python (vectorized) and C++ (greedy lowest-id) BPE3encoders on real text.45Encodes a text sample with tools/prepare_data.py's encoder, dumps the ids,6then runs test_tokenizer with FORGE_TOK_* pointing at them so the C++ side7compares. Exits nonzero on disagreement.89Usage:10  python3 tests/tokenizer_agreement.py --model data/tinystories/tok4096.model \11      --text data/tinystories/TinyStoriesV2-GPT4-valid.txt --bytes 200000 \12      --binary build/test_tokenizer13"""14import argparse15import os16import subprocess17import sys18import tempfile1920sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "tools"))21import prepare_data as pd  # noqa: E402222324def main():25    ap = argparse.ArgumentParser()26    ap.add_argument("--model", required=True)27    ap.add_argument("--text", required=True)28    ap.add_argument("--bytes", type=int, default=200000)29    ap.add_argument("--binary", default="build/test_tokenizer")30    args = ap.parse_args()3132    with open(args.text, "rb") as f:33        sample = f.read(args.bytes)34    # cut at a story boundary so the comparison matches how .bin is built35    sep = b"<|endoftext|>"36    if sep in sample:37        sample = sample[: sample.rindex(sep)]3839    pd._MERGES = pd.load_merges(args.model)40    ids = pd.encode_block(sample)41    print(f"python: {len(sample)} bytes -> {len(ids)} tokens")4243    with tempfile.TemporaryDirectory() as td:44        text_path = os.path.join(td, "text.bin")45        ids_path = os.path.join(td, "ids.txt")46        with open(text_path, "wb") as f:47            f.write(sample)48        with open(ids_path, "w") as f:49            f.write("\n".join(str(int(i)) for i in ids))5051        env = dict(os.environ)52        env["FORGE_TOK_MODEL"] = os.path.abspath(args.model)53        env["FORGE_TOK_TEXT"] = text_path54        env["FORGE_TOK_IDS"] = ids_path55        rc = subprocess.call([os.path.abspath(args.binary)], env=env,56                             cwd=os.path.dirname(os.path.abspath(args.binary)))57    sys.exit(rc)585960if __name__ == "__main__":61    main()62