| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tokenizer for Kimi-Audio using TikToken.""" |
| 4 | |
| 5 | import contextlib |
| 6 | import json |
| 7 | from collections.abc import Sequence |
| 8 | from pathlib import Path |
| 9 | from typing import Any, overload |
| 10 | |
| 11 | import pybase64 |
| 12 | import tiktoken |
| 13 | from transformers import AddedToken, BatchEncoding |
| 14 | from transformers.utils import chat_template_utils as hf_chat_utils |
| 15 | |
| 16 | from vllm.entrypoints.chat_utils import ChatCompletionMessageParam |
| 17 | from vllm.logger import init_logger |
| 18 | from vllm.tokenizers.protocol import TokenizerLike |
| 19 | from vllm.transformers_utils.repo_utils import hf_api |
| 20 | |
| 21 | logger = init_logger(__name__) |
| 22 | |
| 23 | |
| 24 | def _load_tiktoken_encoding( |
| 25 | vocab_file: Path, special_tokens: dict[str, int] |
| 26 | ) -> tuple[Any, dict[str, int]]: |
| 27 | """Load TikToken encoding from vocab file.""" |
| 28 | mergeable_ranks: dict[bytes, int] = {} |
| 29 | with open(vocab_file, encoding="utf-8") as f: |
| 30 | for line in f: |
| 31 | line = line.strip() |
| 32 | if not line: |
| 33 | continue |
| 34 | parts = line.split() |
| 35 | if len(parts) == 2: |
| 36 | token_b64 = parts[0] |
| 37 | rank = int(parts[1]) |
| 38 | token_bytes = pybase64.b64decode(token_b64) |
| 39 | mergeable_ranks[token_bytes] = rank |
| 40 | |
| 41 | tokenizer = tiktoken.Encoding( |
| 42 | name=str(vocab_file), |
| 43 | pat_str=r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}|""" |
| 44 | r""" ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+""", |
| 45 | mergeable_ranks=mergeable_ranks, |
| 46 | special_tokens=special_tokens, |
| 47 | ) |
| 48 | |
| 49 | return tokenizer, special_tokens |
| 50 | |
| 51 | |
| 52 | class KimiAudioTokenizer(TokenizerLike): |
| 53 | """TikToken tokenizer for Kimi-Audio.""" |
| 54 | |
| 55 | @classmethod |
| 56 | def from_pretrained( |
| 57 | cls, |
| 58 | path_or_repo_id: str | Path, |
| 59 | *args, |
| 60 | trust_remote_code: bool = False, |
| 61 | revision: str | None = None, |
| 62 | download_dir: str | None = None, |
| 63 | **kwargs, |
| 64 | ) -> "KimiAudioTokenizer": |
| 65 | if args: |
| 66 | logger.debug_once("Ignoring extra positional args for KimiAudioTokenizer.") |
| 67 | |
| 68 | path = Path(path_or_repo_id) |
| 69 | if path.is_file(): |
| 70 | vocab_file = path |
| 71 | elif path.is_dir(): |
| 72 | vocab_file = path / "tiktoken.model" |
| 73 | if not vocab_file.is_file(): |
| 74 | vocab_file = path / "tokenizer.model" |
| 75 | else: |
| 76 | # Download from HuggingFace Hub |
| 77 | repo_id = str(path_or_repo_id) |
| 78 | |
| 79 | # Try to download tiktoken.model or tokenizer.model |
| 80 | try: |
| 81 | vocab_path = hf_api().hf_hub_download( |
| 82 | repo_id=repo_id, |
| 83 | filename="tiktoken.model", |
| 84 | revision=revision, |
| 85 | local_dir=download_dir, |
| 86 | ) |
| 87 | vocab_file = Path(vocab_path) |
| 88 | except Exception: |
| 89 | try: |
| 90 | vocab_path = hf_api().hf_hub_download( |
| 91 | repo_id=repo_id, |
| 92 | filename="tokenizer.model", |
| 93 | revision=revision, |
| 94 | local_dir=download_dir, |
| 95 | ) |
| 96 | vocab_file = Path(vocab_path) |
| 97 | except Exception as exc: |
| 98 | raise ValueError( |
| 99 | f"Could not find tiktoken.model or tokenizer.model in {repo_id}" |
| 100 | ) from exc |
| 101 | |
| 102 | # Also download tokenizer_config.json if available |
| 103 | with contextlib.suppress(Exception): |
| 104 | hf_api().hf_hub_download( |
| 105 | repo_id=repo_id, |
| 106 | filename="tokenizer_config.json", |
| 107 | revision=revision, |
| 108 | local_dir=download_dir, |
| 109 | ) |
| 110 | |
| 111 | if not vocab_file.is_file(): |
| 112 | raise FileNotFoundError(f"tiktoken.model not found at {vocab_file}.") |
| 113 | |
| 114 | return cls( |
| 115 | vocab_file=vocab_file, |
| 116 | name_or_path=str(path_or_repo_id), |
| 117 | truncation_side=kwargs.get("truncation_side", "left"), |
| 118 | ) |
| 119 | |
| 120 | def __init__( |
| 121 | self, |
| 122 | *, |
| 123 | vocab_file: Path, |
| 124 | name_or_path: str, |
| 125 | truncation_side: str, |
| 126 | ) -> None: |
| 127 | super().__init__() |
| 128 | self.name_or_path = name_or_path |
| 129 | self._truncation_side = truncation_side |
| 130 | self._vocab_file = vocab_file |
| 131 | |
| 132 | # Load special tokens from tokenizer_config.json |
| 133 | special_tokens: dict[str, int] = {} |
| 134 | tokenizer_config = vocab_file.parent / "tokenizer_config.json" |
| 135 | if tokenizer_config.is_file(): |
| 136 | with open(tokenizer_config, encoding="utf-8") as f: |
| 137 | config = json.load(f) |
| 138 | # Extract special tokens from added_tokens_decoder |
| 139 | added_tokens = config.get("added_tokens_decoder", {}) |
| 140 | for token_id_str, token_info in added_tokens.items(): |
| 141 | token_id = int(token_id_str) |
| 142 | content = token_info.get("content", "") |
| 143 | if content: |
| 144 | special_tokens[content] = token_id |
| 145 | |
| 146 | self._tokenizer, self._special_tokens = _load_tiktoken_encoding( |
| 147 | vocab_file, special_tokens |
| 148 | ) |
| 149 | |
| 150 | # Build token <-> ID mappings |
| 151 | self._token_to_id: dict[str, int] = {} |
| 152 | self._id_to_token: dict[int, str] = {} |
| 153 | for token_bytes, token_id in self._tokenizer._mergeable_ranks.items(): |
| 154 | token_str = token_bytes.decode("utf-8", errors="replace") |
| 155 | self._token_to_id[token_str] = token_id |
| 156 | self._id_to_token[token_id] = token_str |
| 157 | |
| 158 | # Initialize added_tokens_decoder before adding special tokens |
| 159 | self._added_tokens_decoder: dict[int, Any] = {} |
| 160 | |
| 161 | # Add Kimi-Audio special tokens |
| 162 | self._add_kimiaudio_special_tokens() |
| 163 | |
| 164 | # Set default special token IDs (will be updated when special tokens are added) |
| 165 | self._bos_token_id = 151643 # Kimi-Audio BOS |
| 166 | self._eos_token_id = 151644 # Kimi-Audio EOS |
| 167 | self._pad_token_id = self._eos_token_id |
| 168 | self._unk_token_id = self._pad_token_id |
| 169 | |
| 170 | self._max_chars_per_token = max( |
| 171 | (len(tok) for tok in self._token_to_id), default=10 |
| 172 | ) |
| 173 | |
| 174 | def _add_kimiaudio_special_tokens(self) -> None: |
| 175 | """Add Kimi-Audio special tokens to the tokenizer.""" |
| 176 | # Tokens should already be in self._special_tokens from tokenizer_config.json |
| 177 | # Just add them to added_tokens_decoder for compatibility |
| 178 | kimiaudio_special_tokens = { |
| 179 | "<|im_media_begin|>": 151661, |
| 180 | "<|im_media_end|>": 151663, |
| 181 | "<|im_kimia_text_blank|>": 151666, |
| 182 | "<|im_msg_end|>": 151645, |
| 183 | "<|im_kimia_user_msg_start|>": 151670, |
| 184 | "<|im_kimia_assistant_msg_start|>": 151671, |
| 185 | } |
| 186 | |
| 187 | for token_str, token_id in kimiaudio_special_tokens.items(): |
| 188 | # Only add if not already present |
| 189 | if token_id not in self._added_tokens_decoder: |
| 190 | self._added_tokens_decoder[token_id] = AddedToken( |
| 191 | token_str, single_word=True, normalized=False, special=True |
| 192 | ) |
| 193 | # Also ensure it's in _token_to_id and _id_to_token |
| 194 | if token_str not in self._token_to_id: |
| 195 | self._token_to_id[token_str] = token_id |
| 196 | if token_id not in self._id_to_token: |
| 197 | self._id_to_token[token_id] = token_str |
| 198 | |
| 199 | def num_special_tokens_to_add(self) -> int: |
| 200 | return 0 |
| 201 | |
| 202 | @property |
| 203 | def all_special_tokens(self) -> list[str]: |
| 204 | return list(self._added_tokens_decoder.values()) |
| 205 | |
| 206 | @property |
| 207 | def all_special_ids(self) -> list[int]: |
| 208 | return list(self._added_tokens_decoder.keys()) |
| 209 | |
| 210 | @property |
| 211 | def bos_token_id(self) -> int: |
| 212 | return self._bos_token_id |
| 213 | |
| 214 | @property |
| 215 | def eos_token_id(self) -> int: |
| 216 | return self._eos_token_id |
| 217 | |
| 218 | @property |
| 219 | def pad_token_id(self) -> int: |
| 220 | return self._pad_token_id |
| 221 | |
| 222 | @property |
| 223 | def is_fast(self) -> bool: |
| 224 | return False |
| 225 | |
| 226 | @property |
| 227 | def vocab_size(self) -> int: |
| 228 | return self._tokenizer.n_vocab |
| 229 | |
| 230 | @property |
| 231 | def max_token_id(self) -> int: |
| 232 | return self._tokenizer.n_vocab - 1 |
| 233 | |
| 234 | @property |
| 235 | def max_chars_per_token(self) -> int: |
| 236 | return self._max_chars_per_token |
| 237 | |
| 238 | @property |
| 239 | def truncation_side(self) -> str: |
| 240 | return self._truncation_side |
| 241 | |
| 242 | @property |
| 243 | def added_tokens_decoder(self) -> dict[int, Any]: |
| 244 | return self._added_tokens_decoder |
| 245 | |
| 246 | @added_tokens_decoder.setter |
| 247 | def added_tokens_decoder(self, value: dict[int, Any]) -> None: |
| 248 | """Set added tokens decoder and update special token IDs.""" |
| 249 | self._added_tokens_decoder = value |
| 250 | # Update special token IDs if known tokens are added |
| 251 | for token_id, token in value.items(): |
| 252 | token_str = str(token) if hasattr(token, "__str__") else token |
| 253 | if "<|im_kimia_user_msg_start|>" in token_str: |
| 254 | self._bos_token_id = token_id |
| 255 | elif "<|im_msg_end|>" in token_str or "<|im_end|>" in token_str: |
| 256 | self._eos_token_id = token_id |
| 257 | |
| 258 | def get_vocab(self) -> dict[str, int]: |
| 259 | return dict(self._token_to_id) |
| 260 | |
| 261 | def __len__(self) -> int: |
| 262 | """Return vocab size for compatibility with HF tokenizer interface.""" |
| 263 | return self._tokenizer.n_vocab |
| 264 | |
| 265 | def get_added_vocab(self) -> dict[str, int]: |
| 266 | return { |
| 267 | str(token): token_id |
| 268 | for token_id, token in self._added_tokens_decoder.items() |
| 269 | } |
| 270 | |
| 271 | def _maybe_truncate(self, tokens: list[int], max_length: int | None) -> list[int]: |
| 272 | if max_length is None or len(tokens) <= max_length: |
| 273 | return tokens |
| 274 | if self.truncation_side == "left": |
| 275 | return tokens[-max_length:] |
| 276 | return tokens[:max_length] |
| 277 | |
| 278 | def encode( |
| 279 | self, |
| 280 | text: str, |
| 281 | truncation: bool | None = None, |
| 282 | max_length: int | None = None, |
| 283 | add_special_tokens: bool = True, |
| 284 | **kwargs, |
| 285 | ) -> list[int]: |
| 286 | del add_special_tokens |
| 287 | # Allow Kimi-Audio special tokens to be encoded |
| 288 | tokens = self._tokenizer.encode( |
| 289 | text, |
| 290 | allowed_special={ |
| 291 | "<|im_media_begin|>", |
| 292 | "<|im_media_end|>", |
| 293 | "<|im_kimia_text_blank|>", |
| 294 | "<|im_msg_end|>", |
| 295 | "<|im_kimia_user_msg_start|>", |
| 296 | "<|im_kimia_assistant_msg_start|>", |
| 297 | }, |
| 298 | ) |
| 299 | if truncation: |
| 300 | tokens = self._maybe_truncate(tokens, max_length) |
| 301 | return tokens |
| 302 | |
| 303 | def decode( |
| 304 | self, ids: Sequence[int] | int, skip_special_tokens: bool = False |
| 305 | ) -> str: |
| 306 | """Decode token IDs to text, optionally skipping special tokens.""" |
| 307 | if isinstance(ids, int): |
| 308 | ids = [ids] |
| 309 | if skip_special_tokens: |
| 310 | # Skip tokens that are in special_tokens (loaded from config) |
| 311 | special_ids = set(self._special_tokens.values()) |
| 312 | ids = [token_id for token_id in ids if token_id not in special_ids] |
| 313 | return self._tokenizer.decode(ids) |
| 314 | |
| 315 | @overload |
| 316 | def convert_tokens_to_ids(self, tokens: str) -> int: ... |
| 317 | |
| 318 | @overload |
| 319 | def convert_tokens_to_ids(self, tokens: list[str]) -> list[int]: ... |
| 320 | |
| 321 | def convert_tokens_to_ids(self, tokens: str | list[str]) -> int | list[int]: |
| 322 | if isinstance(tokens, str): |
| 323 | return self._token_to_id.get(tokens, self._unk_token_id) |
| 324 | return [self._token_to_id.get(token, self._unk_token_id) for token in tokens] |
| 325 | |
| 326 | def convert_ids_to_tokens( |
| 327 | self, ids: Sequence[int], skip_special_tokens: bool = False |
| 328 | ) -> list[str]: |
| 329 | tokens = [] |
| 330 | for token_id in ids: |
| 331 | if skip_special_tokens and token_id in self._added_tokens_decoder: |
| 332 | continue |
| 333 | tokens.append(self._id_to_token.get(token_id, "<|unk|>")) |
| 334 | return tokens |
| 335 | |
| 336 | def convert_tokens_to_string(self, tokens: list[str]) -> str: |
| 337 | token_ids = self.convert_tokens_to_ids(tokens) |
| 338 | return self.decode(token_ids, skip_special_tokens=False) |
| 339 | |
| 340 | def __call__( |
| 341 | self, |
| 342 | text: str | list[str], |
| 343 | text_pair: str | None = None, |
| 344 | add_special_tokens: bool = True, |
| 345 | truncation: bool = False, |
| 346 | max_length: int | None = None, |
| 347 | **kwargs, |
| 348 | ) -> BatchEncoding: |
| 349 | if text_pair is not None: |
| 350 | raise NotImplementedError( |
| 351 | "text_pair is not supported for KimiAudioTokenizer." |
| 352 | ) |
| 353 | |
| 354 | if isinstance(text, list): |
| 355 | input_ids_batch: list[list[int]] = [ |
| 356 | self.encode( |
| 357 | item, |
| 358 | truncation=truncation, |
| 359 | max_length=max_length, |
| 360 | add_special_tokens=add_special_tokens, |
| 361 | ) |
| 362 | for item in text |
| 363 | ] |
| 364 | attention_mask_batch = [[1] * len(ids) for ids in input_ids_batch] |
| 365 | return BatchEncoding( |
| 366 | {"input_ids": input_ids_batch, "attention_mask": attention_mask_batch} |
| 367 | ) |
| 368 | |
| 369 | input_ids = self.encode( |
| 370 | text, |
| 371 | truncation=truncation, |
| 372 | max_length=max_length, |
| 373 | add_special_tokens=add_special_tokens, |
| 374 | ) |
| 375 | attention_mask = [1] * len(input_ids) |
| 376 | return BatchEncoding({"input_ids": input_ids, "attention_mask": attention_mask}) |
| 377 | |
| 378 | def get_chat_template( |
| 379 | self, chat_template: str | None, tools: list[dict[str, Any]] | None = None |
| 380 | ) -> str | None: |
| 381 | del tools |
| 382 | return chat_template |
| 383 | |
| 384 | def apply_chat_template( |
| 385 | self, |
| 386 | messages: list[ChatCompletionMessageParam] | None = None, |
| 387 | tools: list[dict[str, Any]] | None = None, |
| 388 | chat_template: str | None = None, |
| 389 | tokenize: bool = False, |
| 390 | **kwargs, |
| 391 | ) -> str | list[int]: |
| 392 | # Handle both 'messages' (protocol) and 'conversation' (caller) parameter names |
| 393 | conversation = messages if messages is not None else kwargs.get("conversation") |
| 394 | if conversation is None: |
| 395 | raise ValueError("Either 'messages' or 'conversation' must be provided.") |
| 396 | template = self.get_chat_template(chat_template, tools=tools) |
| 397 | if template is None: |
| 398 | raise ValueError( |
| 399 | "No chat template available. Provide `chat_template` explicitly." |
| 400 | ) |
| 401 | # Use render_jinja_template instead of apply_chat_template |
| 402 | # Note: render_jinja_template returns ([prompts], [generation_indices]) |
| 403 | rendered, _ = hf_chat_utils.render_jinja_template( |
| 404 | conversation, |
| 405 | chat_template=template, |
| 406 | tools=tools, |
| 407 | **kwargs, |
| 408 | ) |
| 409 | # Extract the first (and usually only) prompt |
| 410 | prompt = rendered[0] if rendered else "" |
| 411 | if tokenize: |
| 412 | return self.encode(prompt, add_special_tokens=False) |
| 413 | return prompt |