llmcompressor.modeling.kimi_k3.tokenization_kimi

Classes:

TikTokenTokenizer –

Tokenizing and encoding/decoding text using the Tiktoken tokenizer. See megatron/tokenizer/tiktoken_tokenizer.py.

TikTokenTokenizer

TikTokenTokenizer(
    vocab_file,
    bos_token: Union[str, AddedToken] = "[BOS]",
    eos_token: Union[str, AddedToken] = "[EOS]",
    unk_token: Union[str, AddedToken, None] = None,
    pad_token: Union[str, AddedToken, None] = None,
    additional_special_tokens: List[str] = None,
    added_tokens_decoder: Optional[dict] = None,
    **kwargs,
)

Bases: PreTrainedTokenizer

Tokenizing and encoding/decoding text using the Tiktoken tokenizer. See megatron/tokenizer/tiktoken_tokenizer.py.

This tokenizer inherits from [PreTrainedTokenizer] which contains most of the main methods. Users should refer to this superclass for more information regarding those methods.

Args: vocab_file (str): The path to the Tiktoken model file. bos_token (str or tokenizers.AddedToken, optional, defaults to "<|begin_of_text|>",): The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token. eos_token (str or tokenizers.AddedToken, optional, defaults to "<|end_of_text|>"): The end of sequence token. unk_token (str or tokenizers.AddedToken, optional, defaults to "<|reserved_special_token_249|>"): The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this token instead. The second to last item in special_tokens. pad_token (str or tokenizers.AddedToken, optional, defaults to "<|reserved_special_token_250|>"): The token used for padding, for example when batching sequences of different lengths. additional_special_tokens (list of str, optional): A tuple or a list of additional tokens, which will be marked as special, meaning that they will be skipped when decoding if skip_special_tokens is set to True.

Methods:

decode –

Decodes a list of token IDs into a string.
encode –

Encodes a string into a list of token IDs.

Source code in src/llmcompressor/modeling/kimi_k3/tokenization_kimi.py

def __init__(
    self,
    vocab_file,
    bos_token: Union[str, AddedToken] = "[BOS]",
    eos_token: Union[str, AddedToken] = "[EOS]",
    unk_token: Union[str, AddedToken, None] = None,
    pad_token: Union[str, AddedToken, None] = None,
    additional_special_tokens: List[str] = None,
    added_tokens_decoder: Optional[dict] = None,
    **kwargs,
):
    assert os.path.isfile(vocab_file), vocab_file

    if additional_special_tokens is None:
        additional_special_tokens = [
            "<|im_end|>",
            "<|im_user|>",
            "<|im_assistant|>",
            "<|start_header_id|>",
            "<|end_header_id|>",
            "[EOT]",
            "<|im_system|>",
            "<|im_middle|>",
        ]

    if added_tokens_decoder:
        special_tokens_mapping = {
            i: added_tokens_decoder[i].content for i in added_tokens_decoder
        }
    else:
        special_tokens_mapping = {}

    self.vocab_file = vocab_file
    mergeable_ranks = load_tiktoken_bpe(vocab_file)
    num_base_tokens = len(mergeable_ranks)
    self.special_tokens = {
        special_tokens_mapping.get(i, f"<|reserved_token_{i}|>"): i
        for i in range(
            num_base_tokens, num_base_tokens + self.num_reserved_special_tokens
        )
    }

    self.model = tiktoken.Encoding(
        name=Path(vocab_file).name,
        pat_str=self.pat_str,
        mergeable_ranks=mergeable_ranks,
        special_tokens=self.special_tokens,
    )
    logger.info(f"Reloaded tiktoken model from {vocab_file}")

    self.n_words: int = self.model.n_vocab
    # BOS / EOS token IDs
    self.bos_id: int = self.special_tokens[str(bos_token)]
    self.eos_id: int = self.special_tokens[str(eos_token)]
    logger.info(
        f"#words: {self.n_words} - BOS ID: {self.bos_id} - EOS ID: {self.eos_id}"
    )

    self.pad_id: int = self.special_tokens[str(pad_token)]
    self.unk_id: int = self.special_tokens[str(unk_token)]

    self.byte_encoder = bytes_to_unicode()
    self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}

    self.decoder = {}
    for i in range(self.n_words):
        # Taken from https://gist.github.com/xenova/a452a6474428de0182b17605a98631ee
        decoding = "".join(
            [
                self.byte_encoder[ord(char)]
                for char in self.model.decode_single_token_bytes(i).decode(
                    "latin-1"
                )
            ]
        )
        self.decoder[i] = decoding

    self.encoder = {}
    for i in range(self.n_words):
        if i in self.decoder:
            self.encoder[self.decoder[i]] = i

    super().__init__(
        bos_token=bos_token,
        eos_token=eos_token,
        unk_token=unk_token,
        pad_token=pad_token,
        additional_special_tokens=additional_special_tokens,
        added_tokens_decoder=added_tokens_decoder,
        **kwargs,
    )
    self.all_special_ids_set = set(self.all_special_ids)

decode

decode(token_ids: Union[int, List[int]], **kwargs) -> str

Decodes a list of token IDs into a string.

Args: token_ids (List[int]): The list of token IDs to be decoded.

Returns: str: The decoded string.

Source code in src/llmcompressor/modeling/kimi_k3/tokenization_kimi.py

def decode(self, token_ids: Union[int, List[int]], **kwargs) -> str:
    """
    Decodes a list of token IDs into a string.

    Args:
        token_ids (List[int]): The list of token IDs to be decoded.

    Returns:
        str: The decoded string.
    """
    # If there are other args, we should call super().decode because there are a lot of code
    # to handle those args. supper().encode finally will call convert_tokens_to_string and _convert_id_to_token.
    if len(kwargs) > 0:
        return super().decode(token_ids, **kwargs)

    if type(token_ids) is int:
        token_ids = [token_ids]

    return self.model.decode(cast(List[int], token_ids))

encode

encode(
    text: str, allow_special_tokens: bool = True, **kwargs
) -> List[int]

Encodes a string into a list of token IDs.

Args: text (str): The input string to be encoded.

Returns: list[int]: A list of token IDs.

Source code in src/llmcompressor/modeling/kimi_k3/tokenization_kimi.py

def encode(
    self, text: str, allow_special_tokens: bool = True, **kwargs
) -> List[int]:
    """
    Encodes a string into a list of token IDs.

    Args:
        text (str): The input string to be encoded.

    Returns:
        list[int]: A list of token IDs.
    """
    # If there are other args, we should call super().encode because there are a lot of code
    # to handle those args. supper().encode finally will call _tokenize and _convert_token_to_id.
    # NOTE: our encode method is not compatible with the super().encode method,
    #   e.g. split_special_tokens' default is True in our encode method.
    if len(kwargs) > 0:
        logger.warning(f"Calling super().encode with {kwargs}")
        return super().encode(text, **kwargs)

    assert type(text) is str
    return self._encode_text_piece(text, allow_special_tokens=allow_special_tokens)