Upload benchmark_tiktoken_style.py with huggingface_hub

benchmark_tiktoken_style.py

@@ -0,0 +1,264 @@
+#!/usr/bin/env python3
+"""
+Tiktoken-style benchmark comparing SARFTokenizer vs tiktoken vs HuggingFace.
+
+Measures throughput in MB/s with proper thread isolation using multiprocessing.
+
+Usage:
+    python benchmark_tiktoken_style.py --samples 1000000 --threads 1 2 4 8
+"""
+
+import os
+import sys
+import time
+import argparse
+from pathlib import Path
+from typing import List, Tuple
+from multiprocessing import Process, Queue, cpu_count
+
+import pyarrow.parquet as pq
+
+# Add parent to path
+sys.path.insert(0, str(Path(__file__).parent))
+
+# Configuration
+DATA_DIR = "/root/.cache/deeplatent/base_data/"
+HF_TOKENIZER_PATH = os.path.expanduser("~/.cache/deeplatent/tokenizers/SARFTokenizer")
+DEFAULT_THREADS = [2**i for i in range(8) if 2**i <= cpu_count()]
+
+
+def format_byte_size(num_bytes: float) -> Tuple[str, str]:
+    """Convert bytes to human-readable format."""
+    for unit in ["B", "KB", "MB", "GB", "TB"]:
+        if num_bytes < 1024:
+            return f"{num_bytes:.2f} {unit}", unit
+        num_bytes /= 1024
+    return f"{num_bytes:.2f} PB", "PB"
+
+
+def load_samples(data_dir: str, num_samples: int) -> Tuple[List[str], int]:
+    """Load samples from parquet files."""
+    import re
+    AR_DETECT = re.compile(r'[\u0600-\u06FF]')
+
+    parquet_files = sorted(Path(data_dir).glob("shard_*.parquet"))
+    if not parquet_files:
+        raise FileNotFoundError(f"No parquet files found in {data_dir}")
+
+    samples = []
+    target = num_samples
+
+    for pq_file in parquet_files:
+        if len(samples) >= target:
+            break
+
+        table = pq.read_table(pq_file, columns=["text"])
+        texts = table.column("text").to_pylist()
+
+        for text in texts:
+            if len(samples) >= target:
+                break
+            if text and isinstance(text, str):
+                samples.append(text)
+
+    total_bytes = sum(len(t.encode('utf-8')) for t in samples)
+    return samples, total_bytes
+
+
+def benchmark_sarf(documents: List[str], num_threads: int, result_queue: Queue):
+    """Benchmark SARFTokenizer."""
+    from deeplatent import SARFTokenizer
+
+    os.environ["RAYON_NUM_THREADS"] = str(num_threads)
+
+    tok = SARFTokenizer.from_pretrained(HF_TOKENIZER_PATH)
+    num_bytes = sum(len(d.encode('utf-8')) for d in documents)
+
+    # Warmup
+    tok.encode(documents[0])
+
+    # Benchmark
+    start = time.perf_counter_ns()
+    if hasattr(tok, 'encode_batch'):
+        tok.encode_batch(documents)
+    else:
+        for d in documents:
+            tok.encode(d)
+    end = time.perf_counter_ns()
+
+    elapsed_ns = end - start
+    bytes_per_sec = num_bytes / elapsed_ns * 1e9
+    texts_per_sec = len(documents) / elapsed_ns * 1e9
+
+    result_queue.put(("SARFTokenizer", bytes_per_sec, texts_per_sec))
+
+
+def benchmark_tiktoken(documents: List[str], num_threads: int, encoding: str, result_queue: Queue):
+    """Benchmark tiktoken."""
+    import tiktoken
+
+    os.environ["RAYON_NUM_THREADS"] = str(num_threads)
+
+    enc = tiktoken.get_encoding(encoding)
+    num_bytes = sum(len(d.encode('utf-8')) for d in documents)
+
+    # Warmup
+    enc.encode(documents[0])
+
+    # Benchmark
+    start = time.perf_counter_ns()
+    enc.encode_ordinary_batch(documents, num_threads=num_threads)
+    end = time.perf_counter_ns()
+
+    elapsed_ns = end - start
+    bytes_per_sec = num_bytes / elapsed_ns * 1e9
+    texts_per_sec = len(documents) / elapsed_ns * 1e9
+
+    result_queue.put((f"tiktoken ({encoding})", bytes_per_sec, texts_per_sec))
+
+
+def benchmark_hf_tokenizers(documents: List[str], num_threads: int, result_queue: Queue):
+    """Benchmark HuggingFace tokenizers."""
+    from tokenizers import Tokenizer
+
+    os.environ["RAYON_NUM_THREADS"] = str(num_threads)
+
+    # Load the SARFTokenizer's underlying HF tokenizer
+    tokenizer_path = os.path.join(HF_TOKENIZER_PATH, "tokenizer.json")
+    tok = Tokenizer.from_file(tokenizer_path)
+    num_bytes = sum(len(d.encode('utf-8')) for d in documents)
+
+    # Warmup
+    tok.encode(documents[0])
+
+    # Benchmark
+    start = time.perf_counter_ns()
+    tok.encode_batch_fast(documents)
+    end = time.perf_counter_ns()
+
+    elapsed_ns = end - start
+    bytes_per_sec = num_bytes / elapsed_ns * 1e9
+    texts_per_sec = len(documents) / elapsed_ns * 1e9
+
+    result_queue.put(("HF tokenizers", bytes_per_sec, texts_per_sec))
+
+
+def run_benchmark(documents: List[str], num_threads: int, num_bytes: int):
+    """Run benchmarks for all tokenizers with given thread count."""
+    readable_size, _ = format_byte_size(num_bytes)
+    avg_len = sum(len(d) for d in documents) / len(documents)
+
+    print(f"\n{'='*70}")
+    print(f"Threads: {num_threads}, Data: {readable_size}, Documents: {len(documents):,}, Avg Length: {avg_len:.0f}")
+    print(f"{'='*70}")
+
+    results = []
+
+    # SARFTokenizer
+    q = Queue()
+    p = Process(target=benchmark_sarf, args=(documents, num_threads, q))
+    p.start()
+    p.join()
+    if not q.empty():
+        name, bps, tps = q.get()
+        readable, _ = format_byte_size(bps)
+        print(f"{name:<20}\t{readable}/s\t({tps:,.0f} texts/s)")
+        results.append((name, bps, tps))
+
+    # tiktoken o200k_base
+    q = Queue()
+    p = Process(target=benchmark_tiktoken, args=(documents, num_threads, "o200k_base", q))
+    p.start()
+    p.join()
+    if not q.empty():
+        name, bps, tps = q.get()
+        readable, _ = format_byte_size(bps)
+        print(f"{name:<20}\t{readable}/s\t({tps:,.0f} texts/s)")
+        results.append((name, bps, tps))
+
+    # tiktoken cl100k_base
+    q = Queue()
+    p = Process(target=benchmark_tiktoken, args=(documents, num_threads, "cl100k_base", q))
+    p.start()
+    p.join()
+    if not q.empty():
+        name, bps, tps = q.get()
+        readable, _ = format_byte_size(bps)
+        print(f"{name:<20}\t{readable}/s\t({tps:,.0f} texts/s)")
+        results.append((name, bps, tps))
+
+    # HuggingFace tokenizers
+    q = Queue()
+    p = Process(target=benchmark_hf_tokenizers, args=(documents, num_threads, q))
+    p.start()
+    p.join()
+    if not q.empty():
+        name, bps, tps = q.get()
+        readable, _ = format_byte_size(bps)
+        print(f"{name:<20}\t{readable}/s\t({tps:,.0f} texts/s)")
+        results.append((name, bps, tps))
+
+    return results
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Tiktoken-style tokenizer benchmark")
+    parser.add_argument("--samples", type=int, default=10000, help="Number of samples")
+    parser.add_argument("--threads", type=int, nargs="+", default=DEFAULT_THREADS, help="Thread counts")
+    parser.add_argument("--data-dir", type=str, default=DATA_DIR, help="Data directory")
+    args = parser.parse_args()
+
+    print("=" * 70)
+    print("TIKTOKEN-STYLE TOKENIZER BENCHMARK")
+    print("=" * 70)
+    print(f"CPU count: {cpu_count()}")
+    print(f"Samples: {args.samples:,}")
+    print(f"Threads: {args.threads}")
+
+    # Load data
+    print("\nLoading data...")
+    documents, total_bytes = load_samples(args.data_dir, args.samples)
+    readable_size, _ = format_byte_size(total_bytes)
+    print(f"Loaded {len(documents):,} documents ({readable_size})")
+
+    # Run benchmarks
+    all_results = {}
+    for num_threads in args.threads:
+        results = run_benchmark(documents, num_threads, total_bytes)
+        all_results[num_threads] = results
+
+    # Summary table
+    print("\n" + "=" * 100)
+    print("SUMMARY TABLE (MB/s)")
+    print("=" * 100)
+
+    # Header
+    header = f"{'Tokenizer':<25}"
+    for t in args.threads:
+        header += f"{t}T".rjust(15)
+    print(header)
+    print("-" * 100)
+
+    # Collect by tokenizer name
+    tokenizers = {}
+    for threads, results in all_results.items():
+        for name, bps, tps in results:
+            if name not in tokenizers:
+                tokenizers[name] = {}
+            tokenizers[name][threads] = bps / 1024 / 1024  # Convert to MB/s
+
+    # Print rows
+    for name, thread_results in tokenizers.items():
+        row = f"{name:<25}"
+        for t in args.threads:
+            if t in thread_results:
+                row += f"{thread_results[t]:>14.2f}"
+            else:
+                row += "N/A".rjust(15)
+        print(row)
+
+    print("=" * 100)
+
+
+if __name__ == "__main__":
+    main()