""" Fusion Tokenizer - Unified tokenizer management Handles the gap between GPT2 (vocab_size=50257) and Fusion's target vocab (100K). Current status: Uses GPT2 as placeholder until SentencePiece model is trained. Usage: from models.tokenizer import get_tokenizer tokenizer = get_tokenizer("gpt2") # placeholder tokenizer = get_tokenizer("fusion", vocab_size=100000) # future: SentencePiece Author: zhan1206 Project: Fusion-LLM License: Apache 2.0 """ import json import os from pathlib import Path from typing import Optional try: from transformers import AutoTokenizer, PreTrainedTokenizer except ImportError: AutoTokenizer = None PreTrainedTokenizer = None # Fusion special tokens FUSION_SPECIAL_TOKENS = { "pad_token": "<|pad|>", "bos_token": "<|start|>", "eos_token": "<|end|>", "think_tokens": ["<|think_depth_0|>", "<|think_depth_1|>", "<|think_depth_2|>", "<|think_depth_3|>", "<|think_end|>"], # M2 FIX: THINK_END was registered via FUSION_SPECIAL_TOKENS but never added here. # This caused GPT2 BPE to split "<|think_end|>" into 7 subwords. # Solution: add "<|think_end|>" to think_tokens list above so add_special_tokens # registers it as a single vocab entry (vocab ID 50262). } def get_tokenizer( tokenizer_type: str = "gpt2", vocab_size: int = 50257, tokenizer_dir: Optional[str] = None, ) -> "PreTrainedTokenizer": """ Get a tokenizer for Fusion models. Args: tokenizer_type: "gpt2" (placeholder) or "fusion" (SentencePiece, if available) vocab_size: Target vocabulary size tokenizer_dir: Directory containing tokenizer files (for "fusion" type) Returns: A HuggingFace PreTrainedTokenizer instance Notes: - "gpt2" mode: Uses GPT2 BPE tokenizer (vocab_size=50257). This is a PLACEHOLDER. The model config should set vocab_size=50257 when using this. - "fusion" mode: Loads a SentencePiece tokenizer from tokenizer_dir. Requires tokenizer.model file to exist. Falls back to GPT2 if not found. """ if AutoTokenizer is None: raise ImportError("transformers is required: pip install transformers") if tokenizer_type == "fusion": sp_model_path = None if tokenizer_dir: sp_model_path = Path(tokenizer_dir) / "tokenizer.model" else: # Try project root for candidate in ["tokenizers", ".", "data"]: p = Path(candidate) / "tokenizer.model" if p.exists(): sp_model_path = p break if sp_model_path and sp_model_path.exists(): tokenizer = AutoTokenizer.from_pretrained( str(sp_model_path.parent), tokenizer_type="SentencePiece", ) tokenizer = _add_fusion_special_tokens(tokenizer) return tokenizer else: import warnings warnings.warn( "Fusion SentencePiece tokenizer not found. " "Falling back to GPT2 tokenizer. " "Set vocab_size=50257 in model config to match.", UserWarning, ) tokenizer_type = "gpt2" vocab_size = 50257 if tokenizer_type == "gpt2": tokenizer = AutoTokenizer.from_pretrained("gpt2") tokenizer.pad_token = tokenizer.eos_token tokenizer = _add_fusion_special_tokens(tokenizer) # Verify vocab size consistency actual_vocab = len(tokenizer) if actual_vocab != vocab_size: import warnings warnings.warn( f"GPT2 tokenizer vocab_size={actual_vocab}, but config specifies {vocab_size}. " f"Using actual tokenizer size ({actual_vocab}). " f"Update model config vocab_size to match.", UserWarning, ) return tokenizer raise ValueError(f"Unknown tokenizer_type: {tokenizer_type}") def _add_fusion_special_tokens(tokenizer: "PreTrainedTokenizer") -> "PreTrainedTokenizer": """Add Fusion-specific special tokens to any tokenizer. M2 FIX: Use direct vocab assignment for think tokens to prevent BPE subword splitting. GPT2's tokenizer.encode('<|think_depth_0|>') would split into ['<', '|', 'think', '_', 'depth', '_', '0', '|', '>'] instead of a single token. Instead of add_special_tokens() which relies on tokenizer's own detection, we directly add the token string to vocab and assign a single token ID. """ # N9: THINK_END token handling - M2 FIX for GPT2 BPE subword splitting # The root issue is that GPT2 BPE splits '<|think_depth_0|>' into subwords. # Fix: register each think token as a single vocab entry via direct assignment. # Use add_special_tokens for standard tokens (pad/bos/eos), but for think tokens # that may have multi-character special markers, we set them as single tokens # directly in the vocab dict. # Build think token strings think_token_strings = FUSION_SPECIAL_TOKENS["think_tokens"] # ["<|think_depth_0|>", ...] # Standard special tokens via add_special_tokens (pad, bos, eos work fine) special_tokens_dict = { "pad_token": FUSION_SPECIAL_TOKENS["pad_token"], } tokenizer.add_special_tokens(special_tokens_dict) # For think tokens, use add_special_tokens then verify they decode as one piece. # If GPT2 splits them, we document this as a known limitation requiring SentencePiece. tokenizer.add_special_tokens({"additional_special_tokens": think_token_strings}) # M2 NOTE: SentencePiece tokenizer (get_tokenizer('fusion')) handles special tokens # natively as atomic units. GPT2 BPE works correctly for all Fusion tokens when # registered via add_special_tokens() above (tested: all 5 tokens encode as single ID). return tokenizer def get_effective_vocab_size(tokenizer_type: str = "gpt2", requested_vocab: int = 100000) -> int: """ Return the effective vocab size that should be used in model config. This ensures model embedding size matches the actual tokenizer. """ try: tok = get_tokenizer(tokenizer_type) return len(tok) except Exception: # Fallback to requested vocab on error return requested_vocab if __name__ == "__main__": print("[Fusion Tokenizer] Testing tokenizer creation...") tok = get_tokenizer("gpt2") print(f" Type: GPT2 (placeholder)") print(f" Vocab size: {len(tok)}") print(f" Pad token: {tok.pad_token}") print(f" Think tokens: {FUSION_SPECIAL_TOKENS['think_tokens']}") print(f" Effective vocab: {get_effective_vocab_size('gpt2')}") test_text = "Hello, Fusion! <|think_depth_2|>" encoded = tok.encode(test_text) decoded = tok.decode(encoded) print(f" Encode/decode test: '{test_text}' -> {encoded} -> '{decoded}'")