tokenization_atom.py

"""Remote-code tokenizer for Atom GPT checkpoints.

The tokenizer is intentionally HF-compatible: generic callers can use
``AutoTokenizer.from_pretrained(..., trust_remote_code=True)``. Arithmetic digit
spans are reversed before tokenization so the model receives LSD-first numbers,
matching pretraining.
"""

from __future__ import annotations

import re
from typing import Any

from transformers import PreTrainedTokenizerFast


class AtomTokenizer(PreTrainedTokenizerFast):
    vocab_files_names = {"tokenizer_file": "tokenizer.json"}
    model_input_names = ["input_ids", "attention_mask"]
    slow_tokenizer_class = None
    _digit_span_re = re.compile(r"\d+")

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        kwargs.setdefault("bos_token", "<|bos|>")
        kwargs.setdefault("eos_token", "<|eos|>")
        kwargs.setdefault("unk_token", "<|unk|>")
        kwargs.setdefault("pad_token", "<|pad|>")
        super().__init__(*args, **kwargs)

    @classmethod
    def _reverse_digit_spans(cls, text: str) -> str:
        return cls._digit_span_re.sub(lambda match: match.group(0)[::-1], text)

    @classmethod
    def _transform_text(cls, value: Any) -> Any:
        if isinstance(value, str):
            return cls._reverse_digit_spans(value)
        if isinstance(value, tuple):
            return tuple(cls._transform_text(item) for item in value)
        if isinstance(value, list):
            return [cls._transform_text(item) for item in value]
        return value

    def __call__(self, text=None, text_pair=None, *args: Any, **kwargs: Any):
        return super().__call__(
            self._transform_text(text),
            self._transform_text(text_pair),
            *args,
            **kwargs,
        )

    def encode(self, text, text_pair=None, *args: Any, **kwargs: Any):
        return super().encode(
            self._transform_text(text),
            self._transform_text(text_pair),
            *args,
            **kwargs,
        )

    def batch_encode_plus(self, batch_text_or_text_pairs, *args: Any, **kwargs: Any):
        return super().batch_encode_plus(
            self._transform_text(batch_text_or_text_pairs),
            *args,
            **kwargs,
        )

    def _decode(self, token_ids, skip_special_tokens: bool = False, **kwargs: Any) -> str:
        text = super()._decode(
            token_ids,
            skip_special_tokens=skip_special_tokens,
            **kwargs,
        )
        return self._reverse_digit_spans(text)