"""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 = "" BOS_TOKEN = "" EOS_TOKEN = "" 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)