import os import shutil from typing import List, Optional, Tuple import sentencepiece as spm from transformers import PreTrainedTokenizer class UnModelTokenizer(PreTrainedTokenizer): """SentencePiece tokenizer for the SAGE POND UnModel encoder. The special-token arguments are explicit constructor parameters. This is important for Transformers releases whose PythonBackend receives tokenizer configuration values through ``from_pretrained``: values such as ``bos_token`` must not also remain in ``kwargs`` when passed to ``super``. """ vocab_files_names = {"vocab_file": "new_bpe.model"} model_input_names = ["input_ids", "attention_mask"] def __init__( self, vocab_file, add_bos_token=False, add_eos_token=False, model_max_length=600, pad_token=None, bos_token=None, eos_token=None, unk_token=None, clean_up_tokenization_spaces=False, **kwargs, ): if not vocab_file or not os.path.isfile(vocab_file): raise FileNotFoundError( "The tokenizer requires the SentencePiece model file 'new_bpe.model'." ) self.vocab_file = os.path.abspath(vocab_file) self.sp_model = spm.SentencePieceProcessor(model_file=self.vocab_file) self.add_bos_token = bool(add_bos_token) self.add_eos_token = bool(add_eos_token) vocab_size = self.sp_model.get_piece_size() if vocab_size != 40_000: raise ValueError( f"Expected a 40,000-piece SentencePiece model, got {vocab_size:,}." ) def special_id(method_name, fallback): value = getattr(self.sp_model, method_name)() return value if value >= 0 else fallback # These fallbacks match the IDs used by the original training script. self._sp_pad_id = special_id("pad_id", 0) self._sp_bos_id = special_id("bos_id", 2) self._sp_eos_id = special_id("eos_id", 3) self._sp_unk_id = special_id("unk_id", 1) self._sp_special_tokens = { "pad_token": self.sp_model.id_to_piece(self._sp_pad_id), "bos_token": self.sp_model.id_to_piece(self._sp_bos_id), "eos_token": self.sp_model.id_to_piece(self._sp_eos_id), "unk_token": self.sp_model.id_to_piece(self._sp_unk_id), } for token_name, token_id in ( ("pad_token", self._sp_pad_id), ("bos_token", self._sp_bos_id), ("eos_token", self._sp_eos_id), ("unk_token", self._sp_unk_id), ): if token_id < 0 or token_id >= vocab_size: raise ValueError( f"{token_name} ID {token_id} is outside vocabulary size {vocab_size}." ) # Validate values restored from tokenizer_config.json against the # SentencePiece model. Keeping the same strings guarantees that token # IDs remain aligned with the model's embedding table. supplied_tokens = { "pad_token": pad_token, "bos_token": bos_token, "eos_token": eos_token, "unk_token": unk_token, } resolved_tokens = {} for name, supplied in supplied_tokens.items(): expected = self._sp_special_tokens[name] if supplied is None: resolved_tokens[name] = expected continue supplied_content = getattr(supplied, "content", str(supplied)) if supplied_content != expected: raise ValueError( f"Configured {name}={supplied_content!r} does not match the " f"SentencePiece token {expected!r}. The tokenizer config and " "new_bpe.model must come from the same training run." ) resolved_tokens[name] = supplied super().__init__( pad_token=resolved_tokens["pad_token"], bos_token=resolved_tokens["bos_token"], eos_token=resolved_tokens["eos_token"], unk_token=resolved_tokens["unk_token"], model_max_length=model_max_length, clean_up_tokenization_spaces=clean_up_tokenization_spaces, **kwargs, ) # Ensure save_pretrained writes the dynamic AutoTokenizer mapping and # the options needed when the tokenizer is reconstructed from the Hub. self.init_kwargs["auto_map"] = { "AutoTokenizer": ["tokenization_un_model.UnModelTokenizer", None] } self.init_kwargs["tokenizer_class"] = self.__class__.__name__ self.init_kwargs["add_bos_token"] = self.add_bos_token self.init_kwargs["add_eos_token"] = self.add_eos_token @property def vocab_size(self): return self.sp_model.get_piece_size() def get_vocab(self): vocab = { self.sp_model.id_to_piece(token_id): token_id for token_id in range(self.vocab_size) } vocab.update(self.added_tokens_encoder) return vocab def _tokenize(self, text, **kwargs): return self.sp_model.encode(text, out_type=str) def _convert_token_to_id(self, token): token_id = self.sp_model.piece_to_id(token) if token_id < 0 or token_id >= self.vocab_size: return self._sp_unk_id return token_id def _convert_id_to_token(self, index): if index < 0 or index >= self.vocab_size: return self.unk_token return self.sp_model.id_to_piece(index) def convert_tokens_to_string(self, tokens: List[str]) -> str: return self.sp_model.decode_pieces(tokens) def build_inputs_with_special_tokens( self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, ) -> List[int]: output = list(token_ids_0) if self.add_bos_token: output.insert(0, self.bos_token_id) if self.add_eos_token: output.append(self.eos_token_id) if token_ids_1 is not None: second = list(token_ids_1) if self.add_bos_token: second.insert(0, self.bos_token_id) if self.add_eos_token: second.append(self.eos_token_id) output.extend(second) return output def get_special_tokens_mask( self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False, ) -> List[int]: if already_has_special_tokens: special_ids = set(self.all_special_ids) return [1 if token_id in special_ids else 0 for token_id in token_ids_0] mask = [] if self.add_bos_token: mask.append(1) mask.extend([0] * len(token_ids_0)) if self.add_eos_token: mask.append(1) if token_ids_1 is not None: if self.add_bos_token: mask.append(1) mask.extend([0] * len(token_ids_1)) if self.add_eos_token: mask.append(1) return mask def create_token_type_ids_from_sequences( self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, ) -> List[int]: combined = self.build_inputs_with_special_tokens(token_ids_0, token_ids_1) return [0] * len(combined) def save_vocabulary( self, save_directory, filename_prefix: Optional[str] = None, ) -> Tuple[str, ...]: if not os.path.isdir(save_directory): raise ValueError(f"Vocabulary save directory does not exist: {save_directory}") filename = ( f"{filename_prefix}-new_bpe.model" if filename_prefix else "new_bpe.model" ) destination = os.path.join(save_directory, filename) if os.path.abspath(self.vocab_file) != os.path.abspath(destination): shutil.copyfile(self.vocab_file, destination) return (destination,) UnModelTokenizer.register_for_auto_class("AutoTokenizer")