SlayerLab/tokenizers

1#!/usr/bin/env python32"""Benchmark every tokenizer JSON row on one fixed Polish suite.3 4The benchmark is diagnostic, not a downstream-language-model quality score.5All tokenizers see exactly the same strings and are loaded from the lossless6``tokenizer_json`` column, so no contributor-specific files are required.7"""8 9from __future__ import annotations10 11import argparse12import hashlib13import importlib.metadata14import json15import platform16import re17import statistics18import time19from collections import Counter20from pathlib import Path21 22import pyarrow as pa23import pyarrow.parquet as pq24from tokenizers import Tokenizer25 26from custom_tokenizer_adapters import AdaptedTokenizer, load_custom_tokenizer_document27 28 29WORD_RE = re.compile(r"\w+", re.UNICODE)30 31 32def parse_args() -> argparse.Namespace:33    parser = argparse.ArgumentParser()34    parser.add_argument("--input", type=Path, default=Path("data/train-00000-of-00001.parquet"))35    parser.add_argument("--suite", type=Path, default=Path("evaluation/polish_suite.jsonl"))36    parser.add_argument("--output", type=Path, default=Path("results/tokenizer_benchmark.parquet"))37    parser.add_argument("--summary-json", type=Path, default=Path("results/tokenizer_benchmark.json"))38    parser.add_argument("--repeats", type=int, default=15, help="Timed batch repetitions")39    return parser.parse_args()40 41 42def load_suite(path: Path) -> list[dict[str, str]]:43    records = [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]44    if not records or any(set(record) != {"domain", "text"} for record in records):45        raise ValueError("suite must contain non-empty JSONL records with exactly domain and text")46    return records47 48 49def timed_median(function, repeats: int) -> float:50    samples = []51    for _ in range(repeats):52        start = time.perf_counter_ns()53        function()54        samples.append((time.perf_counter_ns() - start) / 1e9)55    return statistics.median(samples)56 57 58def encode_batch(tokenizer: Tokenizer | AdaptedTokenizer, texts: list[str]) -> list[list[int]]:59    if isinstance(tokenizer, Tokenizer):60        return [encoding.ids for encoding in tokenizer.encode_batch(texts, add_special_tokens=False)]61    return [tokenizer.encode(text) for text in texts]62 63 64def decode_batch(tokenizer: Tokenizer | AdaptedTokenizer, batches: list[list[int]]) -> list[str]:65    if isinstance(tokenizer, Tokenizer):66        return tokenizer.decode_batch(batches, skip_special_tokens=False)67    return [tokenizer.decode(ids) for ids in batches]68 69 70def counts_for(tokenizer: Tokenizer | AdaptedTokenizer, records: list[dict[str, str]]) -> tuple[dict, list, list]:71    texts = [record["text"] for record in records]72    id_batches = encode_batch(tokenizer, texts)73    unk_id = tokenizer.token_to_id("<unk>") if isinstance(tokenizer, Tokenizer) else None74    model_unk = None75    try:76        model_unk = json.loads(tokenizer.to_str()).get("model", {}).get("unk_token") if isinstance(tokenizer, Tokenizer) else None77    except (TypeError, json.JSONDecodeError):78        pass79    if model_unk:80        unk_id = tokenizer.token_to_id(model_unk)81 82    result = Counter()83    failures = []84    decoded_batch = decode_batch(tokenizer, id_batches)85    for record, ids, decoded in zip(records, id_batches, decoded_batch):86        text = record["text"]87        result["texts"] += 188        result["chars"] += len(text)89        result["bytes"] += len(text.encode("utf-8"))90        result["words"] += len(WORD_RE.findall(text))91        result["tokens"] += len(ids)92        if unk_id is not None:93            result["unks"] += sum(token_id == unk_id for token_id in ids)94        if decoded != text:95            result["roundtrip_failures"] += 196            failures.append({"domain": record["domain"], "text": text, "decoded": decoded})97    return dict(result), id_batches, failures98 99 100def ratios(counts: dict) -> dict:101    tokens = counts.get("tokens", 0)102    words = counts.get("words", 0)103    return {104        "text_count": counts.get("texts", 0),105        "word_count": words,106        "token_count": tokens,107        "tokens_per_word": tokens / words if words else None,108        "chars_per_token": counts.get("chars", 0) / tokens if tokens else None,109        "bytes_per_token": counts.get("bytes", 0) / tokens if tokens else None,110        "unk_rate": counts.get("unks", 0) / tokens if tokens else None,111        "roundtrip_failures": counts.get("roundtrip_failures", 0),112        "roundtrip_pass": counts.get("roundtrip_failures", 0) == 0,113    }114 115 116def main() -> None:117    args = parse_args()118    if args.repeats < 1:119        raise ValueError("--repeats must be positive")120    suite = load_suite(args.suite)121    table = pq.read_table(args.input)122    rows = table.to_pylist()123    texts = [record["text"] for record in suite]124    total_bytes = sum(len(text.encode("utf-8")) for text in texts)125    results = []126 127    for source in rows:128        base = {key: source[key] for key in ("author", "name", "size", "sha256", "source_path")}129        try:130            if source["hf_loadable"]:131                tokenizer: Tokenizer | AdaptedTokenizer = Tokenizer.from_str(source["tokenizer_json"])132                adapter_status = "native"133                adapter_fidelity = "exact"134                adapter_note = "Native Hugging Face Tokenizers artifact; no adapter used."135                adapter_source_format = "hf_tokenizers"136                runtime = "rust_tokenizers"137            else:138                adapted = load_custom_tokenizer_document(139                    json.loads(source["tokenizer_json"]), source["source_path"]140                )141                tokenizer = adapted142                adapter_status = "custom_adapter"143                # Only fully serialized behavior is exact. Documented-but-not-144                # serialized pre-tokenization remains core_only for judging.145                adapter_fidelity = "exact" if adapted.fidelity == "exact" else "core_only"146                adapter_note = adapted.fidelity_note147                adapter_source_format = adapted.source_format148                runtime = "python_reference_adapter"149            encode_batch(tokenizer, texts)  # warm-up150            total, id_batches, failures = counts_for(tokenizer, suite)151            encode_seconds = timed_median(152                lambda: encode_batch(tokenizer, texts), args.repeats153            )154            decode_seconds = timed_median(155                lambda: decode_batch(tokenizer, id_batches), args.repeats156            )157            domains = {}158            for domain in sorted({record["domain"] for record in suite}):159                domain_counts, _, _ = counts_for(160                    tokenizer, [record for record in suite if record["domain"] == domain]161                )162                domains[domain] = ratios(domain_counts)163            result = {164                **base,165                "status": "ok",166                "adapter_status": adapter_status,167                "adapter_fidelity": adapter_fidelity,168                "adapter_note": adapter_note,169                "adapter_source_format": adapter_source_format,170                "runtime": runtime,171                **ratios(total),172                "encode_mb_per_s": total_bytes / 1_000_000 / encode_seconds,173                "decode_mb_per_s": total_bytes / 1_000_000 / decode_seconds,174                "serialized_bytes": len(source["tokenizer_json"].encode("utf-8")),175                "domain_metrics_json": json.dumps(domains, ensure_ascii=False, sort_keys=True),176                "roundtrip_examples_json": json.dumps(failures[:3], ensure_ascii=False),177                "error": "",178            }179        except Exception as exc:  # retain failed rows instead of silently changing the cohort180            result = {**base, "status": "error", "error": f"{type(exc).__name__}: {exc}"}181        results.append(result)182 183    metadata = {184        "suite": str(args.suite),185        "suite_sha256": hashlib.sha256(args.suite.read_bytes()).hexdigest(),186        "suite_records": len(suite),187        "suite_bytes": total_bytes,188        "repeats": args.repeats,189        "timer": "median wall-clock batch time after one warm-up",190        "python": platform.python_version(),191        "tokenizers": importlib.metadata.version("tokenizers"),192        "pyarrow": importlib.metadata.version("pyarrow"),193        "platform": platform.platform(),194    }195    args.output.parent.mkdir(parents=True, exist_ok=True)196    args.summary_json.parent.mkdir(parents=True, exist_ok=True)197    pq.write_table(pa.Table.from_pylist(results), args.output, compression="zstd")198    args.summary_json.write_text(199        json.dumps({"benchmark_metadata": metadata, "results": results}, ensure_ascii=False, indent=2),200        encoding="utf-8",201    )202    print(json.dumps(metadata, indent=2))203    for result in sorted(results, key=lambda row: (row.get("tokens_per_word", float("inf")), row["author"])):204        print(205            result["author"], result["name"], result["status"],206            f"tpw={result.get('tokens_per_word', float('nan')):.4f}",207            f"roundtrip_failures={result.get('roundtrip_failures', '-')}"208        )209 210 211if __name__ == "__main__":212    main()213