trimbed 0.0.1__py3-none-any.whl

trimbed/__init__.py

@@ -0,0 +1,75 @@
+"""Trim a tokenizer's vocabulary, and optionally its model, down to what you need.
+
+Typical use:
+
+    >>> from trimbed import TrimConfig, TrimPipeline
+    >>> config = TrimConfig.from_yaml("config.yaml")
+    >>> report = TrimPipeline(config).run()
+"""
+
+from __future__ import annotations
+
+from importlib.metadata import version
+
+from trimbed.backends import VocabBackend, get_backend, register_backend, supported_model_types
+from trimbed.config import (
+    CorpusConfig,
+    DatasetSpec,
+    EmbeddingTrimConfig,
+    SelectionConfig,
+    TrimConfig,
+    load_config,
+    parse_overrides,
+)
+from trimbed.counting import CorpusCounter, CorpusCounts
+from trimbed.exceptions import MissingDependencyError
+from trimbed.model_trim import trim_model
+from trimbed.pipeline import TrimPipeline
+from trimbed.presets import available_presets, register_preset, resolve_preset
+from trimbed.remap import IdRemap
+from trimbed.report import ModelReport, ModelVerificationReport, TrimReport, VerificationReport
+from trimbed.selection import Selection, select_tokens
+from trimbed.sidecar import DEFAULT_SIDECAR_PATTERNS, copy_sidecar_files
+from trimbed.spec import TokenizerSpec
+from trimbed.tokenizer_trim import TrimmedTokenizer, trim_tokenizer
+from trimbed.verify import verify_model, verify_tokenizer
+
+
+__version__ = version("trimbed")
+
+__all__ = [
+    "DEFAULT_SIDECAR_PATTERNS",
+    "CorpusConfig",
+    "CorpusCounter",
+    "CorpusCounts",
+    "DatasetSpec",
+    "EmbeddingTrimConfig",
+    "IdRemap",
+    "MissingDependencyError",
+    "ModelReport",
+    "ModelVerificationReport",
+    "Selection",
+    "SelectionConfig",
+    "TokenizerSpec",
+    "TrimConfig",
+    "TrimPipeline",
+    "TrimReport",
+    "TrimmedTokenizer",
+    "VerificationReport",
+    "VocabBackend",
+    "__version__",
+    "available_presets",
+    "copy_sidecar_files",
+    "get_backend",
+    "load_config",
+    "parse_overrides",
+    "register_backend",
+    "register_preset",
+    "resolve_preset",
+    "select_tokens",
+    "supported_model_types",
+    "trim_model",
+    "trim_tokenizer",
+    "verify_model",
+    "verify_tokenizer",
+]

trimbed/_logging.py

@@ -0,0 +1,45 @@
+"""Logging helpers shared by the library and the scripts.
+
+`configure_logging()` should be called once at the start of a script to set up the package logger,
+e.g. to set it to quiet mode.
+"""
+
+from __future__ import annotations
+
+import logging
+
+
+LOGGER_NAME = "trimbed"
+
+
+def get_logger(name: str | None = None) -> logging.Logger:
+    """Return the package logger, or a child of it.
+
+    Args:
+        name: Optional dotted suffix, typically `__name__`.
+
+    Returns:
+        The `trimbed` logger or the named child logger.
+    """
+    if name is None or name == LOGGER_NAME:
+        return logging.getLogger(LOGGER_NAME)
+    suffix = name.removeprefix(f"{LOGGER_NAME}.")
+    return logging.getLogger(f"{LOGGER_NAME}.{suffix}")
+
+
+def configure_logging(verbose: bool = False, quiet: bool = False) -> None:
+    """Attach a single stderr handler to the package logger.
+
+    Repeated calls replace the previous handler rather than stacking.
+
+    Args:
+        verbose: Emit DEBUG-level records.
+        quiet: Emit only warnings and above. Ignored when `verbose` is set.
+    """
+    logger = logging.getLogger(LOGGER_NAME)
+    logger.handlers.clear()
+    handler = logging.StreamHandler()
+    handler.setFormatter(logging.Formatter("%(levelname)-8s %(name)s: %(message)s"))
+    logger.addHandler(handler)
+    logger.setLevel(logging.DEBUG if verbose else logging.WARNING if quiet else logging.INFO)
+    logger.propagate = False

