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

from __future__ import annotations

import argparse
import hashlib
import sys
import time
from collections.abc import Sequence
from pathlib import Path

from uni_tokenizer import BpeTrainer

from common import SPECIAL_TOKENS
from common import load_words


def file_sha256_prefix(path: Path, length: int = 12) -> str:
  digest = hashlib.sha256()
  with path.open("rb") as f:
    for chunk in iter(lambda: f.read(1024 * 1024), b""):
      digest.update(chunk)
  return digest.hexdigest()[:length]


def train(words: Sequence[tuple[str, int]], vocab_size: int, unit: str) -> BpeTrainer:
  trainer = BpeTrainer(
    SPECIAL_TOKENS,
    unit=unit,
    initial_alphabet="byte_level" if unit == "byte" else None,
  )
  trainer.add_words(words)
  trainer.train(vocab_size)
  return trainer


def main(argv: Sequence[str] | None = None) -> int:
  parser = argparse.ArgumentParser(description="Train and save a golden tokenizer model from a word-frequency inventory.")
  parser.add_argument("--words", type=Path, required=True, help="JSON word-frequency inventory.")
  parser.add_argument("--dataset-name", required=True, help="Dataset directory name under out/models/golden.")
  parser.add_argument("--vocab-size", type=int, required=True)
  parser.add_argument("--unit", choices=["byte", "unicode"], default="unicode", help="BPE unit used for training.")
  parser.add_argument("--format", choices=["gpt2", "unitoken"], default="unitoken")
  parser.add_argument("--out-dir", type=Path, default=Path("out") / "models" / "golden")
  parser.add_argument("--max-occurrences", type=int, help="Truncate the weighted corpus for a smaller smoke model.")
  args = parser.parse_args(argv)

  if args.vocab_size < 1:
    parser.error("--vocab-size must be at least 1")
  if args.unit == "unicode" and args.format != "unitoken":
    parser.error("--format must be unitoken when --unit=unicode")

  words = load_words(args.words, args.max_occurrences)
  started = time.perf_counter()
  trainer = train(words, args.vocab_size, args.unit)
  train_s = time.perf_counter() - started

  input_hash = file_sha256_prefix(args.words)
  model_dir = args.out_dir / f"{args.dataset_name}.{input_hash}"
  model_dir.mkdir(parents=True, exist_ok=True)
  suffix = f"vocab{args.vocab_size}.{args.format}"
  vocab_path = model_dir / f"vocab.{suffix}.json"
  merges_path = model_dir / f"merges.{suffix}.txt"
  model = trainer.validate_model()
  model.save_files(vocab_path, merges_path, format=args.format)

  print(f"saved {vocab_path}")
  print(f"saved {merges_path}")
  print(f"train_s={train_s:.3f}")
  return 0


if __name__ == "__main__":
  raise SystemExit(main(sys.argv[1:]))