wehub-resource-sync
84 lines
2.8 KiB
Python
84 lines
2.8 KiB
Python
"""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)
|