trimbed/backends/__init__.py

@@ -0,0 +1,84 @@
+"""Registry of per-tokenizer-family trimming adapters.
+
+Adding support for a new tokenizer family means adding one module here and decorating
+its class with `register_backend`, following one of the existing backends. Nothing else
+in the package needs to change.
+"""
+
+from __future__ import annotations
+
+from trimbed.backends.base import VocabBackend
+
+
+_REGISTRY: dict[str, VocabBackend] = {}
+
+
+def register_backend[BackendT: type[VocabBackend]](cls: BackendT) -> BackendT:
+    """Register a backend adapter under its declared `model_type`.
+
+    Args:
+        cls: A concrete [`VocabBackend`][trimbed.backends.base.VocabBackend] subclass.
+
+    Returns:
+        The class unchanged, so this can be used as a decorator.
+
+    Raises:
+        ValueError: If the class has no `model_type` or the type is already taken.
+    """
+    model_type = getattr(cls, "model_type", None)
+    if not model_type:
+        raise ValueError(f"{cls.__name__} must declare a `model_type` class attribute")
+    if model_type in _REGISTRY:
+        raise ValueError(f"a backend for model type {model_type!r} is already registered")
+    _REGISTRY[model_type] = cls()
+    return cls
+
+
+def get_backend(model_type: str) -> VocabBackend:
+    """Look up the adapter for a `tokenizer.json` model type.
+
+    Args:
+        model_type: Value of `model.type`, e.g. `"BPE"` for codefuse-ai/F2LLM-v2-160M
+            or `"WordPiece"` for google-bert/bert-base-cased.
+
+    Returns:
+        The registered adapter instance, e.g. a [`BpeBackend`][trimbed.backends.bpe.BpeBackend].
+
+    Raises:
+        KeyError: If no adapter handles that model type.
+    """
+    try:
+        return _REGISTRY[model_type]
+    except KeyError:
+        supported = ", ".join(sorted(_REGISTRY)) or "none"
+        raise KeyError(
+            f"no trimming backend registered for tokenizer model type {model_type!r} (supported: {supported})"
+        ) from None
+
+
+def supported_model_types() -> tuple[str, ...]:
+    """Return the tokenizer model types that can currently be trimmed.
+
+    E.g. `("BPE", "Unigram", "WordLevel", "WordPiece")`.
+    """
+    return tuple(sorted(_REGISTRY))
+
+
+# Importing the adapters is what populates the registry, since that is when their
+# `register_backend` decorators run.
+from trimbed.backends.bpe import BpeBackend
+from trimbed.backends.unigram import UnigramBackend
+from trimbed.backends.wordlevel import WordLevelBackend
+from trimbed.backends.wordpiece import WordPieceBackend
+
+
+__all__ = [
+    "BpeBackend",
+    "UnigramBackend",
+    "VocabBackend",
+    "WordLevelBackend",
+    "WordPieceBackend",
+    "get_backend",
+    "register_backend",
+    "supported_model_types",
+]

trimbed/backends/base.py

