Spaces:
Sleeping
Sleeping
| """ | |
| Gemma3 Tokenizer. | |
| Wraps the HuggingFace tokenizers library with Gemma3-specific special | |
| token handling. Uses a hybrid tokenization approach: special tokens | |
| (<start_of_turn>, <end_of_turn>, <think>, etc.) are mapped directly | |
| to their IDs, while regular text is encoded by the underlying BPE | |
| tokenizer with add_special_tokens=False to prevent unwanted <bos> | |
| insertion between fragments. | |
| Chat template follows Gemma3's turn-based format: | |
| <start_of_turn>user\\n{content}<end_of_turn>\\n | |
| <start_of_turn>model\\n | |
| """ | |
| import re | |
| import os | |
| from typing import List, Dict, Union | |
| from pathlib import Path | |
| from tokenizers import Tokenizer | |
| from huggingface_hub import hf_hub_download | |
| from src.utils.model_store import get_model_dir | |
| class Gemma3Tokenizer: | |
| # 1. Define Gemma-specific specials | |
| _SPECIALS = [ | |
| "<bos>", "<eos>", "<unk>", "<pad>", | |
| "<start_of_turn>", "<end_of_turn>", | |
| "<think>", "</think>" | |
| ] | |
| def __init__(self, model_path_or_repo: str): | |
| # ... (Same init logic as before) ... | |
| # 1. Handle Auto-Download | |
| if os.path.exists(model_path_or_repo): | |
| if (Path(model_path_or_repo) / "tokenizer.json").exists(): | |
| file_path = Path(model_path_or_repo) / "tokenizer.json" | |
| else: | |
| file_path = Path(model_path_or_repo) / "tokenizer.model" | |
| else: | |
| # Check centralized model store first | |
| repo_name = Path(model_path_or_repo).parts[-1] | |
| central_path = get_model_dir(repo_name) / "tokenizer.json" | |
| if central_path.exists(): | |
| file_path = central_path | |
| else: | |
| print(f"⬇️ Downloading tokenizer from {model_path_or_repo}...") | |
| file_path = hf_hub_download( | |
| repo_id=model_path_or_repo, | |
| filename="tokenizer.json", | |
| local_dir=str(get_model_dir(repo_name)), | |
| ) | |
| # 2. Load the Tokenizer | |
| self._tok = Tokenizer.from_file(str(file_path)) | |
| # 3. Build Special Token Map | |
| self._special_to_id = {} | |
| self._id_to_special = {} | |
| for t in self._SPECIALS: | |
| tid = self._tok.token_to_id(t) | |
| if tid is not None: | |
| self._special_to_id[t] = tid | |
| self._id_to_special[tid] = t | |
| # 4. Compile Regex | |
| pattern = "|".join(map(re.escape, self._SPECIALS)) | |
| self._split_re = re.compile(f"({pattern})") | |
| # 5. Set Important IDs | |
| self.pad_token_id = self._special_to_id.get("<pad>", self._special_to_id.get("<eos>")) | |
| self.eos_token_id = self._special_to_id.get("<eos>", None) | |
| def apply_chat_template( | |
| self, | |
| messages: List[Dict[str, str]], | |
| add_generation_prompt: bool = False, | |
| add_thinking: bool = False | |
| ) -> str: | |
| formatted_text = "" | |
| # Optional: Add <bos> ONLY ONCE at the very start if needed | |
| # formatted_text += "<bos>" | |
| for msg in messages: | |
| role = msg["role"] | |
| content = msg["content"] | |
| formatted_text += f"<start_of_turn>{role}\n{content}<end_of_turn>\n" | |
| if add_generation_prompt: | |
| formatted_text += "<start_of_turn>model\n" | |
| if add_thinking: | |
| formatted_text += "<think>\n" | |
| return formatted_text | |
| def encode( | |
| self, | |
| text_or_messages: Union[str, List[Dict]], | |
| add_generation_prompt: bool = False, | |
| add_thinking: bool = False | |
| ): | |
| # 1. Handle Chat Logic | |
| if isinstance(text_or_messages, list): | |
| text = self.apply_chat_template( | |
| text_or_messages, | |
| add_generation_prompt=add_generation_prompt, | |
| add_thinking=add_thinking | |
| ) | |
| else: | |
| text = text_or_messages | |
| # 2. Hybrid Tokenization | |
| ids = [] | |
| parts = self._split_re.split(text) | |
| for part in parts: | |
| if not part: continue | |
| if part in self._special_to_id: | |
| ids.append(self._special_to_id[part]) | |
| else: | |
| # --- FIX IS HERE --- | |
| # We disable add_special_tokens so it doesn't add <bos> | |
| # to every single text fragment (like "user", "model", etc.) | |
| ids.extend(self._tok.encode(part, add_special_tokens=False).ids) | |
| return ids | |
| def decode(self, ids, skip_special_tokens=False): | |
| return self._tok.decode(ids, skip_special_tokens=skip_special_tokens) | |
| def vocab_size(self): | |
| return self._tok.get_vocab_size() |