chore: import upstream snapshot with attribution

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)