@@ -0,0 +1,85 @@
+"""The base interface for different tokenizer families.
+
+We luckily can rely on `skeletoken` to do the vocabulary surgery for every tokenizer
+family, so these adapters only carry what a trimmer has to know and a serialiser does
+not: which tokens can never be dropped, and which tokens depend on which.
+"""
+
+from __future__ import annotations
+
+from abc import ABC
+from typing import TYPE_CHECKING, ClassVar
+
+
+if TYPE_CHECKING:
+    from trimbed.spec import TokenizerSpec
+
+
+class VocabBackend(ABC):
+    """Describes the structural constraints of one `tokenizers` model type.
+
+    Subclasses answer which tokens can never be removed without breaking encoding
+    outright and which tokens can only be produced if some other token survives.
+    """
+
+    model_type: ClassVar[str]
+    """Value of `model.type` in tokenizer.json that this adapter handles.
+
+    E.g. `"BPE"` for codefuse-ai/F2LLM-v2-160M, `"WordPiece"` for
+    google-bert/bert-base-cased and `"Unigram"` for google-t5/t5-small.
+    """
+
+    def structural_tokens(self, spec: TokenizerSpec) -> set[str]:
+        """Return tokens that must survive or the tokenizer stops working.
+
+        The default covers the unknown token. Byte-level backends extend it with the byte
+        alphabet.
+
+        Args:
+            spec: The tokenizer being trimmed.
+
+        Returns:
+            Token strings that may never be dropped. E.g. `{"[UNK]"}` for
+            google-bert/bert-base-cased, and an empty set for codefuse-ai/F2LLM-v2-160M,
+            which declares no unknown token at all.
+        """
+        unk = spec.unk_token
+        return {unk} if unk else set()
+
+    def dependencies(self, spec: TokenizerSpec) -> dict[int, tuple[int, ...]]:
+        """Return which other tokens each token needs in order to stay reachable.
+
+        Keeping a token while dropping something it is assembled from leaves it in the
+        vocabulary but unreachable, so text silently fragments.
+        Backends with no such structure return an empty mapping.
+
+        Args:
+            spec: The tokenizer being trimmed.
+
+        Returns:
+            A mapping of token id to the ids it directly depends on. Empty here, since
+            BPE is the only family that builds tokens out of other tokens. See
+            [`BpeBackend.dependencies`][trimbed.backends.bpe.BpeBackend.dependencies] for
+            how it fills this in.
+        """
+        return {}
+
+    def _byte_fallback_tokens(self, spec: TokenizerSpec) -> set[str]:
+        """Return the `<0xNN>` tokens present in the vocabulary, if byte fallback is on.
+
+        Args:
+            spec: The tokenizer being trimmed.
+
+        Returns:
+            The byte-fallback tokens, e.g. `{"<0x00>", "<0x01>", ..., "<0xFF>"}`, or an
+            empty set when the feature is off, as it is for google-t5/t5-small and
+            FacebookAI/xlm-roberta-base. A byte-level BPE never needs it: its alphabet
+            already covers all 256 bytes.
+        """
+        if not getattr(spec.model, "byte_fallback", False):
+            return set()
+        return {f"<0x{byte:02X}>" for byte in range(256)} & spec.vocabulary.keys()
+
+    def __repr__(self) -> str:
+        """Return a short debug representation."""
+        return f"<{type(self).__name__} model_type={self.model_type!r}>"

trimbed/backends/bpe.py

