tokenization_nexus.py

"""Nexus-Erebus digit-atomic, least-significant-digit-first tokenizer.

Digits are never merged by BPE, and every maximal run of digits is reversed at
encode time so the model reads and writes numbers least-significant-digit first.
This aligns carry propagation with the left-to-right direction the model reads,
which is what lets a tiny model do integer arithmetic.

The transform is an involution, so decoding simply applies it again to restore
ordinary left-to-right numbers. Callers see normal text in and normal text out.

Load with:
    AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
"""
import re
from transformers import PreTrainedTokenizerFast

_DIGIT_RUN = re.compile(r"\d+")


def rev_digits(text: str) -> str:
    """Reverse each maximal run of digits. Involution: rev(rev(x)) == x."""
    return _DIGIT_RUN.sub(lambda m: m.group(0)[::-1], text)


class NexusLSDTokenizer(PreTrainedTokenizerFast):
    """PreTrainedTokenizerFast that applies the LSD-first digit transform."""

    def _t(self, x):
        if x is None:
            return None
        if isinstance(x, str):
            return rev_digits(x)
        if isinstance(x, (list, tuple)):
            return type(x)(self._t(i) for i in x)
        return x

    def _batch_encode_plus(self, batch_text_or_text_pairs, *args, **kwargs):
        return super()._batch_encode_plus(self._t(batch_text_or_text_pairs), *args, **kwargs)

    def _encode_plus(self, text, text_pair=None, *args, **kwargs):
        return super()._encode_plus(self._t(text), self._t(text_pair), *args, **kwargs)

    # NOTE: do not override tokenize(); it routes through _encode_plus, so the
    # transform would be applied twice and cancel out.

    def _decode(self, *args, **kwargs):
        return rev_digits(super()._decode(*args, **kwargs))