chore: import upstream snapshot with attribution

python/mlc_llm/tokenizers/tokenizers.py

@@ -0,0 +1,127 @@
+"""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))