tokenization_dendro_omni.py

"""Deterministic UTF-8 byte tokenizer for Dendro Omni."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Iterable, Sequence

import torch

try:  # pragma: no cover - optional Hugging Face branch.
    from transformers import PreTrainedTokenizer

    _HF_TOKENIZER = True
except Exception:  # pragma: no cover - fallback is covered.
    PreTrainedTokenizer = object  # type: ignore[assignment,misc]
    _HF_TOKENIZER = False


PAD_TOKEN = "<pad>"
BOS_TOKEN = "<bos>"
EOS_TOKEN = "<eos>"
BYTE_OFFSET = 3


def _byte_token(value: int) -> str:
    return f"<0x{value:02X}>"


def _token_byte(token: str) -> int | None:
    if len(token) == 6 and token.startswith("<0x") and token.endswith(">"):
        try:
            value = int(token[3:5], 16)
        except ValueError:
            return None
        return value if 0 <= value <= 255 else None
    return None


class _ByteTokenizerLogic:
    pad_token_id = 0
    bos_token_id = 1
    eos_token_id = 2
    byte_offset = BYTE_OFFSET
    model_input_names = ["input_ids", "attention_mask"]

    @property
    def vocab_size(self) -> int:
        return self.byte_offset + 256

    def raw_encode(self, text: str, *, add_bos: bool = False, add_eos: bool = False) -> list[int]:
        ids: list[int] = []
        if add_bos:
            ids.append(self.bos_token_id)
        ids.extend(byte + self.byte_offset for byte in text.encode("utf-8"))
        if add_eos:
            ids.append(self.eos_token_id)
        return ids

    def raw_decode(self, token_ids: Iterable[int], *, skip_special_tokens: bool = True) -> str:
        data = bytearray()
        literal: list[str] = []
        for item in token_ids:
            token_id = int(item)
            if token_id in (self.pad_token_id, self.bos_token_id, self.eos_token_id):
                if skip_special_tokens:
                    continue
                literal.append({0: PAD_TOKEN, 1: BOS_TOKEN, 2: EOS_TOKEN}[token_id])
                continue
            value = token_id - self.byte_offset
            if 0 <= value <= 255:
                if literal:
                    data.extend("".join(literal).encode("utf-8"))
                    literal.clear()
                data.append(value)
        if literal:
            data.extend("".join(literal).encode("utf-8"))
        return data.decode("utf-8", errors="replace")


if _HF_TOKENIZER:
    class DendroByteTokenizer(_ByteTokenizerLogic, PreTrainedTokenizer):
        """Hugging Face tokenizer with lossless byte coverage and no learned files."""

        vocab_files_names = {"vocab_file": "dendro_vocab.json"}
        model_input_names = ["input_ids", "attention_mask"]

        def __init__(
            self,
            vocab_file: str | None = None,
            *,
            model_max_length: int = 131_072,
            **kwargs: Any,
        ) -> None:
            del vocab_file
            # ``from_pretrained`` restores these values through ``kwargs``.  Use
            # defaults instead of passing a second explicit copy so this remains
            # compatible with both legacy tokenizers and the Transformers 5
            # Python backend.
            kwargs.setdefault("pad_token", PAD_TOKEN)
            kwargs.setdefault("bos_token", BOS_TOKEN)
            kwargs.setdefault("eos_token", EOS_TOKEN)
            kwargs.setdefault("model_max_length", model_max_length)
            kwargs.setdefault("clean_up_tokenization_spaces", False)
            super().__init__(**kwargs)

        def _tokenize(self, text: str, **kwargs: Any) -> list[str]:
            del kwargs
            return [_byte_token(byte) for byte in text.encode("utf-8")]

        def _convert_token_to_id(self, token: str) -> int:
            if token == PAD_TOKEN:
                return self.pad_token_id
            if token == BOS_TOKEN:
                return self.bos_token_id
            if token == EOS_TOKEN:
                return self.eos_token_id
            value = _token_byte(token)
            return self.eos_token_id if value is None else value + self.byte_offset

        def _convert_id_to_token(self, index: int) -> str:
            index = int(index)
            if index == self.pad_token_id:
                return PAD_TOKEN
            if index == self.bos_token_id:
                return BOS_TOKEN
            if index == self.eos_token_id:
                return EOS_TOKEN
            value = index - self.byte_offset
            return _byte_token(value) if 0 <= value <= 255 else EOS_TOKEN

        def convert_tokens_to_string(self, tokens: list[str]) -> str:
            data = bytearray()
            text_parts: list[str] = []
            for token in tokens:
                value = _token_byte(token)
                if value is None:
                    if data:
                        text_parts.append(data.decode("utf-8", errors="replace"))
                        data.clear()
                    if token not in {PAD_TOKEN, BOS_TOKEN, EOS_TOKEN}:
                        text_parts.append(token)
                else:
                    data.append(value)
            if data:
                text_parts.append(data.decode("utf-8", errors="replace"))
            return "".join(text_parts)

        def get_vocab(self) -> dict[str, int]:
            vocab = {PAD_TOKEN: 0, BOS_TOKEN: 1, EOS_TOKEN: 2}
            vocab.update({_byte_token(value): value + self.byte_offset for value in range(256)})
            return vocab

        def build_inputs_with_special_tokens(
            self,
            token_ids_0: list[int],
            token_ids_1: list[int] | None = None,
        ) -> list[int]:
            result = [self.bos_token_id, *token_ids_0, self.eos_token_id]
            if token_ids_1 is not None:
                result.extend([*token_ids_1, self.eos_token_id])
            return result

        def get_special_tokens_mask(
            self,
            token_ids_0: list[int],
            token_ids_1: list[int] | None = None,
            already_has_special_tokens: bool = False,
        ) -> list[int]:
            if already_has_special_tokens:
                return [int(token in {0, 1, 2}) for token in token_ids_0]
            mask = [1] + [0] * len(token_ids_0) + [1]
            if token_ids_1 is not None:
                mask += [0] * len(token_ids_1) + [1]
            return mask

        def save_vocabulary(
            self,
            save_directory: str,
            filename_prefix: str | None = None,
        ) -> tuple[str]:
            directory = Path(save_directory)
            directory.mkdir(parents=True, exist_ok=True)
            name = f"{filename_prefix + '-' if filename_prefix else ''}dendro_vocab.json"
            path = directory / name
            path.write_text(json.dumps(self.get_vocab(), indent=2, sort_keys=True), encoding="utf-8")
            return (str(path),)

        def encode_bytes(self, text: str, *, add_bos: bool = False, add_eos: bool = False) -> list[int]:
            return self.raw_encode(text, add_bos=add_bos, add_eos=add_eos)

        def decode_bytes(self, ids: Iterable[int], *, skip_special_tokens: bool = True) -> str:
            return self.raw_decode(ids, skip_special_tokens=skip_special_tokens)

else:
    class DendroByteTokenizer(_ByteTokenizerLogic):
        """PyTorch-only fallback implementing the common tokenizer call surface."""

        def __init__(
            self,
            vocab_file: str | None = None,
            *,
            model_max_length: int = 131_072,
            padding_side: str = "right",
            **_: Any,
        ) -> None:
            del vocab_file
            self.model_max_length = int(model_max_length)
            self.padding_side = str(padding_side)
            self.pad_token = PAD_TOKEN
            self.bos_token = BOS_TOKEN
            self.eos_token = EOS_TOKEN

        def get_vocab(self) -> dict[str, int]:
            vocab = {PAD_TOKEN: 0, BOS_TOKEN: 1, EOS_TOKEN: 2}
            vocab.update({_byte_token(value): value + self.byte_offset for value in range(256)})
            return vocab

        def encode(
            self,
            text: str,
            *,
            add_special_tokens: bool = True,
            add_bos: bool | None = None,
            add_eos: bool | None = None,
            **_: Any,
        ) -> list[int]:
            return self.raw_encode(
                text,
                add_bos=add_special_tokens if add_bos is None else add_bos,
                add_eos=add_special_tokens if add_eos is None else add_eos,
            )

        def decode(
            self,
            token_ids: Iterable[int] | torch.Tensor,
            *,
            skip_special_tokens: bool = True,
            **_: Any,
        ) -> str:
            if torch.is_tensor(token_ids):
                token_ids = token_ids.detach().cpu().tolist()
            return self.raw_decode(token_ids, skip_special_tokens=skip_special_tokens)

        def batch_decode(
            self,
            sequences: Sequence[Sequence[int] | torch.Tensor],
            *,
            skip_special_tokens: bool = True,
            **kwargs: Any,
        ) -> list[str]:
            return [self.decode(row, skip_special_tokens=skip_special_tokens, **kwargs) for row in sequences]

        def __call__(
            self,
            text: str | Sequence[str],
            *,
            add_special_tokens: bool = True,
            padding: bool | str = False,
            truncation: bool = False,
            max_length: int | None = None,
            return_tensors: str | None = None,
            **_: Any,
        ) -> dict[str, Any]:
            texts = [text] if isinstance(text, str) else list(text)
            rows = [self.encode(item, add_special_tokens=add_special_tokens) for item in texts]
            limit = self.model_max_length if max_length is None else int(max_length)
            if truncation:
                rows = [row[:limit] for row in rows]
            target = max((len(row) for row in rows), default=0) if padding else None
            if padding == "max_length":
                target = limit
            masks: list[list[int]] = []
            if target is not None:
                padded: list[list[int]] = []
                for row in rows:
                    amount = max(0, target - len(row))
                    if self.padding_side == "left":
                        padded.append([self.pad_token_id] * amount + row[:target])
                        masks.append([0] * amount + [1] * min(len(row), target))
                    else:
                        padded.append(row[:target] + [self.pad_token_id] * amount)
                        masks.append([1] * min(len(row), target) + [0] * amount)
                rows = padded
            else:
                masks = [[1] * len(row) for row in rows]
            result: dict[str, Any] = {"input_ids": rows, "attention_mask": masks}
            if return_tensors is not None:
                if return_tensors != "pt":
                    raise ValueError("The local tokenizer supports return_tensors='pt' only")
                if not padding and len({len(row) for row in rows}) > 1:
                    raise ValueError("Batch tensor output requires padding")
                result = {key: torch.tensor(value, dtype=torch.long) for key, value in result.items()}
            if isinstance(text, str) and return_tensors is None:
                result = {key: value[0] for key, value in result.items()}
            return result

        def save_pretrained(self, save_directory: str | Path, **_: Any) -> tuple[str, ...]:
            directory = Path(save_directory)
            directory.mkdir(parents=True, exist_ok=True)
            vocab_path = directory / "dendro_vocab.json"
            config_path = directory / "tokenizer_config.json"
            special_path = directory / "special_tokens_map.json"
            vocab_path.write_text(json.dumps(self.get_vocab(), indent=2, sort_keys=True), encoding="utf-8")
            config_path.write_text(
                json.dumps(
                    {
                        "tokenizer_class": "DendroByteTokenizer",
                        "model_max_length": self.model_max_length,
                        "padding_side": self.padding_side,
                        "auto_map": {"AutoTokenizer": ["tokenization_dendro_omni.DendroByteTokenizer", None]},
                    },
                    indent=2,
                ),
                encoding="utf-8",
            )
            special_path.write_text(
                json.dumps(
                    {"pad_token": PAD_TOKEN, "bos_token": BOS_TOKEN, "eos_token": EOS_TOKEN},
                    indent=2,
                ),
                encoding="utf-8",
            )
            return str(vocab_path), str(config_path), str(special_path)

        @classmethod
        def from_pretrained(cls, path: str | Path, **kwargs: Any) -> "DendroByteTokenizer":
            directory = Path(path)
            config_path = directory / "tokenizer_config.json"
            if config_path.exists():
                config = json.loads(config_path.read_text(encoding="utf-8"))
                kwargs.setdefault("model_max_length", config.get("model_max_length", 131_072))
                kwargs.setdefault("padding_side", config.get("padding_side", "right"))
            return cls(vocab_file=str(directory / "dendro_vocab.json"), **kwargs)