unitoken 0.1.5

Fast BPE tokenizer/trainer with a Rust core and Python bindings
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111

from __future__ import annotations

import argparse
import gc
import json
import sys
import time
from collections.abc import Sequence
from pathlib import Path
from typing import Any

from uni_tokenizer import BpeTrainer

from common import SPECIAL_TOKENS
from common import add_report_args
from common import benchmark_metadata
from common import bucket_steps
from common import duration_summary
from common import load_words
from common import resolve_report_path
from common import write_report


SCRIPT_NAME = "profile_training_core"


def profile_training_core(words: Sequence[tuple[str, int]], vocab_size: int, bucket_size: int, unit: str) -> dict[str, Any]:
  gc.collect()
  trainer = BpeTrainer(
    SPECIAL_TOKENS,
    unit=unit,
    initial_alphabet="byte_level" if unit == "byte" else None,
  )

  started = time.perf_counter()
  trainer.add_words(words)
  add_words_s = time.perf_counter() - started
  initial_vocab_size = trainer.vocab_size

  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,
    "initial_vocab_size": initial_vocab_size,
    "add_words_s": add_words_s,
    "init_training_s": init_training_s,
    "step_summary": duration_summary(step_times),
    "step_buckets": bucket_steps(step_times, bucket_size),
    "total_train_s": add_words_s + init_training_s + sum(step_times),
  }


def main(argv: Sequence[str] | None = None) -> int:
  parser = argparse.ArgumentParser(description="Profile unitoken BPE training core from a compressed word-frequency inventory.")
  parser.add_argument("--words", type=Path, required=True, help="JSON word-frequency inventory.")
  parser.add_argument("--vocab-size", type=int, default=10000)
  parser.add_argument("--unit", choices=["byte", "unicode"], default="byte", help="BPE unit used for training.")
  parser.add_argument("--max-occurrences", type=int, help="Truncate the weighted corpus for a faster smoke profile.")
  parser.add_argument("--bucket-size", type=int, default=500)
  add_report_args(parser)
  args = parser.parse_args(argv)

  if args.vocab_size < 1:
    parser.error("--vocab-size must be at least 1")
  if args.bucket_size < 1:
    parser.error("--bucket-size must be at least 1")

  words = load_words(args.words, args.max_occurrences)
  result = {
    "metadata": benchmark_metadata(
      contract="fixed_words_unitoken_training_core_profile",
      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.",
        "This isolates training core phases and excludes pretokenization and external library comparisons.",
      ],
    ),
    "source": {
      "input_kind": "words_json",
      "words": str(args.words),
      "unique_words": len(words),
      "occurrences": sum(freq for _, freq in words),
      "unitoken_input_kind": "compressed_word_counts",
      "unit": args.unit,
    },
    "target_vocab_size": args.vocab_size,
    "unitoken": profile_training_core(words, args.vocab_size, args.bucket_size, args.unit),
  }

  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:]))