unitoken 0.1.4

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
112
113
114
115
116
117
118
119
120
121
122

from collections.abc import Mapping, Sequence
from os import PathLike
from pathlib import Path
from typing import Literal
from ._lib import BpeTrainer_Character_CharIdx, BpeTrainer_u8_Idx

CharLevel = Literal["char", "u8"]
OutputFormat = Literal["uni", "gpt2"]
InitialAlphabet = Literal["raw", "byte_level"]
TieBreak = Literal["smallest_pair_id", "largest_content"]

class BpeTrainer:
  """Train a BPE model from a word-frequency inventory.

  This wraps the Rust trainer classes exposed via the extension module.

  Parameters
  ----------
  special_tokens:
    Sequence of special tokens. The first token is treated as the end-of-text marker.
  ch:
    Character level: `"u8"` (byte-level, GPT-2-style) or `"char"` (Unicode-aware, Uni spec).
  output_format:
    Output spec for serialization (`"gpt2"` or `"uni"`). For `ch="char"`, only `"uni"` is supported.
  """
  def __init__(
    self,
    special_tokens: Sequence[str],
    *,
    ch: CharLevel = "u8",
    output_format: OutputFormat | None = None,
    initial_alphabet: InitialAlphabet | None = None,
    tie_break: TieBreak | None = None,
    parallel_merge_min_occurs_in: int | None = None,
  ) -> None:
    # super().__init__()
    self._ch: CharLevel = ch
    self.output_format: OutputFormat = "uni"
    if ch == "char":
      self._trainer = BpeTrainer_Character_CharIdx(
        special_tokens=special_tokens,
        initial_alphabet=initial_alphabet,
        tie_break=tie_break,
        parallel_merge_min_occurs_in=parallel_merge_min_occurs_in,
      )
    else:
      self.output_format = output_format or "gpt2"
      self._trainer = BpeTrainer_u8_Idx(
        special_tokens=special_tokens,
        initial_alphabet=initial_alphabet,
        tie_break=tie_break,
        parallel_merge_min_occurs_in=parallel_merge_min_occurs_in,
      )

  @property
  def vocab_size(self) -> int:
    """Current vocabulary size."""
    return self._trainer.vocab_size()

  @property
  def char_level(self) -> CharLevel:
    """Character level used by this trainer."""
    return self._ch

  @property
  def vocabs(self):
    """Vocabulary view.

    Returns a `Vocabs` object (from the extension module) supporting `len`, `get`, and `items`.
    """
    return self._trainer.get_vocabs()

  def add_words(self, words: Mapping[str, int] | Sequence[tuple[str, int]]) -> None:
    """Add training data.

    Accepts either a mapping `{word: freq}` or an explicit sequence of `(word, freq)` pairs.
    """
    if isinstance(words, Mapping):
      words = list(words.items())
    self._trainer.add_words(words)

  def init_training(self) -> None:
    """Initialize internal training state."""
    self._trainer.init_training()

  def train(self, vocab_size: int) -> None:
    """Train until the vocab reaches `vocab_size` entries.

    Training runs inside Rust until the target is reached.
    """
    try:
      self._trainer.train_until(vocab_size)
    except:
      return

  def step(self) -> int | None:
    """Perform one training step.

    Returns the updated vocabulary size, or `None` if no further merges can be made.
    """
    try:
      return self._trainer.step() or self._trainer.vocab_size()
    except:
      return None

  def save(self, name: str, *, outdir: str | PathLike = ".", output_format: OutputFormat | None = None) -> None:
    """Save `vocab.{name}[{ch}].json` and `merges.{name}[{ch}].txt` into `outdir`."""
    vocab_path = Path(outdir) / f"vocab.{name}[{self.char_level}].json"
    merges_path = Path(outdir) / f"merges.{name}[{self.char_level}].txt"
    self.save_files(vocab_path, merges_path, output_format=output_format)

  def save_files(
    self,
    vocab_path: str | PathLike,
    merges_path: str | PathLike,
    *,
    output_format: OutputFormat | None = None,
  ) -> None:
    """Save vocab and merges to explicit paths."""
    spec = output_format or self.output_format
    self._trainer.save_vocab(vocab_path, spec)
    self._trainer.save_merges_txt(merges_path, spec)