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