from __future__ import annotations
import argparse
import gc
import json
import statistics
import sys
import time
from collections.abc import Callable, Iterable, Sequence
from pathlib import Path
from typing import Any, cast
from tokenizers import Tokenizer
from tokenizers import models
from tokenizers import pre_tokenizers
from tokenizers import trainers
from ffbpe import BpeTrainer
from ffbpe import BoundaryMode
from ffbpe import PreTokenizer
from ffbpe import WordCounter
from common import DEFAULT_CHUNK_SIZE
from common import REPO_ROOT
from common import SPECIAL_TOKENS
from common import add_report_args
from common import benchmark_metadata
from common import bigram_frequency_guard
from common import load_words
from common import load_word_inventory_manifest
from common import resolve_report_path
from common import word_inventory_manifest_path
from common import write_report
DEFAULT_WORDS = REPO_ROOT / "fixtures" / "_words.tinystories_sample_5M.json"
DEFAULT_HOT_PAIR_WINDOW_SIZE = 4096
SCRIPT_NAME = "compare_hf_training"
def bytes_to_unicode() -> dict[int, str]:
bs = (
list(range(ord("!"), ord("~") + 1))
+ list(range(ord("¡"), ord("¬") + 1))
+ list(range(ord("®"), ord("ÿ") + 1))
)
cs = bs[:]
n = 0
for byte in range(256):
if byte not in bs:
bs.append(byte)
cs.append(256 + n)
n += 1
return dict(zip(bs, map(chr, cs)))
BYTE_ENCODER = bytes_to_unicode()
def to_byte_level_token(token: bytes) -> str:
return "".join(BYTE_ENCODER[byte] for byte in token)
def expanded_words(words: Sequence[tuple[str, int]]) -> Iterable[str]:
for word, freq in words:
for _ in range(freq):
yield word
def new_unitoken_trainer(hot_pair_window_size: int) -> BpeTrainer:
return BpeTrainer(
SPECIAL_TOKENS,
unit="byte",
initial_alphabet="byte_level",
hot_pair_window_size=hot_pair_window_size,
)
def unitoken_result(trainer: BpeTrainer) -> dict[str, Any]:
vocab = {
to_byte_level_token(token): rank
for token, rank in trainer.vocab.items()
}
return {
"vocab": vocab,
"vocab_size": trainer.vocab_size,
"final_merge_freq": trainer.last_merge_freq,
"hot_pair_window_stats": trainer.hot_pair_window_stats,
}
def train_unitoken(
words: Sequence[tuple[str, int]],
vocab_size: int,
hot_pair_window_size: int,
) -> dict[str, Any]:
trainer = new_unitoken_trainer(hot_pair_window_size)
trainer.add_words(words)
trainer.train(vocab_size)
return unitoken_result(trainer)
def train_unitoken_from_counter(
counter: WordCounter,
vocab_size: int,
hot_pair_window_size: int,
) -> dict[str, Any]:
trainer = new_unitoken_trainer(hot_pair_window_size)
trainer.add_word_counter(counter)
trainer.train(vocab_size)
return unitoken_result(trainer)
def train_hugging_face(words: Sequence[tuple[str, int]], vocab_size: int) -> dict[str, Any]:
tokenizer = Tokenizer(models.BPE())
tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
trainer = trainers.BpeTrainer(
vocab_size=vocab_size,
special_tokens=SPECIAL_TOKENS,
initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
)
tokenizer.train_from_iterator(
expanded_words(words),
trainer=trainer,
length=sum(freq for _, freq in words),
)
return {
"vocab": tokenizer.get_vocab(),
"vocab_size": tokenizer.get_vocab_size(),
}
def iter_text_segments(path: Path, segments: Sequence[tuple[int, int]]) -> Iterable[str]:
with path.open("rb") as file:
for offset, length in segments:
file.seek(offset)
yield file.read(length).decode("utf-8")
def iter_text_chunks(path: Path, chunk_bytes: int) -> Iterable[str]:
pending = b""
with path.open("rb") as file:
while True:
chunk = file.read(chunk_bytes)
if not chunk:
break
data = pending + chunk
boundary = len(data)
while boundary > 0:
try:
text = data[:boundary].decode("utf-8")
break
except UnicodeDecodeError:
boundary -= 1
else:
pending = data
continue
yield text
pending = data[boundary:]
if pending:
yield pending.decode("utf-8")
def train_unitoken_from_text(
path: Path,
vocab_size: int,
chunk_size: int,
boundary: BoundaryMode,
hot_pair_window_size: int,
) -> dict[str, Any]:
started = time.perf_counter()
pretokenizer = PreTokenizer(SPECIAL_TOKENS, SPECIAL_TOKENS[0])
segments = pretokenizer.find_chunk_boundaries(path, chunk_size=chunk_size, boundary=boundary)
counter = pretokenizer.word_counter()
counter.add_source(iter_text_segments(path, segments))
pretokenize_s = time.perf_counter() - started
unique_words = counter.len
started = time.perf_counter()
train_result = train_unitoken_from_counter(counter, vocab_size, hot_pair_window_size)
train_s = time.perf_counter() - started
train_result.update({
"pretokenize_s": pretokenize_s,
"train_s": train_s,
"unique_words": unique_words,
"occurrences": None,
})
return train_result
def train_hugging_face_from_text(
path: Path,
vocab_size: int,
*,
chunk_bytes: int | None,
segments: Sequence[tuple[int, int]] | None,
) -> dict[str, Any]:
tokenizer = Tokenizer(models.BPE())
tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
trainer = trainers.BpeTrainer(
vocab_size=vocab_size,
special_tokens=SPECIAL_TOKENS,
initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
)
iterator = iter_text_segments(path, segments) if segments is not None else iter_text_chunks(path, chunk_bytes or 1)
tokenizer.train_from_iterator(
iterator,
trainer=trainer,
)
return {
"vocab": tokenizer.get_vocab(),
"vocab_size": tokenizer.get_vocab_size(),
}
def time_call(label: str, fn: Callable[[], dict[str, Any]], repeats: int) -> dict[str, Any]:
samples = []
result = None
for _ in range(repeats):
gc.collect()
started = time.perf_counter()
result = fn()
samples.append(time.perf_counter() - started)
assert result is not None
timed = {
"label": label,
"repeats": repeats,
"min_s": min(samples),
"median_s": statistics.median(samples),
"mean_s": statistics.mean(samples),
"vocab_size": result["vocab_size"],
"vocab": result["vocab"],
}
for key, value in result.items():
if key in timed or key == "vocab":
continue
timed[key] = value
return timed
def without_vocab(result: dict[str, Any]) -> dict[str, Any]:
return {key: value for key, value in result.items() if key != "vocab"}
def run_words(args: argparse.Namespace) -> dict[str, Any]:
words = load_words(args.words, args.max_occurrences)
manifest = load_word_inventory_manifest(args.words)
occurrence_count = sum(freq for _, freq in words)
unitoken = time_call(
"unitoken.train",
lambda: train_unitoken(words, args.vocab_size, args.hot_pair_window_size),
args.repeats,
)
guard = (
bigram_frequency_guard(unitoken["final_merge_freq"], manifest)
if args.max_occurrences is None
else None
)
if guard is not None:
unitoken["bigram_frequency_guard"] = guard
hf = time_call(
"huggingface.train",
lambda: train_hugging_face(words, args.vocab_size),
args.repeats,
)
same_vocab = unitoken["vocab"] == hf["vocab"]
unitoken_median = unitoken["median_s"]
hf_median = hf["median_s"]
speedup = hf_median / unitoken_median if unitoken_median else None
results = {
"metadata": benchmark_metadata(
contract="fixed_words_unitoken_vs_hf_expanded_iterator",
script_name=SCRIPT_NAME,
dataset_name=args.dataset_name,
config_name=args.config_name,
experiment_name=args.experiment_name,
notes=[
"Unitoken receives compressed (word, frequency) pairs.",
"Unitoken bounds persistent pair occurrence postings with the configured hot-pair window.",
"Hugging Face receives an expanded iterator of repeated words because the Python tokenizers API does not accept compressed counts.",
"Use raw_text_unitoken_vs_hf for end-to-end implementation comparison.",
],
),
"source": {
"input_kind": "words_json",
"words": str(args.words),
"unique_words": len(words),
"occurrences": occurrence_count,
"huggingface_input_kind": "expanded_word_iterator",
"unitoken_input_kind": "compressed_word_counts",
"word_inventory_manifest_path": (
str(word_inventory_manifest_path(args.words))
if manifest is not None
else None
),
"word_inventory_manifest": manifest,
},
"target_vocab_size": args.vocab_size,
"hot_pair_window_size": args.hot_pair_window_size,
"same_vocab": same_vocab,
"speedup_hf_over_unitoken_median": speedup,
"benchmarks": [
without_vocab(unitoken),
without_vocab(hf),
],
}
if not same_vocab:
unitoken_vocab = unitoken["vocab"]
hf_vocab = hf["vocab"]
results["vocab_diff"] = {
"unitoken_only": sorted(set(unitoken_vocab) - set(hf_vocab))[:args.diff_limit],
"huggingface_only": sorted(set(hf_vocab) - set(unitoken_vocab))[:args.diff_limit],
"rank_mismatches": [
{
"token": token,
"unitoken_rank": unitoken_vocab[token],
"huggingface_rank": hf_vocab[token],
}
for token in sorted(set(unitoken_vocab) & set(hf_vocab))
if unitoken_vocab[token] != hf_vocab[token]
][:args.diff_limit],
}
return results
def run_text(args: argparse.Namespace) -> dict[str, Any]:
boundary = cast(BoundaryMode, args.boundary)
hf_segments = None
hf_chunking = f"fixed {args.hf_chunk_bytes} byte chunks"
if args.hf_chunk_bytes is None:
pretokenizer = PreTokenizer(SPECIAL_TOKENS, SPECIAL_TOKENS[0])
hf_segments = pretokenizer.find_chunk_boundaries(
args.text,
chunk_size=args.chunk_size,
boundary=boundary,
)
hf_chunking = "FFBPE chunk boundaries"
unitoken = time_call(
"unitoken.raw_train",
lambda: train_unitoken_from_text(
args.text,
args.vocab_size,
args.chunk_size,
boundary,
args.hot_pair_window_size,
),
args.repeats,
)
hf = time_call(
"huggingface.raw_train",
lambda: train_hugging_face_from_text(
args.text,
args.vocab_size,
chunk_bytes=args.hf_chunk_bytes,
segments=hf_segments,
),
args.repeats,
)
same_vocab = unitoken["vocab"] == hf["vocab"]
unitoken_median = unitoken["median_s"]
hf_median = hf["median_s"]
speedup = hf_median / unitoken_median if unitoken_median else None
unitoken_train_s = unitoken.get("train_s")
train_phase_speedup = hf_median / unitoken_train_s if unitoken_train_s else None
return {
"metadata": benchmark_metadata(
contract="raw_text_unitoken_vs_hf",
script_name=SCRIPT_NAME,
dataset_name=args.dataset_name,
config_name=args.config_name,
experiment_name=args.experiment_name,
notes=[
"Raw-text mode compares end-to-end training contracts.",
"FFBPE timing includes native WordCounter pretokenization plus bounded-window training.",
"The native WordCounter is consumed directly without materializing the full inventory in Python.",
"Raw FFBPE occurrence count is null because computing it through the current Python API would copy the full inventory.",
"By default Hugging Face receives FFBPE chunk boundaries so vocab parity reflects tokenizer/trainer behavior instead of iterator boundary differences.",
],
),
"source": {
"input_kind": "raw_text",
"text": str(args.text),
"text_bytes": args.text.stat().st_size,
"boundary": boundary,
"chunk_size": args.chunk_size,
"huggingface_chunking": hf_chunking,
"unitoken_input_kind": "native_word_counter",
},
"huggingface_chunking": hf_chunking,
"target_vocab_size": args.vocab_size,
"hot_pair_window_size": args.hot_pair_window_size,
"same_vocab": same_vocab,
"speedup_hf_over_unitoken_median": speedup,
"speedup_hf_over_unitoken_train_phase": train_phase_speedup,
"benchmarks": [
without_vocab(unitoken),
without_vocab(hf),
],
"note": "Raw-text mode includes FFBPE pretokenization plus training. By default Hugging Face receives FFBPE's chunk boundaries so vocab parity reflects tokenizer/trainer behavior instead of iterator boundary differences.",
}
def main(argv: Sequence[str] | None = None) -> int:
parser = argparse.ArgumentParser(description="Benchmark FFBPE training against Hugging Face tokenizers.")
input_group = parser.add_mutually_exclusive_group()
input_group.add_argument("--words", type=Path, help="JSON word-frequency fixture.")
input_group.add_argument("--text", type=Path, help="Raw UTF-8 text file for end-to-end training.")
parser.add_argument("--vocab-size", type=int, default=2000)
parser.add_argument("--repeats", type=int, default=3)
parser.add_argument(
"--hot-pair-window-size",
type=int,
default=DEFAULT_HOT_PAIR_WINDOW_SIZE,
help="Retain occurrence postings for only the exact top-K FFBPE pair window.",
)
parser.add_argument("--max-occurrences", type=int, help="Truncate the weighted corpus for a faster smoke benchmark.")
parser.add_argument("--chunk-size", type=int, default=DEFAULT_CHUNK_SIZE, help="Target FFBPE pretokenizer chunk size in bytes for --text mode.")
parser.add_argument("--boundary", choices=["auto", "eot", "line", "utf8"], default="auto", help="Boundary strategy for FFBPE chunking in --text mode.")
parser.add_argument("--hf-chunk-bytes", type=int, help="Force fixed byte chunks for Hugging Face in --text mode. Defaults to FFBPE chunk boundaries.")
parser.add_argument("--diff-limit", type=int, default=10)
add_report_args(parser)
args = parser.parse_args(argv)
if args.words is None and args.text is None:
args.words = DEFAULT_WORDS
if args.repeats < 1:
parser.error("--repeats must be at least 1")
if args.hot_pair_window_size < 1:
parser.error("--hot-pair-window-size must be at least 1")
if args.chunk_size < 1:
parser.error("--chunk-size must be at least 1")
if args.hf_chunk_bytes is not None and args.hf_chunk_bytes < 1:
parser.error("--hf-chunk-bytes must be at least 1")
results = run_text(args) if args.text else run_words(args)
rendered = json.dumps(results, indent=2)
if not args.quiet:
print(rendered)
write_report(resolve_report_path(args, script_name=SCRIPT_NAME, vocab_size=args.vocab_size), rendered)
return 0 if args.text or results["same_vocab"] else 1
if __name__ == "__main__":
raise SystemExit(main(sys.argv[1:]))