@@ -0,0 +1,63 @@
+"""Selection constraints for byte-pair-encoding tokenizers (GPT-2, Qwen, Llama, ...)."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from trimbed.backends import register_backend
+from trimbed.backends.base import VocabBackend
+from trimbed.bytelevel import byte_level_alphabet
+
+
+if TYPE_CHECKING:
+    from trimbed.spec import TokenizerSpec
+
+
+@register_backend
+class BpeBackend(VocabBackend):
+    """Keeps a BPE vocabulary encodable and its merge chains intact."""
+
+    model_type = "BPE"
+
+    def structural_tokens(self, spec: TokenizerSpec) -> set[str]:
+        """Return the unk token plus, for byte-level BPE, the whole byte alphabet.
+
+        Args:
+            spec: The tokenizer being trimmed.
+
+        Returns:
+            Token strings that may never be dropped. For codefuse-ai/F2LLM-v2-160M that is
+            the 256 alphabet characters (`"!"`, `"Ġ"` for a space, `"Ċ"` for a
+            newline, ...) and nothing else, since it has no unknown token.
+        """
+        required = super().structural_tokens(spec)
+        if spec.uses_byte_level:
+            # add the byte-level alphabet, but only the bytes that are actually in the vocabulary
+            required |= set(byte_level_alphabet()) & spec.vocabulary.keys()
+        return required | self._byte_fallback_tokens(spec)
+
+    def dependencies(self, spec: TokenizerSpec) -> dict[int, tuple[int, ...]]:
+        """Map each merged token to the pair it is assembled from.
+
+        codefuse-ai/F2LLM-v2-160M builds `"Ġthe"` by merging `"Ġth"` with `"e"`, so
+        dropping `"Ġth"` leaves `"Ġthe"` in the vocabulary but unreachable and the text
+        quietly tokenizes into characters instead. The selector uses this mapping to pull
+        in whatever a kept token is built from.
+
+        Args:
+            spec: The tokenizer being trimmed.
+
+        Returns:
+            A mapping of merged token id to its two parent ids, e.g. `Ġthe -> (Ġth, e)`
+            and `Ġworld -> (Ġw, orld)`. That is 151,387 entries for
+            codefuse-ai/F2LLM-v2-160M, one per merge rule whose three tokens all survive.
+        """
+        vocabulary = spec.vocabulary
+        dependencies: dict[int, tuple[int, ...]] = {}
+        for left, right in spec.model.merges.root:
+            # concatenate the strings and check if the merged token is actually in vocab
+            merged = left + right
+            if merged in vocabulary and left in vocabulary and right in vocabulary:
+                # The first rule producing a token is the one BPE actually applies.
+                dependencies.setdefault(vocabulary[merged], (vocabulary[left], vocabulary[right]))
+        return dependencies

trimbed/backends/unigram.py

@@ -0,0 +1,35 @@
+"""Selection constraints for Unigram tokenizers (XLM-R, mT5, ALBERT, ...)."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from trimbed.backends import register_backend
+from trimbed.backends.base import VocabBackend
+
+
+if TYPE_CHECKING:
+    from trimbed.spec import TokenizerSpec
+
+
+@register_backend
+class UnigramBackend(VocabBackend):
+    """Keeps a Unigram vocabulary usable.
+
+    Unigram scores whole candidate pieces independently rather than composing them from
+    merges.
+    """
+
+    model_type = "Unigram"
+
+    def structural_tokens(self, spec: TokenizerSpec) -> set[str]:
+        """Return the unk token plus any byte-fallback tokens the model relies on.
+
+        Args:
+            spec: The tokenizer being trimmed.
+
+        Returns:
+            Token strings that may never be dropped, e.g. `{"<unk>"}` for
+            google-t5/t5-small and FacebookAI/xlm-roberta-base.
+        """
+        return super().structural_tokens(spec) | self._byte_fallback_tokens(spec)

trimbed/backends/wordlevel.py

@@ -0,0 +1,17 @@
+"""Selection constraints for WordLevel tokenizers."""
+
+from __future__ import annotations
+
+from trimbed.backends import register_backend
+from trimbed.backends.base import VocabBackend
+
+
+@register_backend
+class WordLevelBackend(VocabBackend):
+    """Keeps a flat word-level vocabulary usable.
+
+    The simplest case: no merges, no subwords, so no dependencies. Anything dropped
+    becomes unk, which makes the unk token the only structural requirement.
+    """
+
+    model_type = "WordLevel"

trimbed/backends/wordpiece.py

@@ -0,0 +1,18 @@
+"""Selection constraints for WordPiece tokenizers (BERT and friends)."""
+
+from __future__ import annotations
+
+from trimbed.backends import register_backend
+from trimbed.backends.base import VocabBackend
+
+
+@register_backend
+class WordPieceBackend(VocabBackend):
+    """Keeps a WordPiece vocabulary usable.
+
+    WordPiece greedily matches the longest prefix in the vocabulary and needs no merge
+    table, so tokens carry no dependencies: dropping `"##ing"` costs coverage but
+    never makes another vocabulary token unreachable.
+    """
+
+    model_type = "WordPiece"

trimbed/bytelevel.py

