chore: import upstream snapshot with attribution

python/mlc_llm/tokenizers/__init__.py

@@ -0,0 +1,4 @@
+"""Namespace for tokenizer rleated utilities"""
+
+from .streamer import StopStrHandler, TextStreamer
+from .tokenizers import Tokenizer

python/mlc_llm/tokenizers/_ffi_api.py

@@ -0,0 +1,7 @@
+"""FFI APIs for mlc_llm"""
+
+import tvm_ffi
+
+# Exports functions registered via TVM_FFI_REGISTER_GLOBAL with the "mlc" prefix.
+# e.g. TVM_FFI_REGISTER_GLOBAL("mlc.Tokenizer")
+tvm_ffi.init_ffi_api("mlc.tokenizers", __name__)

python/mlc_llm/tokenizers/streamer.py

@@ -0,0 +1,83 @@
+"""Streamers in MLC LLM."""
+
+from typing import List, Union  # noqa: UP035
+
+import tvm_ffi
+from tvm.runtime import Object
+from tvm_ffi import Shape
+
+from . import _ffi_api
+from .tokenizers import Tokenizer
+
+
+@tvm_ffi.register_object("mlc.TextStreamer")
+class TextStreamer(Object):
+    """The class that streams back validated utf-8 text strings
+    that generated by tokenizer.
+    """
+
+    def __init__(self, tokenizer: Tokenizer) -> None:
+        """Create the text streamer from tokenizer"""
+        self.__init_handle_by_constructor__(
+            _ffi_api.TextStreamer,
+            tokenizer,
+        )
+
+    def put(self, delta_tokens: Union[List[int], Shape]) -> str:  # noqa: UP006
+        """Put new delta tokens into the streamer, and get the UTF-8-valid
+        delta string. The text streamer may hold some of the input delta tokens
+        which cannot decode into valid UTF-8 strings. The returned string
+        is always guaranteed to be UTF-8 valid.
+
+        Parameters
+        ----------
+        delta_tokens : Union[List[int], Shape]
+            The new tokens to put into the streamer.
+
+        Returns
+        -------
+        delta_text : str
+            The decoded delta string after putting the input new tokens.
+        """
+        if isinstance(delta_tokens, list):
+            delta_tokens = Shape(delta_tokens)
+        return _ffi_api.TextStreamerPut(self, delta_tokens)
+
+    def finish(self) -> str:
+        """Return the string decoded by remaining tokens."""
+        return _ffi_api.TextStreamerFinish(self)
+
+
+@tvm_ffi.register_object("mlc.StopStrHandler")
+class StopStrHandler(Object):
+    """The stop string handler in MLC LLM, which takes input delta tokens
+    one at a time, and return the output delta token before stopping due to
+    stop strings."""
+
+    def __init__(
+        self,
+        stop_strs: List[str],  # noqa: UP006
+        tokenizer: Tokenizer,
+    ) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.StopStrHandler,
+            stop_strs,
+            tokenizer,
+        )
+
+    def put(self, token_id: int) -> List[int]:  # noqa: UP006
+        """Add new input delta token to the handler, return output
+        delta tokens before stopping. The stop string handler may hold
+        some of the input delta token which may be part of a stop string.
+        The returned tokens are always guaranteed not to be part of stop string.
+        """
+        return list(_ffi_api.StopStrHandlerPut(self, token_id))
+
+    def finish(self) -> List[int]:  # noqa: UP006
+        """Stop string handling has finished, return remaining cached token ids."""
+        return list(_ffi_api.StopStringHandlerFinish(self))
+
+    @property
+    def stop_triggered(self) -> bool:
+        """Check if the generation has stopped due to stop string."""
+        return _ffi_api.StopStrHandlerStopTriggered(self)

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))