""" Tokenizer module with BPE implementation and auto-loading support. """ from dataclasses import dataclass import json from pathlib import Path from typing import Any, Dict, List, Optional, Union from tokenizers import Tokenizer from tokenizers import decoders, processors, normalizers, pre_tokenizers from tokenizers.models import BPE from tokenizers.trainers import BpeTrainer from jinja2 import Template DEFAULT_SPECIAL_TOKENS = { "bos_token": "<｜begin▁of▁sentence｜>", "eos_token": "<｜end▁of▁sentence｜>", "pad_token": "<｜▁pad▁｜>", } CONTROL_TOKENS = [ "<｜begin▁of▁sentence｜>", "<｜end▁of▁sentence｜>", "<｜▁pad▁｜>", ] SPECIAL_TOKENS = ["<｜im▁start｜>", "<｜im▁end｜>"] def train_bpe_tokenizer( files: List[str], vocab_size: int, min_freq: int = 2, reserved_token_size: int = 100, max_token_length: int = 18, ) -> Tokenizer: reserved_tokens = [ f"<｜reserve{i:02d}｜>" for i in range(reserved_token_size - len(SPECIAL_TOKENS)) ] detail_vocab_size = vocab_size - (len(reserved_tokens) + len(SPECIAL_TOKENS)) alphabet = pre_tokenizers.ByteLevel.alphabet() min_size = len(alphabet) + len(CONTROL_TOKENS) assert detail_vocab_size > min_size tokenizer = Tokenizer(BPE()) tokenizer.normalizer = normalizers.Sequence([normalizers.NFC(), normalizers.Strip()]) tokenizer.pre_tokenizer = pre_tokenizers.Sequence([ pre_tokenizers.UnicodeScripts(), pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=True), ]) tokenizer.decoder = decoders.ByteLevel() tokenizer.post_processor = processors.ByteLevel(trim_offsets=True) trainer = BpeTrainer( vocab_size=detail_vocab_size, min_frequency=min_freq, limit_alphabet=detail_vocab_size // 6, max_token_length=max_token_length, special_tokens=CONTROL_TOKENS + SPECIAL_TOKENS, initial_alphabet=alphabet, show_progress=True, ) tokenizer.train(files=files, trainer=trainer) tokenizer.add_special_tokens(CONTROL_TOKENS + SPECIAL_TOKENS + reserved_tokens) return tokenizer # Message type for chat messages type MessageType = Dict[str, Any] @dataclass class ChatTemplate: """A chat template with Jinja2 rendering support. Attributes: name: Unique identifier for the template. template_str: Jinja2 template string. description: Optional description. default_variables: Optional dictionary of default variable values that will be passed to the template if not overridden during rendering. special_tokens: Optional dictionary mapping token names to their string values. These tokens are automatically added to the template variables. """ name: str template_str: str description: str = "" default_variables: Dict[str, Any] = None special_tokens: Dict[str, str] = None def __post_init__(self): if self.default_variables is None: self.default_variables = {} if self.special_tokens is None: self.special_tokens = {} @classmethod def from_string( cls, template_str: str, description: str = "", default_variables: Optional[Dict[str, Any]] = None, special_tokens: Optional[Dict[str, str]] = None, ) -> "ChatTemplate": """Create a ChatTemplate instance directly from a template string.""" return cls( name="", # empty name for ad‑hoc templates template_str=template_str, description=description, default_variables=default_variables, special_tokens=special_tokens, ) def render( self, messages: List[MessageType], system_prompt: Optional[str] = None, **extra_variables: Any, ) -> str: """Render the template with given messages and variables. Args: messages: List of message dicts with 'role' and 'content'. system_prompt: Optional system prompt string. **extra_variables: Additional variables to pass to the template. These override default_variables and special_tokens. Returns: Rendered prompt string. """ # Merge default variables, special tokens, and extra variables variables = {**self.default_variables, **self.special_tokens, **extra_variables} variables["messages"] = messages if system_prompt is not None: variables["system_prompt"] = system_prompt jinja_template = Template(self.template_str) return jinja_template.render(**variables) class AutoTokenizer: """Base tokenizer class with automatic loading support""" TOKENIZER_CLASSES = {} # Registry for auto-loading def __init__( self, path: Optional[Union[str, Path]] = None, special_token_map: Optional[Dict[str, str]] = None, chat_template: Optional[str] = None, ): self._tokenizer: Tokenizer = None self._chat_template: Optional[ChatTemplate] = None self._special_token_map: Optional[Dict] = special_token_map or {} if chat_template: self.set_chat_template(chat_template) if path: self.load(path) def load(self, path: Union[str, Path]): """Load tokenizer from directory.""" path = Path(path) tokenizer_file = path / "tokenizer.json" config_file = path / "tokenizer_config.json" self._tokenizer = Tokenizer.from_file(str(tokenizer_file)) if config_file.exists(): with open(config_file, "r", encoding="utf-8") as f: config = json.load(f) if "special_tokens" in config: self._special_token_map.update(config["special_tokens"]) # Load chat template from config if "chat_template" in config: self.set_chat_template(config["chat_template"]) @classmethod def from_pretrained(cls, path: Union[str, Path], **kwargs) -> "AutoTokenizer": """Load tokenizer from pretrained directory.""" instance = cls(path) return instance def save_pretrained(self, save_path: str): """ Save tokenizer to pretrained directory. Args: save_path: Path to save the tokenizer """ save_path = Path(save_path) save_path.mkdir(parents=True, exist_ok=True) # Save tokenizer self._tokenizer.save(str(save_path / "tokenizer.json")) # Save tokenizer config config = {} if self._special_token_map is not None: config["special_tokens"] = self._special_token_map if self._chat_template is not None: config["chat_template"] = self._chat_template.template_str with open(save_path / "tokenizer_config.json", "w", encoding="utf-8") as f: json.dump(config, f, ensure_ascii=False, indent=2) @classmethod def register_tokenizer(cls, name: str, tokenizer_class: type): """ Register a new tokenizer class. Args: name: Name to register the tokenizer class under tokenizer_class: The tokenizer class to register """ cls.TOKENIZER_CLASSES[name] = tokenizer_class def encode( self, tokens: Union[str, List[str]], out_ids: bool = True, is_pretokenized: bool = False, add_special_tokens: bool = True, ) -> List: """Encode text to tokens or token IDs.""" if self._tokenizer is None: raise RuntimeError( "Tokenizer not initialized. Load or create a tokenizer first." ) if isinstance(tokens, str): encoded = self._tokenizer.encode( tokens, is_pretokenized=is_pretokenized, add_special_tokens=add_special_tokens, ) return encoded.ids if out_ids else encoded.tokens else: encoded_list = self._tokenizer.encode_batch( tokens, is_pretokenized=is_pretokenized, add_special_tokens=add_special_tokens, ) return [ encoded.ids if out_ids else encoded.tokens for encoded in encoded_list ] def decode(self, tokens: List[int], skip_special_tokens: bool = True) -> str: """Decode token IDs to text.""" if self._tokenizer is None: raise RuntimeError( "Tokenizer not initialized. Load or create a tokenizer first." ) return self._tokenizer.decode(tokens, skip_special_tokens=skip_special_tokens) def __len__(self) -> int: if self._tokenizer is None: return 0 return self._tokenizer.get_vocab_size() def __getattr__(self, key: str): """ Dynamically intercept special token attribute access. Supports three forms: - tokenizer.bos_token → returns string - tokenizer.bos_token_id → returns corresponding integer ID - tokenizer.stop_ids → returns list of corresponding integer IDs for all special tokens """ # Handle stop_ids - return IDs for all special tokens if key == "stop_ids": stop_ids = [] if self._tokenizer is None: return stop_ids for val in self._special_token_map.values(): token_id = self._tokenizer.token_to_id(val) if token_id is not None: stop_ids.append(token_id) return stop_ids # Handle _id suffix (e.g., bos_token_id -> bos_token) if key.endswith("_id"): base_attr = key[:-3] # Remove "_id" token_str = self._special_token_map.get(base_attr) if token_str is None: return None if self._tokenizer is None: raise RuntimeError("Tokenizer not loaded, cannot convert token to id.") return self._tokenizer.token_to_id(token_str) # Handle regular string attributes if key in self._special_token_map: return self._special_token_map.get(key) # Other attributes trigger default AttributeError raise AttributeError(f"'{type(self).__name__}' object has no attribute '{key}'") @property def vocab_size(self) -> int: return len(self) def set_chat_template(self, template: Union[str, ChatTemplate]): """ Set the chat template for the tokenizer. Args: template: Either a template name (str) registered in the global registry, or a ChatTemplate instance, or a Jinja2 template string. Raises: KeyError: If template name is not registered. """ if isinstance(template, str): self._chat_template = ChatTemplate.from_string(template) elif isinstance(template, ChatTemplate): self._chat_template = template else: raise ValueError("Invalid template type, must be str or ChatTemplate.") def apply_chat_template( self, messages: List[Dict[str, str]], system_prompt: Optional[str] = None, tokenize: bool = True, add_generation_prompt: bool = True, **kwargs, ) -> Union[str, List[int]]: """ Apply the chat template to messages and optionally tokenize the result. Args: messages: List of message dicts with 'role' and 'content'. system_prompt: Optional system prompt string (auto-converted to first message). tokenize: Whether to return token IDs (True) or raw string (False). add_generation_prompt: Whether to add the generation prompt (default: True). **kwargs: Additional variables to pass to the template. Returns: Either the rendered string or list of token IDs. Raises: RuntimeError: If chat template is not set. """ if self._chat_template is None: raise RuntimeError( "Chat template not set. Use set_chat_template() to set a template first." ) # Auto-convert system_prompt to first message if provided if system_prompt: messages = [{"role": "system", "content": system_prompt}] + list(messages) # Render the template rendered = self._chat_template.render( messages=messages, add_generation_prompt=add_generation_prompt, **kwargs, ) if tokenize: return self.encode(rendered) return rendered