from __future__ import annotations
import argparse
import json
import statistics
import sys
import time
from collections.abc import Sequence
from pathlib import Path
from typing import Any, cast
from ffbpe import BpeTrainer
from ffbpe import BoundaryMode
from ffbpe import PreTokenizer
from ffbpe import Unit
from ffbpe import UnicodeBigramMixedBoundary
from common import DEFAULT_CHUNK_SIZE
from common import SPECIAL_TOKENS
from common import add_report_args
from common import benchmark_metadata
from common import resolve_report_path
from common import write_word_inventory_manifest
from common import write_report
SCRIPT_NAME = "unicode_bigram_split"
def duration_summary(values: Sequence[float]) -> dict[str, Any]:
if not values:
return {
"count": 0,
"total_s": 0.0,
"min_s": None,
"median_s": None,
"max_s": None,
}
return {
"count": len(values),
"total_s": sum(values),
"min_s": min(values),
"median_s": statistics.median(values),
"max_s": max(values),
}
def word_length_summary(words: dict[str, int]) -> dict[str, Any]:
lengths = [len(word) for word in words]
if not lengths:
return {
"min": None,
"median": None,
"p90": None,
"max": None,
}
lengths.sort()
return {
"min": lengths[0],
"median": statistics.median(lengths),
"p90": lengths[round((len(lengths) - 1) * 0.90)],
"max": lengths[-1],
}
def inventory_summary(words: dict[str, int]) -> dict[str, Any]:
frequencies = list(words.values())
singleton_count = sum(1 for freq in frequencies if freq == 1)
return {
"unique_words": len(words),
"occurrences": sum(frequencies),
"singleton_count": singleton_count,
"singleton_ratio": singleton_count / len(words) if words else None,
"word_length": word_length_summary(words),
}
def save_words(path: Path, words: dict[str, int]) -> None:
sorted_words = dict(sorted(words.items(), key=lambda item: (item[1], item[0]), reverse=True))
path.write_text(json.dumps(sorted_words, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
def train_unitoken(words: dict[str, int], vocab_size: int, unit: Unit) -> dict[str, Any]:
trainer = BpeTrainer(
SPECIAL_TOKENS,
unit=unit,
initial_alphabet="byte_level" if unit == "byte" else None,
)
started = time.perf_counter()
trainer.add_words(list(words.items()))
add_words_s = time.perf_counter() - started
started = time.perf_counter()
trainer.init_training()
init_training_s = time.perf_counter() - started
step_times = []
while trainer.vocab_size < vocab_size:
started = time.perf_counter()
next_vocab_size = trainer.step()
step_times.append(time.perf_counter() - started)
if next_vocab_size is None:
break
return {
"vocab_size": trainer.vocab_size,
"final_merge_freq": trainer.last_merge_freq,
"add_words_s": add_words_s,
"init_training_s": init_training_s,
"step_summary": duration_summary(step_times),
"total_train_s": add_words_s + init_training_s + sum(step_times),
}
def collect_words(pretokenizer: PreTokenizer, path: Path, chunk_size: int, boundary: BoundaryMode) -> tuple[dict[str, int], float]:
started = time.perf_counter()
words = pretokenizer.get_words_from_file(path, chunk_size=chunk_size, boundary=boundary)
return words, time.perf_counter() - started
def main(argv: Sequence[str] | None = None) -> int:
parser = argparse.ArgumentParser(description="Experiment with Unicode bigram pre-splitting.")
parser.add_argument("--text", type=Path, required=True)
parser.add_argument("--chunk-size", type=int, default=DEFAULT_CHUNK_SIZE)
parser.add_argument("--boundary", choices=["auto", "eot", "line", "utf8"], default="auto")
parser.add_argument("--top-k", type=int, default=100_000)
parser.add_argument("--min-freq", type=int, default=16)
parser.add_argument("--unicode-bigram-mixed-boundary", choices=["keep", "split"], default="keep")
parser.add_argument("--unit", choices=["byte", "unicode"], default="unicode")
parser.add_argument("--vocab-size", type=int, help="Optionally train FFBPE on both inventories.")
parser.add_argument("--save-words", type=Path, help="Save the Unicode bigram split word-frequency inventory as JSON.")
add_report_args(parser)
args = parser.parse_args(argv)
if args.chunk_size < 1:
parser.error("--chunk-size must be at least 1")
if args.top_k < 1:
parser.error("--top-k must be at least 1")
if args.min_freq < 1:
parser.error("--min-freq must be at least 1")
boundary = cast(BoundaryMode, args.boundary)
unit = cast(Unit, args.unit)
unicode_bigram_mixed_boundary = cast(UnicodeBigramMixedBoundary, args.unicode_bigram_mixed_boundary)
base = PreTokenizer(SPECIAL_TOKENS, SPECIAL_TOKENS[0])
started = time.perf_counter()
bigram_selection = base.select_unicode_bigrams_from_file(
args.text,
chunk_size=args.chunk_size,
boundary=boundary,
top_k=args.top_k,
min_freq=args.min_freq,
)
unicode_bigrams = bigram_selection.bigrams
build_bigrams_s = time.perf_counter() - started
unicode_bigram_manifest = {
"top_k": args.top_k,
"min_freq": args.min_freq,
"mixed_boundary": unicode_bigram_mixed_boundary,
"retained": len(unicode_bigrams),
"cutoff_freq": bigram_selection.cutoff_freq,
"max_excluded_freq": bigram_selection.max_excluded_freq,
}
unicode_bigram_result = {
**unicode_bigram_manifest,
"build_s": build_bigrams_s,
}
baseline_words, baseline_pretokenize_s = collect_words(base, args.text, args.chunk_size, boundary)
split = PreTokenizer(
SPECIAL_TOKENS,
SPECIAL_TOKENS[0],
unicode_bigrams=unicode_bigrams,
unicode_bigram_mixed_boundary=unicode_bigram_mixed_boundary,
)
split_words, split_pretokenize_s = collect_words(split, args.text, args.chunk_size, boundary)
words_manifest = None
if args.save_words:
args.save_words.parent.mkdir(parents=True, exist_ok=True)
save_words(args.save_words, split_words)
words_manifest = write_word_inventory_manifest(
args.save_words,
source={
"input_kind": "raw_text",
"path": str(args.text),
"bytes": args.text.stat().st_size,
},
pretokenizer={
"special_tokens": SPECIAL_TOKENS,
"end_of_text": SPECIAL_TOKENS[0],
"pat_str": "default",
"chunk_size": args.chunk_size,
"boundary": boundary,
"unicode_bigram_mixed_boundary": unicode_bigram_mixed_boundary,
},
unicode_bigrams=unicode_bigram_manifest,
unique_words=len(split_words),
weighted_occurrences=sum(split_words.values()),
)
result: dict[str, Any] = {
"metadata": benchmark_metadata(
contract="unicode_bigram_split_experiment",
script_name=SCRIPT_NAME,
dataset_name=args.dataset_name,
config_name=args.config_name,
experiment_name=args.experiment_name,
notes=[
"Compares baseline pretokenization with unicode-bigram-guided splitting.",
"final_merge_at_or_above_bigram_cutoff is inclusive because cutoff-frequency ties are retained.",
],
),
"source": {
"input_kind": "raw_text",
"text": str(args.text),
"text_bytes": args.text.stat().st_size,
"chunk_size": args.chunk_size,
"boundary": boundary,
"unit": unit,
},
"unicode_bigram": unicode_bigram_result,
"baseline": {
"pretokenize_s": baseline_pretokenize_s,
**inventory_summary(baseline_words),
},
"unicode_bigram_split": {
"pretokenize_s": split_pretokenize_s,
**inventory_summary(split_words),
"words_manifest": str(words_manifest) if words_manifest else None,
},
}
if args.vocab_size:
result["target_vocab_size"] = args.vocab_size
result["baseline"]["training"] = train_unitoken(
baseline_words,
args.vocab_size,
unit,
)
split_training = train_unitoken(split_words, args.vocab_size, unit)
final_merge_freq = split_training["final_merge_freq"]
cutoff_freq = bigram_selection.cutoff_freq
max_excluded_freq = bigram_selection.max_excluded_freq
split_training["final_merge_at_or_above_bigram_cutoff"] = (
final_merge_freq >= cutoff_freq
if final_merge_freq is not None and cutoff_freq is not None
else None
)
split_training["final_merge_minus_bigram_cutoff"] = (
final_merge_freq - cutoff_freq
if final_merge_freq is not None and cutoff_freq is not None
else None
)
split_training["final_merge_above_max_excluded_freq"] = (
final_merge_freq > max_excluded_freq
if final_merge_freq is not None and max_excluded_freq is not None
else None
)
result["unicode_bigram_split"]["training"] = split_training
rendered = json.dumps(result, 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 __name__ == "__main__":
raise SystemExit(main(sys.argv[1:]))