Trainer

ffbpe.trainer.BpeTrainer

BpeTrainer(special_tokens, *, unit='byte', initial_alphabet=None, tie_break=None, parallel_merge_min_occurs_in=None, hot_pair_window_size=None, bigram_cutoff_freq=None)

Train a BPE model from a word-frequency inventory.

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

Parameters:

Name	Type	Description	Default
special_tokens	Sequence[str]	Sequence of tokens reserved in the vocabulary.	required
unit	Unit	Primary segmentation unit. Unicode models may include UTF-8 byte fallback merges.	'byte'
hot_pair_window_size	int | None	If set, retain occurrence postings for an exact top-K candidate window. Smaller values reduce memory but may require additional inventory scans.	None
bigram_cutoff_freq	int | None	Inclusive minimum frequency for pair merges performed by automatic training. Manual step() calls ignore it, but model validation still enforces it.	None

Source code in python/ffbpe/trainer.py

def __init__(
  self,
  special_tokens: Sequence[str],
  *,
  unit: Unit = "byte",
  initial_alphabet: InitialAlphabet | None = None,
  tie_break: TieBreak | None = None,
  parallel_merge_min_occurs_in: int | None = None,
  hot_pair_window_size: int | None = None,
  bigram_cutoff_freq: int | None = None,
) -> None:
  _validate_unit(unit)
  self._unit = unit
  if unit == "unicode":
    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,
      hot_pair_window_size=hot_pair_window_size,
      bigram_cutoff_freq=bigram_cutoff_freq,
    )
  elif unit == "byte":
    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,
      hot_pair_window_size=hot_pair_window_size,
      bigram_cutoff_freq=bigram_cutoff_freq,
    )

vocab_size property

vocab_size

Current vocabulary size.

last_merge_freq property

last_merge_freq

Frequency of the most recently completed pair merge.

hot_pair_window_stats property

hot_pair_window_stats

Diagnostics for the bounded pair-posting window, if enabled.

unit property

unit

Primary segmentation unit used by this trainer.

vocab property

vocab

Return a snapshot of the current token-to-id vocabulary.

add_words

add_words(words)

Add training data.

Accepts either a mapping {word: freq} or an explicit sequence of (word, freq) pairs.

Source code in python/ffbpe/trainer.py

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)

add_word_counter

add_word_counter(counter)

Replace the training inventory by consuming an exact native word counter.

The counter is empty and reusable after this call. Unlike words(), this transfer does not construct a Python dictionary.

Source code in python/ffbpe/trainer.py

def add_word_counter(self, counter: WordCounter) -> None:
  """Replace the training inventory by consuming an exact native word counter.

  The counter is empty and reusable after this call. Unlike `words()`, this
  transfer does not construct a Python dictionary.
  """
  self._trainer.add_word_counter(counter)

init_training

init_training()

Initialize internal training state.

Source code in python/ffbpe/trainer.py

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

train

train(vocab_size)

Train until the vocab reaches vocab_size entries or no eligible pair remains.

Training may finish below the requested size when the inventory is exhausted or the next pair frequency is below bigram_cutoff_freq. A target smaller than the current vocabulary is rejected.

Source code in python/ffbpe/trainer.py

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

  Training may finish below the requested size when the inventory is
  exhausted or the next pair frequency is below `bigram_cutoff_freq`.
  A target smaller than the current vocabulary is rejected.
  """
  if vocab_size < self.vocab_size:
    raise ValueError(
      f"Target vocabulary size {vocab_size} is smaller than "
      f"the current vocabulary size {self.vocab_size}"
    )
  self._trainer.train_until(vocab_size)

train_with_bbpe_fallback

train_with_bbpe_fallback(vocab_size, *, primary_vocab_ratio=0.9)

Train a Unicode model with a terminal byte-BPE fallback phase.

primary_vocab_ratio allocates a fraction of learned slots to the initial Unicode phase. The mandatory 256-byte alphabet and special tokens are excluded; unused fallback slots return to primary training. A fallback pass must start before ordinary vocabulary growth and finalizes the trainer, so create a new trainer for further training. A ratio of 1.0 delegates to ordinary training and remains extendable; a target at or below the base vocabulary is a no-op. The pair-frequency cutoff may leave the final vocabulary below vocab_size.

Source code in python/ffbpe/trainer.py

def train_with_bbpe_fallback(
  self,
  vocab_size: int,
  *,
  primary_vocab_ratio: float = 0.9,
) -> None:
  """Train a Unicode model with a terminal byte-BPE fallback phase.

  `primary_vocab_ratio` allocates a fraction of learned slots to the initial
  Unicode phase. The mandatory 256-byte alphabet and special tokens are
  excluded; unused fallback slots return to primary training. A fallback pass
  must start before ordinary vocabulary growth and finalizes the trainer, so
  create a new trainer for further training. A ratio of `1.0` delegates to
  ordinary training and remains extendable; a target at or below the base
  vocabulary is a no-op. The pair-frequency cutoff may leave the final
  vocabulary below `vocab_size`.
  """
  if self._unit != "unicode":
    raise ValueError('train_with_bbpe_fallback requires unit="unicode"')
  _validate_primary_vocab_ratio(primary_vocab_ratio)
  assert isinstance(self._trainer, BpeTrainer_Character_CharIdx)
  self._trainer.train_until_with_bbpe_fallback(
    vocab_size,
    primary_vocab_ratio=primary_vocab_ratio,
  )

step

step()

Perform one training step.

Returns the updated vocabulary size.

Source code in python/ffbpe/trainer.py

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

  Returns the updated vocabulary size.
  """
  return self._trainer.step()

validate_model

validate_model()

Validate the trainer state and return an immutable model snapshot.

Source code in python/ffbpe/trainer.py

def validate_model(self) -> "BpeModel":
  """Validate the trainer state and return an immutable model snapshot."""
  from .model import BpeModel
  return BpeModel(self._trainer.validate_model())

save

save(name, *, outdir='.', format=None)

Validate a snapshot and save it under name.

Source code in python/ffbpe/trainer.py

def save(
  self,
  name: str,
  *,
  outdir: str | PathLike = ".",
  format: FileFormat | None = None,
) -> None:
  """Validate a snapshot and save it under `name`."""
  vocab_path = Path(outdir) / f"vocab.{name}[{self.unit}].json"
  merges_path = Path(outdir) / f"merges.{name}[{self.unit}].txt"
  self.save_files(
    vocab_path,
    merges_path,
    format=format,
  )

save_files

save_files(vocab_path, merges_path, *, format=None)

Validate a snapshot and save its files.

Source code in python/ffbpe/trainer.py

def save_files(
  self,
  vocab_path: str | PathLike,
  merges_path: str | PathLike,
  *,
  format: FileFormat | None = None,
) -> None:
  """Validate a snapshot and save its files."""
  resolved_format = _resolve_format(self.unit, format)
  self.validate_model().save_files(vocab_path, merges_path, format=resolved_format)