File size: 2,448 Bytes
3020d73 | 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 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | """TinyBrainBot units-first number tokenizer (byte-level BPE + length-prefixed, reversed digit runs).
Every run of 1-8 digits is written as a length token and then its digits units first:
"62" -> "<len2>26" "1234" -> "<len4>4321"
so a model emits the units digit first and carries flow forward in time. Digit runs touching ":" (clock
times like 3:15) and runs longer than 8 digits are left unchanged. The text is rewritten on encode and
restored on decode, so callers see ordinary numbers. The rewrite is idempotent (an already-rewritten run
is not rewritten again), so it is safe on every encode path.
The idea (units-first digit tokens; a length prefix) follows Atom2.7m and Pascaline-3M. This is our own
implementation.
"""
import re
from transformers import PreTrainedTokenizerFast
MAX_LEN = 8
# a digit run: not part of "<lenN>", not right after a "<lenN>" tag, not touching another digit or ":"
_RUN = re.compile(r"(?<!<len)(?<!<len\d>)(?<![\d:])(\d{1,%d})(?![\d:])" % MAX_LEN)
_TAG = re.compile(r"<len([1-%d])>(\d+)" % MAX_LEN)
def encode_numbers(text):
if not isinstance(text, str):
return text
return _RUN.sub(lambda m: "<len%d>%s" % (len(m.group(1)), m.group(1)[::-1]), text)
def decode_numbers(text):
def back(m):
n = int(m.group(1)); d = m.group(2)
return d[:n][::-1] + d[n:]
return _TAG.sub(back, text)
def _map(x):
if isinstance(x, str):
return encode_numbers(x)
if isinstance(x, tuple):
return tuple(_map(v) for v in x)
if isinstance(x, list):
return [_map(v) for v in x]
return x
class TBBUnitsFirstTokenizerFast(PreTrainedTokenizerFast):
def _batch_encode_plus(self, batch_text_or_text_pairs, *args, **kwargs):
if not kwargs.get("is_split_into_words", False):
batch_text_or_text_pairs = _map(batch_text_or_text_pairs)
return super()._batch_encode_plus(batch_text_or_text_pairs, *args, **kwargs)
def _encode_plus(self, text, text_pair=None, *args, **kwargs):
if not kwargs.get("is_split_into_words", False):
text, text_pair = _map(text), _map(text_pair)
return super()._encode_plus(text, text_pair, *args, **kwargs)
def tokenize(self, text, *args, **kwargs):
return super().tokenize(encode_numbers(text), *args, **kwargs)
def _decode(self, token_ids, *args, **kwargs):
return decode_numbers(super()._decode(token_ids, *args, **kwargs))
|