tokenization_andrea.py

# This is free software for the public good of a permacomputer hosted at
# permacomputer.com, an always-on computer by the people, for the people.
"""Harris successor-variety tokenizer for HuggingFace transformers.

A port of `HarrisTokenizer` in `microgpt/microgpt.py`. Tokenization is
byte-level with a learned morpheme table:

    0 .. 255            raw UTF-8 byte fallback
    256 .. 256+M-1      morpheme segments, in the order they were learned
    256+M               BOS

Encoding splits on a fixed separator set, then takes the longest matching
segment at each position inside a word, falling back to single bytes. Because
two distinct segments of the same length can never match the same span,
"first match in length-sorted order" and "longest match" are the same rule —
so this scans by descending length instead of walking the whole table.

Token strings use the GPT-2 byte-to-unicode mapping, so every byte gets a
printable, reversible character and the vocab stays a plain str -> int dict
even for segments that are not valid UTF-8 on their own.

Round-trip equivalence with `HarrisTokenizer` is asserted over real corpus
text in `microgpt/tests/test_hf_tokenizer.py`.
"""
import json
import os
from functools import lru_cache

from transformers.tokenization_utils import PreTrainedTokenizer

VOCAB_FILES_NAMES = {'vocab_file': 'harris_segments.json'}

# Must stay identical to _HARRIS_SEPARATORS in microgpt/microgpt.py.
SEPARATORS = frozenset(b' \t\n\r.,;:!?()[]{}"\'/\\@#$%^&*+-=<>|~`')

# A one-byte segment would otherwise share a surface form with its byte-fallback
# token — two distinct ids, one string, and a vocab dict that silently loses
# one of them. Segment tokens carry this marker so the mapping stays injective.
# bytes_to_unicode never emits a codepoint above U+0143, so U+2581 cannot clash.
SEGMENT_PREFIX = '▁'


@lru_cache()
def bytes_to_unicode():
    """GPT-2's reversible byte -> printable-codepoint map."""
    bs = (list(range(ord('!'), ord('~') + 1))
          + list(range(ord('\xa1'), ord('\xac') + 1))
          + list(range(ord('\xae'), ord('\xff') + 1)))
    cs = bs[:]
    n = 0
    for b in range(256):
        if b not in bs:
            bs.append(b)
            cs.append(256 + n)
            n += 1
    return dict(zip(bs, (chr(c) for c in cs)))


