wehub-resource-sync
128 lines
4.0 KiB
Python
128 lines
4.0 KiB
Python
"""The tokenizer and related tools in MLC LLM.
|
|
This tokenizer essentially wraps and binds the HuggingFace tokenizer
|
|
library and sentencepiece.
|
|
Reference: https://github.com/mlc-ai/tokenizers-cpp
|
|
"""
|
|
|
|
import json
|
|
from dataclasses import asdict, dataclass
|
|
from typing import List, Literal # noqa: UP035
|
|
|
|
import tvm_ffi
|
|
from tvm.runtime import Object
|
|
|
|
from . import _ffi_api
|
|
|
|
|
|
@dataclass
|
|
class TokenizerInfo:
|
|
"""Useful information of the tokenizer during generation.
|
|
|
|
Attributes
|
|
----------
|
|
token_postproc_method : Literal["byte_fallback", "byte_level"]
|
|
The method to post-process the tokens to their original strings.
|
|
Possible values (each refers to a kind of tokenizer):
|
|
- "byte_fallback": The same as the byte-fallback BPE tokenizer, including LLaMA-2,
|
|
Mixtral-7b, etc. E.g. "▁of" -> " of", "<0x1B>" -> "\x1b".
|
|
This method:
|
|
1) Transform tokens like <0x1B> to hex char byte 1B. (so-called byte-fallback)
|
|
2) Replace \\u2581 "▁" with space.
|
|
- "byte_level": The same as the byte-level BPE tokenizer, including LLaMA-3, GPT-2,
|
|
Phi-2, etc. E.g. "Ġin" -> " in", "ě" -> "\x1b"
|
|
This method inverses the bytes-to-unicode transformation in the encoding process in
|
|
https://github.com/huggingface/transformers/blob/87be06ca77166e6a6215eee5a990ab9f07238a18/src/transformers/models/gpt2/tokenization_gpt2.py#L38-L59
|
|
|
|
prepend_space_in_encode : bool
|
|
Whether to prepend a space during encoding.
|
|
|
|
strip_space_in_decode : bool
|
|
Whether to strip the first space during decoding.
|
|
"""
|
|
|
|
token_postproc_method: Literal["byte_fallback", "byte_level"] = "byte_fallback"
|
|
prepend_space_in_encode: bool = False
|
|
strip_space_in_decode: bool = False
|
|
|
|
def asjson(self) -> str:
|
|
"""Return the config in string of JSON format."""
|
|
return json.dumps(asdict(self))
|
|
|
|
@staticmethod
|
|
def from_json(json_str: str) -> "TokenizerInfo":
|
|
"""Construct a config from JSON string."""
|
|
return TokenizerInfo(**json.loads(json_str))
|
|
|
|
|
|
@tvm_ffi.register_object("mlc.Tokenizer")
|
|
class Tokenizer(Object):
|
|
"""The tokenizer class in MLC LLM."""
|
|
|
|
def __init__(self, tokenizer_path: str) -> None:
|
|
"""Create the tokenizer from tokenizer directory path."""
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.Tokenizer,
|
|
tokenizer_path,
|
|
)
|
|
|
|
def encode(self, text: str) -> List[int]: # noqa: UP006
|
|
"""Encode text into ids.
|
|
|
|
Parameters
|
|
----------
|
|
text : str
|
|
The text string to encode.
|
|
|
|
Returns
|
|
-------
|
|
token_ids : List[int]
|
|
The list of encoded token ids.
|
|
"""
|
|
return list(_ffi_api.TokenizerEncode(self, text))
|
|
|
|
def encode_batch(self, texts: List[str]) -> List[List[int]]: # noqa: UP006
|
|
"""Encode a batch of texts into ids.
|
|
|
|
Parameters
|
|
----------
|
|
texts : List[str]
|
|
The list of text strings to encode.
|
|
|
|
Returns
|
|
-------
|
|
token_ids : List[List[int]]
|
|
The list of list of encoded token ids.
|
|
"""
|
|
return list(_ffi_api.TokenizerEncodeBatch(self, texts))
|
|
|
|
def decode(self, token_ids: List[int]) -> str: # noqa: UP006
|
|
"""Decode token ids into text.
|
|
|
|
Parameters
|
|
----------
|
|
token_ids : List[int]
|
|
The token ids to decode to string.
|
|
|
|
Returns
|
|
-------
|
|
text : str
|
|
The decoded text string.
|
|
"""
|
|
return _ffi_api.TokenizerDecode(self, tvm_ffi.Shape(token_ids))
|
|
|
|
@staticmethod
|
|
def detect_tokenizer_info(tokenizer_path: str) -> TokenizerInfo:
|
|
"""Detect the tokenizer info from the given path of the tokenizer.
|
|
|
|
Parameters
|
|
----------
|
|
tokenizer_path : str
|
|
The tokenizer directory path.
|
|
|
|
Returns
|
|
-------
|
|
tokenizer_info : str
|
|
The detected tokenizer info in JSON string.
|
|
"""
|
|
return TokenizerInfo.from_json(_ffi_api.DetectTokenizerInfo(tokenizer_path))
|