RedHatAI/Kimi-K3-NVFP4

1import os2from logging import getLogger3from pathlib import Path4from shutil import copyfile5from typing import Dict, Iterator, List, Optional, Tuple, Union, cast6 7import tiktoken8from tiktoken.load import load_tiktoken_bpe9from tokenizers import AddedToken10from transformers.convert_slow_tokenizer import bytes_to_unicode11from transformers.tokenization_utils import PreTrainedTokenizer12 13try:14    from .encoding_k3 import build_chat_segments, is_batched_conversation15except ImportError:  # pragma: no cover - supports direct file execution/import.16    from encoding_k3 import build_chat_segments, is_batched_conversation17 18logger = getLogger(__name__)19VOCAB_FILES_NAMES = {"vocab_file": "tiktoken.model"}20 21 22class TikTokenTokenizer(PreTrainedTokenizer):23    """24    Tokenizing and encoding/decoding text using the Tiktoken tokenizer. See megatron/tokenizer/tiktoken_tokenizer.py.25 26    This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to27    this superclass for more information regarding those methods.28 29    Args:30        vocab_file (`str`):31            The path to the Tiktoken model file.32        bos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|begin_of_text|>",`):33            The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token.34        eos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|end_of_text|>"`):35            The end of sequence token.36        unk_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|reserved_special_token_249|>"`):37            The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this38            token instead. The second to last item in special_tokens.39        pad_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|reserved_special_token_250|>"`):40            The token used for padding, for example when batching sequences of different lengths.41        additional_special_tokens (list of `str`, *optional*):42            A tuple or a list of additional tokens, which will be marked as `special`, meaning that they will be43            skipped when decoding if `skip_special_tokens` is set to `True`.44    """45 46    vocab_files_names = VOCAB_FILES_NAMES47 48    model_input_names = ["input_ids", "attention_mask"]49 50    special_tokens: Dict[str, int]51 52    num_reserved_special_tokens = 25653 54    pat_str = "|".join([55        r"""[\p{Han}]+""",56        r"""[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?""",57        r"""[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?""",58        r"""\p{N}{1,3}""",59        r""" ?[^\s\p{L}\p{N}]+[\r\n]*""",60        r"""\s*[\r\n]+""",61        r"""\s+(?!\S)""",62        r"""\s+""",63    ])64 65    def __init__(66        self,67        vocab_file,68        bos_token: Union[str, AddedToken] = "[BOS]",69        eos_token: Union[str, AddedToken] = "[EOS]",70        unk_token: Union[str, AddedToken, None] = None,71        pad_token: Union[str, AddedToken, None] = None,72        additional_special_tokens: List[str] = None,73        added_tokens_decoder: Optional[dict] = None,74        **kwargs,75    ):76        assert os.path.isfile(vocab_file), vocab_file77 78        if additional_special_tokens is None:79            additional_special_tokens = [80                "<|im_end|>",81                "<|im_user|>",82                "<|im_assistant|>",83                "<|start_header_id|>",84                "<|end_header_id|>",85                "[EOT]",86                "<|im_system|>",87                "<|im_middle|>",88            ]89 90        if added_tokens_decoder:91            special_tokens_mapping = {92                i: added_tokens_decoder[i].content93                for i in added_tokens_decoder94            }95        else:96            special_tokens_mapping = {}97 98        self.vocab_file = vocab_file99        mergeable_ranks = load_tiktoken_bpe(vocab_file)100        num_base_tokens = len(mergeable_ranks)101        self.special_tokens = {102            special_tokens_mapping.get(i, f"<|reserved_token_{i}|>"): i103            for i in range(num_base_tokens, num_base_tokens +104                           self.num_reserved_special_tokens)105        }106 107        self.model = tiktoken.Encoding(108            name=Path(vocab_file).name,109            pat_str=self.pat_str,110            mergeable_ranks=mergeable_ranks,111            special_tokens=self.special_tokens,112        )113        logger.info(f"Reloaded tiktoken model from {vocab_file}")114 115        self.n_words: int = self.model.n_vocab116        # BOS / EOS token IDs117        self.bos_id: int = self.special_tokens[str(bos_token)]118        self.eos_id: int = self.special_tokens[str(eos_token)]119        logger.info(120            f"#words: {self.n_words} - BOS ID: {self.bos_id} - EOS ID: {self.eos_id}"121        )122 123        self.pad_id: int = self.special_tokens[str(pad_token)]124        self.unk_id: int = self.special_tokens[str(unk_token)]125 126        self.byte_encoder = bytes_to_unicode()127        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}128 129        self.decoder = {}130        for i in range(self.n_words):131            # Taken from https://gist.github.com/xenova/a452a6474428de0182b17605a98631ee132            decoding = ''.join([133                self.byte_encoder[ord(char)] for char in134                self.model.decode_single_token_bytes(i).decode('latin-1')135            ])136            self.decoder[i] = decoding137 138        self.encoder = {}139        for i in range(self.n_words):140            if i in self.decoder:141                self.encoder[self.decoder[i]] = i142 143        super().__init__(144            bos_token=bos_token,145            eos_token=eos_token,146            unk_token=unk_token,147            pad_token=pad_token,148            additional_special_tokens=additional_special_tokens,149            added_tokens_decoder=added_tokens_decoder,150            **kwargs,151        )152        self.all_special_ids_set = set(self.all_special_ids)153 154    def _encode_text_piece(self, text: str,155                           allow_special_tokens: bool = True) -> List[int]:156        # The tiktoken tokenizer can handle <=400k chars without157        # pyo3_runtime.PanicException.158        TIKTOKEN_MAX_ENCODE_CHARS = 400_000159 160        # https://github.com/openai/tiktoken/issues/195161        # Here we iterate over subsequences and split if we exceed the limit162        # of max consecutive non-whitespace or whitespace characters.163        MAX_NO_WHITESPACES_CHARS = 25_000164 165        t: List[int] = []166        for i in range(0, len(text), TIKTOKEN_MAX_ENCODE_CHARS):167            for substr in self._split_whitespaces_or_nonwhitespaces(168                text[i:i + TIKTOKEN_MAX_ENCODE_CHARS],169                MAX_NO_WHITESPACES_CHARS,170            ):171                if allow_special_tokens:172                    t.extend(173                        # structural markers: encode <|...|> as their special token IDs174                        self.model.encode(175                            substr,176                            allowed_special="all",177                        ))178                else:179                    t.extend(180                        # user/tool text: encode any <|...|> as ordinary BPE tokens (never as control tokens)181                        self.model.encode(182                            substr,183                            disallowed_special=(),184                        ))185 186        return t187 188    def encode(self,189               text: str,190               allow_special_tokens: bool = True,191               **kwargs) -> List[int]:192        """193        Encodes a string into a list of token IDs.194 195        Args:196            text (str): The input string to be encoded.197 198        Returns:199            list[int]: A list of token IDs.200        """201        # If there are other args, we should call super().encode because there are a lot of code202        # to handle those args. supper().encode finally will call _tokenize and _convert_token_to_id.203        # NOTE: our encode method is not compatible with the super().encode method,204        #   e.g. split_special_tokens' default is True in our encode method.205        if len(kwargs) > 0:206            logger.warning(f"Calling super().encode with {kwargs}")207            return super().encode(text, **kwargs)208 209        assert type(text) is str210        return self._encode_text_piece(text,211                                       allow_special_tokens=allow_special_tokens)212 213    def decode(self, token_ids: Union[int, List[int]], **kwargs) -> str:214        """215        Decodes a list of token IDs into a string.216 217        Args:218            token_ids (List[int]): The list of token IDs to be decoded.219 220        Returns:221            str: The decoded string.222        """223        # If there are other args, we should call super().decode because there are a lot of code224        # to handle those args. supper().encode finally will call convert_tokens_to_string and _convert_id_to_token.225        if len(kwargs) > 0:226            return super().decode(token_ids, **kwargs)227 228        if type(token_ids) is int:229            token_ids = [token_ids]230 231        return self.model.decode(cast(List[int], token_ids))232 233    @staticmethod234    def _split_whitespaces_or_nonwhitespaces(235            s: str, max_consecutive_slice_len: int) -> Iterator[str]:236        """237        Splits the string `s` so that each substring contains no more than `max_consecutive_slice_len`238        consecutive whitespaces or consecutive non-whitespaces.239        """240        current_slice_len = 0241        current_slice_is_space = s[0].isspace() if len(s) > 0 else False242        slice_start = 0243 244        for i in range(len(s)):245            is_now_space = s[i].isspace()246 247            if current_slice_is_space ^ is_now_space:248                current_slice_len = 1249                current_slice_is_space = is_now_space250            else:251                current_slice_len += 1252                if current_slice_len > max_consecutive_slice_len:253                    yield s[slice_start:i]254                    slice_start = i255                    current_slice_len = 1256        yield s[slice_start:]257 258    def _encode_chat_segments(self, segments) -> List[int]:259        token_ids: List[int] = []260        for segment in segments:261            token_ids.extend(262                self._encode_text_piece(263                    segment.text,264                    allow_special_tokens=segment.allow_special,265                ))266        return token_ids267 268    @staticmethod269    def _truncate(ids: List[int],270                  truncation: bool = False,271                  max_length: Optional[int] = None) -> List[int]:272        if truncation and max_length is not None:273            return ids[:max_length]274        return ids275 276    def _format_chat_token_output(self,277                                  encoded_inputs: List[List[int]],278                                  *,279                                  is_batched: bool,280                                  padding=False,281                                  truncation: bool = False,282                                  max_length: Optional[int] = None,283                                  return_tensors=None,284                                  return_dict: bool = False):285        encoded_inputs = [286            self._truncate(ids, truncation=truncation, max_length=max_length)287            for ids in encoded_inputs288        ]289 290        needs_batch_encoding = (291            is_batched or padding or return_tensors is not None or return_dict)292        if not needs_batch_encoding:293            return encoded_inputs[0]294 295        features = [{296            "input_ids": ids,297            "attention_mask": [1] * len(ids)298        } for ids in encoded_inputs]299        batch = self.pad(features,300                         padding=padding,301                         max_length=max_length if padding else None,302                         return_attention_mask=True,303                         return_tensors=return_tensors)304 305        if return_dict:306            return batch307        if is_batched:308            return batch["input_ids"]309        return batch["input_ids"][0] if return_tensors is None else batch[310            "input_ids"]311 312    """ ----- Below are the abstract methods required by PreTrainedTokenizer ----- """313 314    @property315    def vocab_size(self) -> int:316        return self.n_words317 318    def get_vocab(self) -> Dict[str, int]:319        return self.encoder320 321    def _tokenize(self, text: str, **kwargs) -> List[str]:322        return [self.decoder[t] for t in self.encode(text)]323 324    def _convert_token_to_id(self, token: str) -> int:325        return self.encoder.get(token, self.unk_id)326 327    def _convert_id_to_token(self, index: int) -> str:328        return self.decoder.get(index)329 330    @staticmethod331    def clean_up_tokenization(out_string: str) -> str:332        return out_string333 334    def convert_tokens_to_string(self, tokens: List[str]) -> str:335        text = ''.join(tokens)336        text = bytearray([self.byte_decoder[c]337                          for c in text]).decode('utf-8', 'replace')338        return text339 340    def save_vocabulary(self,341                        save_directory: str,342                        filename_prefix: Optional[str] = None) -> Tuple[str]:343        if not os.path.isdir(save_directory):344            raise ValueError(345                f"vocabulary path ({save_directory}) should be a directory")346        out_vocab_file = os.path.join(347            save_directory,348            (filename_prefix + "-" if filename_prefix else "") +349            VOCAB_FILES_NAMES["vocab_file"])350 351        if os.path.abspath(self.vocab_file) != os.path.abspath(352                out_vocab_file) and os.path.isfile(self.vocab_file):353            copyfile(self.vocab_file, out_vocab_file)354 355        return (out_vocab_file, )356 357    def apply_chat_template(self,358                            conversation,359                            tools: Optional[list[dict]] = None,360                            tokenize: bool = False,361                            add_generation_prompt: bool = True,362                            thinking: bool = True,363                            padding=False,364                            truncation: bool = False,365                            max_length: Optional[int] = None,366                            return_tensors=None,367                            return_dict: bool = False,368                            **kwargs):369        # Tokenizer-level rendering reorders tool result messages to match370        # assistant tool_calls, normalizes per-call arguments and response371        # schema, then encodes the resulting XTML structure segment-by-segment.372        is_batched = is_batched_conversation(conversation)373        conversations = conversation if is_batched else [conversation]374        image_prompts = kwargs.pop("image_prompts", None)375        if is_batched and image_prompts is not None:376            raise ValueError("image_prompts is only supported for one chat.")377        378        # by default set thinking effort to max379        kwargs.setdefault("thinking_effort", "max")380 381        segment_batches = [382            build_chat_segments(383                messages,384                tools=tools,385                add_generation_prompt=add_generation_prompt,386                thinking=thinking,387                image_prompts=image_prompts,388                **kwargs,389            ) for messages in conversations390        ]391 392        if not tokenize:393            rendered = ["".join(segment.text for segment in segments)394                        for segments in segment_batches]395            return rendered if is_batched else rendered[0]396 397        encoded_inputs = [398            self._encode_chat_segments(segments) for segments in segment_batches399        ]400        return self._format_chat_token_output(401            encoded_inputs,402            is_batched=is_batched,403            padding=padding,404            truncation=truncation,405            max_length=max_length,406            return_tensors=return_tensors,407            return_dict=return_dict,408        )409