File size: 1,805 Bytes
4d848c2 | 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 | """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))
|