class AndreaHarrisTokenizer(PreTrainedTokenizer):
    vocab_files_names = VOCAB_FILES_NAMES
    model_input_names = ['input_ids', 'attention_mask']

    def __init__(self, vocab_file, bos_token='<|bos|>', eos_token='<|bos|>',
                 unk_token=None, pad_token=None, **kwargs):
        with open(vocab_file, encoding='utf-8') as f:
            raw = json.load(f)
        # Accept list-of-strings or list-of-int-lists — both ship in the wild.
        self.segments = [s.encode('utf-8') if isinstance(s, str) else bytes(s)
                         for s in raw]
        self.vocab_file = vocab_file

        self._byte_encoder = bytes_to_unicode()
        self._byte_decoder = {v: k for k, v in self._byte_encoder.items()}

        self._seg_to_id = {seg: 256 + i for i, seg in enumerate(self.segments)}
        self._id_to_seg = {i: bytes([i]) for i in range(256)}
        self._id_to_seg.update({256 + i: seg for i, seg in enumerate(self.segments)})
        self._bos_id = 256 + len(self.segments)
        self._id_to_seg[self._bos_id] = b''
        self._max_seg_len = max((len(s) for s in self.segments), default=1)
        # v2 space-annexed mode: auto-detected — any segment containing a
        # space byte. v1 tables tokenize identically to the pre-v2 path.
        self._v2 = any(0x20 in s for s in self.segments)

        self._id_to_token = {i: self._bytes_to_token(bytes([i])) for i in range(256)}
        self._id_to_token.update({
            256 + i: SEGMENT_PREFIX + self._bytes_to_token(seg)
            for i, seg in enumerate(self.segments)})
        self._id_to_token[self._bos_id] = str(bos_token)
        self._token_to_id = {v: k for k, v in self._id_to_token.items()}
        if len(self._token_to_id) != len(self._id_to_token):
            raise ValueError('segment table produced duplicate token strings; '
                             'vocabulary would not round-trip')

        super().__init__(bos_token=bos_token, eos_token=eos_token,
                         unk_token=unk_token, pad_token=pad_token, **kwargs)

    # ── vocab ────────────────────────────────────────────────────────────
    @property
    def vocab_size(self):
        return 256 + len(self.segments) + 1

    def get_vocab(self):
        return dict(self._token_to_id)

    def _bytes_to_token(self, raw):
        return ''.join(self._byte_encoder[b] for b in raw)

    def _token_to_bytes(self, token):
        if token.startswith(SEGMENT_PREFIX):
            token = token[len(SEGMENT_PREFIX):]
        return bytes(self._byte_decoder[c] for c in token)

    # ── encode ───────────────────────────────────────────────────────────
    def _split_ids(self, raw):
        """Greedy longest-match over segments, byte fallback. Returns token ids."""
        ids = []
        i, n = 0, len(raw)
        while i < n:
            # v2: a space directly followed by a word byte opens a
            # [space + word] match window instead of passing through.
            space_chunk = (self._v2 and raw[i] == 0x20 and i + 1 < n
                           and not (raw[i + 1] < 128
                                    and raw[i + 1] in SEPARATORS))
            if raw[i] < 128 and raw[i] in SEPARATORS and not space_chunk:
                ids.append(raw[i])
                i += 1
                continue
            word_end = i + 1 if space_chunk else i
            while word_end < n:
                b = raw[word_end]
                if b < 128 and b in SEPARATORS:
                    break
                word_end += 1
            pos = i
            while pos < word_end:
                hi = min(self._max_seg_len, word_end - pos)
                for slen in range(hi, 0, -1):
                    sid = self._seg_to_id.get(raw[pos:pos + slen])
                    if sid is not None:
                        ids.append(sid)
                        pos += slen
                        break
                else:
                    ids.append(raw[pos])
                    pos += 1
            i = word_end
        return ids

    def _tokenize(self, text, **kwargs):
        return [self._id_to_token[i] for i in self._split_ids(text.encode('utf-8'))]

    def _convert_token_to_id(self, token):
        got = self._token_to_id.get(token)
        if got is None and self.unk_token is not None:
            return self._token_to_id.get(str(self.unk_token))
        return got

    def _convert_id_to_token(self, index):
        return self._id_to_token.get(index, '')

    def convert_tokens_to_string(self, tokens):
        raw = b''.join(self._token_to_bytes(t) for t in tokens
                       if t in self._token_to_id and t != str(self.bos_token))
        return raw.decode('utf-8', errors='replace')

    # ── special tokens ───────────────────────────────────────────────────
    # Training separates documents with BOS, and our CUDA sampler seeds
    # generation from it. Text encoded without a leading BOS starts the model
    # mid-document, which measurably degrades what comes out — so BOS is
    # prepended whenever special tokens are enabled, as LLaMA-style tokenizers
    # do. Pass add_special_tokens=False for raw parity with HarrisTokenizer.
    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
        out = [self._bos_id] + token_ids_0
        if token_ids_1 is not None:
            out = out + token_ids_1
        return out

    def get_special_tokens_mask(self, token_ids_0, token_ids_1=None,
                                already_has_special_tokens=False):
        if already_has_special_tokens:
            return [1 if t == self._bos_id else 0 for t in token_ids_0]
        mask = [1] + [0] * len(token_ids_0)
        if token_ids_1 is not None:
            mask += [0] * len(token_ids_1)
        return mask

    def save_vocabulary(self, save_directory, filename_prefix=None):
        name = (filename_prefix + '-' if filename_prefix else '') + \
            VOCAB_FILES_NAMES['vocab_file']
        path = os.path.join(save_directory, name)
        with open(path, 'w', encoding='utf-8') as f:
            json.dump([list(s) for s in self.segments], f)
        return (path,)