@@ -0,0 +1,71 @@
+"""The GPT-2-style byte-to-unicode mapping used by ByteLevel tokenizers."""
+
+from __future__ import annotations
+
+from functools import cache
+
+
+@cache
+def bytes_to_unicode() -> dict[int, str]:
+    """Return the reversible byte -> printable-character map used by ByteLevel.
+
+    Borrowed mostly from GPT-2's tokenizer, also see the [transformers implementation](https://github.com/huggingface/transformers/blob/36bc98ef9dd009569366f5e253ec1876ecafd925/src/transformers/convert_slow_tokenizer.py#L1879).
+
+    Returns:
+        A mapping of all 256 byte values to distinct printable characters, e.g.
+        `32 -> "Ġ"` (space), `10 -> "Ċ"` (newline), `0 -> "Ā"`, and `65 -> "A"`
+        for the bytes that are already printable.
+    """
+    printable = (
+        list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
+    )
+    mapped = list(printable)
+    spare = 0
+    for byte in range(256):
+        if byte not in printable:
+            printable.append(byte)
+            mapped.append(256 + spare)
+            spare += 1
+    return {byte: chr(code) for byte, code in zip(printable, mapped, strict=True)}
+
+
+@cache
+def unicode_to_bytes() -> dict[str, int]:
+    """Return the inverse of [`bytes_to_unicode`][trimbed.bytelevel.bytes_to_unicode].
+
+    E.g. `"Ġ" -> 32` and `"Ċ" -> 10`.
+    """
+    return {char: byte for byte, char in bytes_to_unicode().items()}
+
+
+@cache
+def byte_level_alphabet() -> frozenset[str]:
+    """Return the 256 characters a ByteLevel pre-tokenizer can produce.
+
+    Every one of them must stay in the vocabulary of a byte-level tokenizer, otherwise
+    some byte sequences become unencodable.
+    """
+    return frozenset(bytes_to_unicode().values())
+
+
+def decode_byte_level(token: str) -> str | None:
+    """Turn a byte-level token back into the text it represents.
+
+    Args:
+        token: A token as stored in a ByteLevel vocabulary, e.g. `"Ġde"`, which will
+            produce `" de"`. Multi-byte characters arrive as several alphabet
+            characters, so `"é"` produces `"é"`.
+
+    Returns:
+        The decoded text, or `None` if the token is a partial UTF-8 sequence or
+        contains characters outside the byte-level alphabet.
+    """
+    table = unicode_to_bytes()
+    try:
+        raw = bytes(table[char] for char in token)
+    except KeyError:
+        return None
+    try:
+        return raw.decode("utf-8")
+    except UnicodeDecodeError:
+        return None

trimbed/cli/__init__.py

@@ -0,0 +1,7 @@
+"""The command surface: one module per `trimbed` subcommand, plus the router.
+
+The router is [`trimbed.cli.__main__`][trimbed.cli.__main__].
+Each module holds a `run(...)` with the logic and an `add_arguments(parser)` describing
+its command line, so a command is equally usable from the shell and as a plain Python
+call.
+"""

trimbed/cli/__main__.py

@@ -0,0 +1,51 @@
+"""The main `trimbed` command.
+
+    trimbed inspect --model codefuse-ai/F2LLM-v2-160M
+    trimbed trim --config my_config.yaml --dry-run selection.top_k=30000
+    trimbed presets
+
+Each subcommand owns its own arguments, in `add_arguments`, next to the `run` they feed.
+This module holds the table of subcommand names and nothing else, so adding a command is
+a new module plus one line here.
+"""
+
+from __future__ import annotations
+
+import argparse
+from types import ModuleType
+
+from trimbed.cli import count_tokens, inspect_tokenizer, list_presets, trim_vocab
+
+
+# subcommand name -> the module implementing it and the one line `trimbed --help` shows.
+COMMANDS: dict[str, tuple[ModuleType, str]] = {
+    "trim": (trim_vocab, "Trim a tokenizer, and optionally its model, and write the result."),
+    "count": (count_tokens, "Count the corpus once and cache the frequencies to JSON."),
+    "inspect": (inspect_tokenizer, "Describe a tokenizer as JSON, changing nothing."),
+    "presets": (list_presets, "List the registered presets that --keep-preset accepts."),
+}
+
+
+def main() -> None:
+    """Parse the command line and hand the arguments to the chosen subcommand's `run`."""
+    parser = argparse.ArgumentParser(
+        prog="trimbed",
+        description=(
+            "Trim a tokenizer's vocabulary, and optionally its model's embedding table, down"
+            " to the subset a corpus and your must-keep rules actually need."
+        ),
+    )
+    subparsers = parser.add_subparsers(dest="command", required=True, metavar="COMMAND")
+    for name, (module, summary) in COMMANDS.items():
+        subparser = subparsers.add_parser(name, help=summary, description=module.DESCRIPTION)
+        module.add_arguments(subparser)
+        subparser.set_defaults(run=module.run)
+
+    arguments = vars(parser.parse_args())
+    command = arguments.pop("run")
+    del arguments["command"]
+    command(**arguments)
+
+
+if __name__ == "__main__":
+    main()

