unitoken 0.1.7

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

if TYPE_CHECKING:
  from .model import BpeModel

Unit = Literal["byte", "unicode"]
FileFormat = Literal["unitoken", "gpt2"]
InitialAlphabet = Literal["raw", "byte_level"]
TieBreak = Literal["smallest_pair_id", "largest_content"]

def _validate_unit(unit: str) -> None:
  if unit not in ("byte", "unicode"):
    raise ValueError(f"Unknown unit: {unit}")

def _validate_primary_vocab_ratio(primary_vocab_ratio: float) -> None:
  if not math.isfinite(primary_vocab_ratio) or not 0.0 <= primary_vocab_ratio <= 1.0:
    raise ValueError("primary_vocab_ratio must be finite and between 0 and 1")

def _resolve_format(unit: Unit, format: FileFormat | None) -> FileFormat:
  _validate_unit(unit)
  resolved_format = format or ("unitoken" if unit == "unicode" else "gpt2")
  if resolved_format not in ("gpt2", "unitoken"):
    raise ValueError(f"Unknown format: {resolved_format}")
  if unit == "unicode" and resolved_format == "gpt2":
    raise ValueError('format="gpt2" is not compatible with unit="unicode"')
  return resolved_format

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 tokens reserved in the vocabulary.
  unit:
    Primary segmentation unit. Unicode models may include UTF-8 byte fallback merges.
  hot_pair_window_size:
    If set, retain occurrence postings for an exact top-K candidate window.
    Smaller values reduce memory but may require additional inventory scans.
  bigram_cutoff_freq:
    Inclusive minimum frequency for pair merges performed by automatic training.
    Manual `step()` calls ignore it, but model validation still enforces it.
  """
  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,
      )

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

  @property
  def last_merge_freq(self) -> int | None:
    """Frequency of the most recently completed pair merge."""
    return self._trainer.last_merge_freq

  @property
  def hot_pair_window_stats(self) -> dict[str, int] | None:
    """Diagnostics for the bounded pair-posting window, if enabled."""
    stats = self._trainer.hot_pair_window_stats
    return None if stats is None else dict(stats)

  @property
  def unit(self) -> Unit:
    """Primary segmentation unit used by this trainer."""
    return self._unit

  @property
  def vocab(self) -> dict[bytes, int]:
    """Return a snapshot of the current token-to-id vocabulary."""
    return dict(self._trainer.get_vocab().items())

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

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

  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,
    )

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

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

  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())

  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,
    )

  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)