"""SentencePiece tokenizer wrapper and special-token metadata helpers.""" import json from pathlib import Path from typing import Optional BUILTIN_SPECIAL_TOKENS = ("", "", "", "") def get_special_token_metadata_path(tokenizer_path: str | Path) -> Path: """Return the sidecar metadata path for a tokenizer artifact.""" tokenizer_path = Path(tokenizer_path) return tokenizer_path.with_suffix(".special_tokens.json") def load_special_token_metadata(tokenizer_path: str | Path) -> dict[str, int]: """Load resolved special-token IDs saved alongside a tokenizer.""" metadata_path = get_special_token_metadata_path(tokenizer_path) if not metadata_path.exists(): return {} with metadata_path.open("r", encoding="utf-8") as handle: data = json.load(handle) special_tokens = data.get("special_tokens", {}) return {str(token): int(token_id) for token, token_id in special_tokens.items()} def load_special_token_metadata_payload(tokenizer_path: str | Path) -> dict[str, object]: """Load the raw special-token metadata payload, if available.""" metadata_path = get_special_token_metadata_path(tokenizer_path) if not metadata_path.exists(): return {} with metadata_path.open("r", encoding="utf-8") as handle: return json.load(handle) def load_tokenizer_runtime_info( tokenizer_path: str, tokenizer_type: Optional[str] = None, ) -> dict[str, int]: """Load basic runtime tokenizer information needed for model/data validation.""" if tokenizer_type is None: tokenizer_type = "sentencepiece" if tokenizer_path.endswith(".model") else "huggingface" if tokenizer_type == "sentencepiece": try: import sentencepiece as spm except ImportError: raise ImportError("SentencePiece not installed. Install with: pip install sentencepiece") processor = spm.SentencePieceProcessor() processor.Load(tokenizer_path) return { "vocab_size": int(processor.vocab_size()), "unk_id": int(processor.unk_id()), "bos_id": int(processor.bos_id()), "eos_id": int(processor.eos_id()), "pad_id": int(processor.pad_id()), } from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) vocab_size = getattr(tokenizer, "vocab_size", None) if vocab_size is None: vocab_size = len(tokenizer) return { "vocab_size": int(vocab_size), "unk_id": int(getattr(tokenizer, "unk_token_id", -1)), "bos_id": int(getattr(tokenizer, "bos_token_id", -1)), "eos_id": int(getattr(tokenizer, "eos_token_id", -1)), "pad_id": int(getattr(tokenizer, "pad_token_id", -1)), } def resolve_special_token_id(tokenizer, token: str) -> Optional[int]: """Compatibility helper that delegates to tokenizer-native special-token lookup.""" if hasattr(tokenizer, "get_special_token_id"): token_id = tokenizer.get_special_token_id(token) if token_id is not None: return int(token_id) if hasattr(tokenizer, "convert_tokens_to_ids"): token_id = tokenizer.convert_tokens_to_ids(token) unk_token_id = getattr(tokenizer, "unk_token_id", None) if token_id is not None and token_id != unk_token_id: return int(token_id) return None def require_special_token_id(tokenizer, token: str) -> int: """Resolve a special token ID or raise a clear error.""" if hasattr(tokenizer, "require_special_token_id"): return int(tokenizer.require_special_token_id(token)) token_id = resolve_special_token_id(tokenizer, token) if token_id is None: raise ValueError( f"Special token `{token}` is not registered in the tokenizer. " "Retrain or regenerate the tokenizer with that token registered as a special token." ) return int(token_id) class SentencePieceTokenizerWrapper: """Wrapper to make SentencePiece tokenizer compatible with a HuggingFace-like interface.""" def __init__(self, sp_processor, special_token_ids: Optional[dict[str, int]] = None): """Initialize wrapper.""" self.sp = sp_processor self.vocab_size = self.sp.vocab_size() self.special_token_ids = self._build_special_token_ids(special_token_ids or {}) self.pad_token_id = self.special_token_ids[""] self.eos_token_id = self.special_token_ids[""] self.bos_token_id = self.special_token_ids[""] self.unk_token_id = self.special_token_ids[""] def _build_special_token_ids(self, special_token_ids: dict[str, int]) -> dict[str, int]: """Merge processor-native built-ins with any saved custom special-token IDs.""" resolved_ids = { "": int(self.sp.unk_id()), "": int(self.sp.bos_id()), "": int(self.sp.eos_id()), "": int(self.sp.pad_id()), } for token, token_id in special_token_ids.items(): resolved_ids[str(token)] = int(token_id) invalid_builtins = [token for token in BUILTIN_SPECIAL_TOKENS if resolved_ids.get(token, -1) < 0] if invalid_builtins: raise ValueError( "Tokenizer is missing required built-in special token IDs: " f"{invalid_builtins}. Retrain or regenerate the tokenizer with built-in PAD/BOS/EOS/UNK tokens enabled." ) return resolved_ids def get_special_token_id(self, token: str) -> Optional[int]: """Resolve a special token ID from the wrapper registry or exact SentencePiece pieces.""" token_id = self.special_token_ids.get(token) if token_id is not None: return int(token_id) if hasattr(self.sp, "piece_to_id") and hasattr(self.sp, "id_to_piece"): piece_id = self.sp.piece_to_id(token) if piece_id is not None and piece_id >= 0 and self.sp.id_to_piece(piece_id) == token: return int(piece_id) return None def require_special_token_id(self, token: str) -> int: """Require that a special token resolve to a single known token ID.""" token_id = self.get_special_token_id(token) if token_id is None: raise ValueError( f"Special token `{token}` is not registered in the tokenizer wrapper. " "Retrain or regenerate the tokenizer with that token registered as a special token." ) return int(token_id) def __call__(self, text, **kwargs): """Tokenize text.""" is_single = isinstance(text, str) texts = [text] if is_single else text max_length = kwargs.get("max_length") padding = kwargs.get("padding") truncation = kwargs.get("truncation", False) return_attention_mask = kwargs.get("return_attention_mask", True) all_input_ids = [] for t in texts: if isinstance(t, str): special_token_id = self.get_special_token_id(t) if special_token_id is not None: tokens = [special_token_id] else: tokens = self.sp.encode(t, out_type=int) else: tokens = self.sp.encode(t, out_type=int) if truncation and max_length and len(tokens) > max_length: tokens = tokens[:max_length] all_input_ids.append(tokens) if padding or max_length: target_length = max_length or (max(len(ids) for ids in all_input_ids) if all_input_ids else 0) padded_input_ids = [] padded_attention_masks = [] for ids in all_input_ids: pad_length = target_length - len(ids) if pad_length > 0: padded_ids = ids + [self.pad_token_id] * pad_length else: padded_ids = ids[:target_length] padded_input_ids.append(padded_ids) attention_mask = [1] * min(len(ids), target_length) + [0] * max(0, target_length - len(ids)) padded_attention_masks.append(attention_mask) result = { "input_ids": padded_input_ids if not is_single else padded_input_ids[0], } if return_attention_mask: result["attention_mask"] = ( padded_attention_masks if not is_single else padded_attention_masks[0] ) return result result = { "input_ids": all_input_ids[0] if is_single else all_input_ids, } if return_attention_mask: attention_masks = [[1] * len(ids) for ids in all_input_ids] result["attention_mask"] = attention_masks[0] if is_single else attention_masks return result def encode(self, text, return_tensors=None, **kwargs): """Encode text to token IDs.""" result = self(text, **kwargs) input_ids = result["input_ids"] if return_tensors == "pt": import torch if isinstance(input_ids[0], list): input_ids = input_ids[0] return torch.tensor(input_ids, dtype=torch.long).unsqueeze(0) return input_ids def encode_plus(self, text, **kwargs): """Encode text with additional information.""" return self(text, **kwargs) def decode(self, token_ids, skip_special_tokens=False, **kwargs): """Decode token IDs to text.""" if hasattr(token_ids, "tolist"): token_ids = token_ids.tolist() if isinstance(token_ids, (list, tuple)) and token_ids and isinstance(token_ids[0], (list, tuple)): token_ids = token_ids[0] if not isinstance(token_ids, list): token_ids = [int(t) for t in token_ids] if skip_special_tokens: special_token_ids = {int(token_id) for token_id in self.special_token_ids.values()} token_ids = [int(token_id) for token_id in token_ids if int(token_id) not in special_token_ids] return self.sp.decode(token_ids)