trimbed/cli/count_tokens.py

@@ -0,0 +1,85 @@
+"""Count corpus token frequencies once and cache them to JSON.
+
+Counting is the expensive part of a trimming run and it does not depend on the selection
+policy, so do it once here and point `corpus.counts_cache` at the output to reuse it
+across selection experiments.
+
+    trimbed count --config my_config.yaml -o counts.json
+    trimbed count --config my_config.yaml -o counts.json \
+        corpus.batch_size=4000 corpus.num_proc=8
+
+Only the corpus half of the config matters here. The selection and embedding sections are
+still validated, but they are not used.
+"""
+
+from __future__ import annotations
+
+import argparse
+
+from trimbed._logging import configure_logging
+from trimbed.config import load_config, parse_overrides
+from trimbed.counting import CorpusCounter
+from trimbed.loading import load_tokenizer
+
+
+def run(
+    output: str,
+    config: str | None = None,
+    model: str | None = None,
+    overrides: list[str] | None = None,
+    verbose: bool = False,
+    quiet: bool = False,
+) -> None:
+    """Count the configured corpus and write the frequencies to a cache file.
+
+    Args:
+        output: JSON file to write the counts to.
+        config: Path to a YAML configuration file.
+        model: Hub model id or local path, overriding the config.
+        overrides: `key=value` strings applied on top of the config and flags.
+        verbose: Emit debug logging.
+        quiet: Only emit warnings and errors.
+
+    Raises:
+        ValueError: If the resolved config names no datasets to count over.
+    """
+    configure_logging(verbose=verbose, quiet=quiet)
+
+    trim_config = load_config(config, model).with_overrides({"model": model, **parse_overrides(overrides or [])})
+    if not trim_config.corpus.datasets:
+        raise ValueError("no datasets configured under 'corpus.datasets'; nothing to count")
+
+    tokenizer = load_tokenizer(trim_config.model, trim_config.revision, trim_config.trust_remote_code)
+    counter = CorpusCounter(
+        tokenizer, trim_config.corpus, seed=trim_config.seed, sample_size=trim_config.verify_samples
+    )
+    counts = counter.count()
+    counts.save(output)
+    print(
+        f"counted {counts.total_num_tokens:,} tokens ({counts.distinct_tokens:,} distinct) "
+        f"over {counts.num_documents:,} documents -> {output}"
+    )
+
+
+DESCRIPTION = (
+    "Tokenize the configured corpus and write the per-token frequencies to JSON, so"
+    " later trimming runs can reuse them instead of re-reading the corpus."
+)
+
+
+def add_arguments(parser: argparse.ArgumentParser) -> None:
+    """Add the counting arguments to `parser`.
+
+    Args:
+        parser: The `trimbed count` subparser to populate.
+    """
+    parser.add_argument("-c", "--config", help="Path to a YAML configuration file.")
+    parser.add_argument("-m", "--model", help="Hub model id or local path (overrides the config).")
+    parser.add_argument("-o", "--output", required=True, help="JSON file to write the counts to.")
+    parser.add_argument("-v", "--verbose", action="store_true", help="Emit debug logging.")
+    parser.add_argument("-q", "--quiet", action="store_true", help="Only emit warnings and errors.")
+    parser.add_argument(
+        "overrides",
+        nargs="*",
+        help="Optional key=value overrides, e.g. corpus.batch_size=4000",
+    )