spicyneuron/Kimi-K2.7-Code-MLX-3.6bit

1import os2from collections import OrderedDict3from logging import getLogger4from pathlib import Path5from shutil import copyfile6from typing import Any, Dict, Iterator, List, Optional, Tuple, Union, cast7 8import tiktoken9from tiktoken.load import load_tiktoken_bpe10from tokenizers import AddedToken11from transformers.convert_slow_tokenizer import bytes_to_unicode12from transformers.tokenization_utils import PreTrainedTokenizer13 14from .tool_declaration_ts import encode_tools_to_typescript_style15 16logger = getLogger(__name__)17VOCAB_FILES_NAMES = {"vocab_file": "tiktoken.model"}18 19 20class TikTokenTokenizer(PreTrainedTokenizer):21    """22    Tokenizing and encoding/decoding text using the Tiktoken tokenizer. See megatron/tokenizer/tiktoken_tokenizer.py.23 24    This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to25    this superclass for more information regarding those methods.26 27    Args:28        vocab_file (`str`):29            The path to the Tiktoken model file.30        bos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|begin_of_text|>",`):31            The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token.32        eos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|end_of_text|>"`):33            The end of sequence token.34        unk_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|reserved_special_token_249|>"`):35            The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this36            token instead. The second to last item in special_tokens.37        pad_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|reserved_special_token_250|>"`):38            The token used for padding, for example when batching sequences of different lengths.39        additional_special_tokens (list of `str`, *optional*):40            A tuple or a list of additional tokens, which will be marked as `special`, meaning that they will be41            skipped when decoding if `skip_special_tokens` is set to `True`.42    """43 44    vocab_files_names = VOCAB_FILES_NAMES45 46    model_input_names = ["input_ids", "attention_mask"]47 48    special_tokens: Dict[str, int]49 50    num_reserved_special_tokens = 25651 52    pat_str = "|".join([53        r"""[\p{Han}]+""",54        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)?""",55        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)?""",56        r"""\p{N}{1,3}""",57        r""" ?[^\s\p{L}\p{N}]+[\r\n]*""",58        r"""\s*[\r\n]+""",59        r"""\s+(?!\S)""",60        r"""\s+""",61    ])62 63    def __init__(64        self,65        vocab_file,66        bos_token: Union[str, AddedToken] = "[BOS]",67        eos_token: Union[str, AddedToken] = "[EOS]",68        unk_token: Union[str, AddedToken, None] = None,69        pad_token: Union[str, AddedToken, None] = None,70        additional_special_tokens: List[str] = None,71        added_tokens_decoder: Optional[dict] = None,72        **kwargs,73    ):74        assert os.path.isfile(vocab_file), vocab_file75 76        # Transformers ≥5 may supply ``extra_special_tokens`` instead of77        # ``additional_special_tokens``; treat empty dict as absent.78        if additional_special_tokens is None:79            extra = kwargs.pop("extra_special_tokens", None)80            if isinstance(extra, dict) and not extra:81                extra = None82            if isinstance(extra, (list, tuple)):83                additional_special_tokens = list(extra)84 85        if additional_special_tokens is None:86            additional_special_tokens = [87                "<|im_end|>",88                "<|im_user|>",89                "<|im_assistant|>",90                "<|start_header_id|>",91                "<|end_header_id|>",92                "[EOT]",93                "<|im_system|>",94                "<|im_middle|>",95            ]96 97        if added_tokens_decoder:98            special_tokens_mapping = {99                i: added_tokens_decoder[i].content100                for i in added_tokens_decoder101            }102        else:103            special_tokens_mapping = {}104 105        self.vocab_file = vocab_file106        mergeable_ranks = load_tiktoken_bpe(vocab_file)107        num_base_tokens = len(mergeable_ranks)108        self.special_tokens = {109            special_tokens_mapping.get(i, f"<|reserved_token_{i}|>"): i110            for i in range(num_base_tokens, num_base_tokens +111                           self.num_reserved_special_tokens)112        }113 114        self.model = tiktoken.Encoding(115            name=Path(vocab_file).name,116            pat_str=self.pat_str,117            mergeable_ranks=mergeable_ranks,118            special_tokens=self.special_tokens,119        )120        logger.info(f"Reloaded tiktoken model from {vocab_file}")121 122        self.n_words: int = self.model.n_vocab123        # BOS / EOS token IDs124        self.bos_id: int = self.special_tokens[str(bos_token)]125        self.eos_id: int = self.special_tokens[str(eos_token)]126        logger.info(127            f"#words: {self.n_words} - BOS ID: {self.bos_id} - EOS ID: {self.eos_id}"128        )129 130        self.pad_id: int = self.special_tokens[str(pad_token)]131        self.unk_id: int = self.special_tokens[str(unk_token)]132 133        self.byte_encoder = bytes_to_unicode()134        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}135 136        self.decoder = {}137        for i in range(self.n_words):138            # Taken from https://gist.github.com/xenova/a452a6474428de0182b17605a98631ee139            decoding = ''.join([140                self.byte_encoder[ord(char)] for char in141                self.model.decode_single_token_bytes(i).decode('latin-1')142            ])143            self.decoder[i] = decoding144 145        self.encoder = {}146        for i in range(self.n_words):147            if i in self.decoder:148                self.encoder[self.decoder[i]] = i149 150        self._token_config_cache = OrderedDict()151        self._cache_max_size = 128152 153        super().__init__(154            bos_token=bos_token,155            eos_token=eos_token,156            unk_token=unk_token,157            pad_token=pad_token,158            additional_special_tokens=additional_special_tokens,159            added_tokens_decoder=added_tokens_decoder,160            **kwargs,161        )162        self.all_special_ids_set = set(self.all_special_ids)163 164    def encode(self,165               text: str,166               allow_special_tokens: bool = True,167               **kwargs) -> List[int]:168        """169        Encodes a string into a list of token IDs.170 171        Args:172            text (str): The input string to be encoded.173 174        Returns:175            list[int]: A list of token IDs.176        """177        # If there are other args, we should call super().encode because there are a lot of code178        # to handle those args. supper().encode finally will call _tokenize and _convert_token_to_id.179        # NOTE: our encode method is not compatible with the super().encode method,180        #   e.g. split_special_tokens' default is True in our encode method.181        if len(kwargs) > 0:182            logger.warning(f"Calling super().encode with {kwargs}")183            return super().encode(text, **kwargs)184 185        assert type(text) is str186 187        # The tiktoken tokenizer can handle <=400k chars without188        # pyo3_runtime.PanicException.189        TIKTOKEN_MAX_ENCODE_CHARS = 400_000190 191        # https://github.com/openai/tiktoken/issues/195192        # Here we iterate over subsequences and split if we exceed the limit193        # of max consecutive non-whitespace or whitespace characters.194        MAX_NO_WHITESPACES_CHARS = 25_000195 196        texts = self.pre_tokenizer_process(text)197 198        all_substrs = []199        for text in texts:200            substrs = (201                substr for i in range(0, len(text), TIKTOKEN_MAX_ENCODE_CHARS)202                for substr in self._split_whitespaces_or_nonwhitespaces(203                    text[i:i +204                         TIKTOKEN_MAX_ENCODE_CHARS], MAX_NO_WHITESPACES_CHARS))205            all_substrs.extend(substrs)206 207        t: List[int] = []208        for substr in all_substrs:209            if allow_special_tokens:210                t.extend(211                    # we should consider special token as a common token212                    self.model.encode(213                        substr,214                        allowed_special="all",215                    ))216            else:217                t.extend(218                    # we should consider special token as a common token219                    self.model.encode(220                        substr,221                        disallowed_special=(),222                    ))223 224        return t225 226    def decode(self, token_ids: Union[int, List[int]], **kwargs) -> str:227        """228        Decodes a list of token IDs into a string.229 230        Args:231            token_ids (List[int]): The list of token IDs to be decoded.232 233        Returns:234            str: The decoded string.235        """236        # If there are other args, we should call super().decode because there are a lot of code237        # to handle those args. supper().encode finally will call convert_tokens_to_string and _convert_id_to_token.238        if len(kwargs) > 0:239            return super().decode(token_ids, **kwargs)240 241        if type(token_ids) is int:242            token_ids = [token_ids]243 244        return self.model.decode(cast(List[int], token_ids))245 246    @staticmethod247    def _split_whitespaces_or_nonwhitespaces(248            s: str, max_consecutive_slice_len: int) -> Iterator[str]:249        """250        Splits the string `s` so that each substring contains no more than `max_consecutive_slice_len`251        consecutive whitespaces or consecutive non-whitespaces.252        """253        current_slice_len = 0254        current_slice_is_space = s[0].isspace() if len(s) > 0 else False255        slice_start = 0256 257        for i in range(len(s)):258            is_now_space = s[i].isspace()259 260            if current_slice_is_space ^ is_now_space:261                current_slice_len = 1262                current_slice_is_space = is_now_space263            else:264                current_slice_len += 1265                if current_slice_len > max_consecutive_slice_len:266                    yield s[slice_start:i]267                    slice_start = i268                    current_slice_len = 1269        yield s[slice_start:]270 271    def pre_tokenizer_process(self, text: str) -> List[str]:272        """273        pre-tokenizes the input text into a list of tokens.274        This method is used to split the input text into smaller chunks for internal processing.275        """276        return [text]277 278    """ ----- Below are the abstract methods required by PreTrainedTokenizer ----- """279 280    @property281    def vocab_size(self) -> int:282        return self.n_words283 284    def get_vocab(self) -> Dict[str, int]:285        return self.encoder286 287    def _tokenize(self, text: str, **kwargs) -> List[str]:288        return [self.decoder[t] for t in self.encode(text)]289 290    def _convert_token_to_id(self, token: str) -> int:291        return self.encoder.get(token, self.unk_id)292 293    def _convert_id_to_token(self, index: int) -> str:294        return self.decoder.get(index)295 296    @staticmethod297    def clean_up_tokenization(out_string: str) -> str:298        return out_string299 300    def convert_tokens_to_string(self, tokens: List[str]) -> str:301        text = ''.join(tokens)302        text = bytearray([self.byte_decoder[c]303                          for c in text]).decode('utf-8', 'replace')304        return text305 306    def save_vocabulary(self,307                        save_directory: str,308                        filename_prefix: Optional[str] = None) -> Tuple[str]:309        if not os.path.isdir(save_directory):310            raise ValueError(311                f"vocabulary path ({save_directory}) should be a directory")312        out_vocab_file = os.path.join(313            save_directory,314            (filename_prefix + "-" if filename_prefix else "") +315            VOCAB_FILES_NAMES["vocab_file"])316 317        if os.path.abspath(self.vocab_file) != os.path.abspath(318                out_vocab_file) and os.path.isfile(self.vocab_file):319            copyfile(self.vocab_file, out_vocab_file)320 321        return (out_vocab_file, )322 323    def apply_chat_template(self,324                            conversation,325                            tools: Optional[list[dict]] = None,326                            tokenize: bool = False,327                            add_generation_prompt: bool = True,328                            thinking: bool = True,329                            preserve_thinking: bool = True,330                            **kwargs):331 332        tools = deep_sort_dict(tools)333 334        # Convert tools to TypeScript style string if tools are provided335        tools_ts_str = None336        if tools:337            try:338                tools_ts_str = encode_tools_to_typescript_style(tools)339 340            except Exception as e:341                print(f"Failed to convert tools to TypeScript style: {e}")342                tools_ts_str = None343 344        # Store the TypeScript string in kwargs so it can be accessed by the template345        if tools_ts_str is not None:346            kwargs['tools_ts_str'] = tools_ts_str347 348        if not thinking:349            logger.warning("thinking=False is not supported, overriding to True")350            thinking = True351 352        if not preserve_thinking:353            logger.warning("preserve_thinking=False is not supported, overriding to True")354            preserve_thinking = True355 356        return super().apply_chat_template(357            conversation,358            tools=tools,359            tokenize=tokenize,360            add_generation_prompt=add_generation_prompt,361            thinking=thinking,362            preserve_thinking=preserve_thinking,363            **kwargs)364 365 366def deep_sort_dict(obj: Any) -> Any:367    if isinstance(obj, dict):368        return {k: deep_sort_dict(v) for k, v in sorted(obj.items())}369    if isinstance(obj, list):370        return [deep_sort_dict(item) for item in obj]371    return obj372