codingutf8 2023 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license base class from which generate streamers should inherit function that is called by generate to push new tokens raise notimplementederror def endself simple text streamer that prints the tokens to stdout as soon as entire words are formed tip warningtrue the api for the streamer classes is still under development and may change in the future tip parameters tokenizer autotokenizer the tokenized used to decode the tokens skipprompt bool optional defaults to false whether to skip the prompt to generate or not useful e g for chatbots decodekwargs dict optional additional keyword arguments to pass to the tokenizer s decode method examples python from transformers import automodelforcausallm autotokenizer textstreamer tok autotokenizer frompretrainedgpt2 model automodelforcausallm frompretrainedgpt2 inputs tokan increasing sequence one returntensorspt streamer textstreamertok despite returning the usual output the streamer will also print the generated text to stdout model generateinputs streamerstreamer maxnewtokens20 an increasing sequence one two three four five six seven eight nine ten eleven variables used in the streaming process receives tokens decodes them and prints them to stdout as soon as they form entire words add the new token to the cache and decodes the entire thing after the symbol for a new line we flush the cache if the last token is a cjk character we print the characters otherwise prints until the last space char simple heuristic to avoid printing incomplete words which may change with the subsequent token there are probably smarter ways to do this flushes any remaining cache and prints a newline to stdout flush the cache if it exists if lenself tokencache 0 text self tokenizer decodeself tokencache self decodekwargs printabletext textself printlen self tokencache self printlen 0 else printabletext self nexttokensareprompt true self onfinalizedtextprintabletext streamendtrue def onfinalizedtextself text str streamend bool false checks whether cp is the codepoint of a cjk character this defines a chinese character as anything in the cjk unicode block https en wikipedia orgwikicjkunifiedideographsunicodeblock note that the cjk unicode block is not all japanese and korean characters despite its name the modern korean hangul alphabet is a different block as is japanese hiragana and katakana those alphabets are used to write spaceseparated words so they are not treated specially and handled like the all of the other languages if cp 0x4e00 and cp 0x9fff or cp 0x3400 and cp 0x4dbf or cp 0x20000 and cp 0x2a6df or cp 0x2a700 and cp 0x2b73f or cp 0x2b740 and cp 0x2b81f or cp 0x2b820 and cp 0x2ceaf or cp 0xf900 and cp 0xfaff or cp 0x2f800 and cp 0x2fa1f return true return false class textiteratorstreamertextstreamer def init self tokenizer autotokenizer skipprompt bool false timeout optionalfloat none decodekwargs super inittokenizer skipprompt decodekwargs self textqueue queue self stopsignal none self timeout timeout def onfinalizedtextself text str streamend bool false coding utf 8 2023 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license base class from which generate streamers should inherit function that is called by generate to push new tokens function that is called by generate to signal the end of generation simple text streamer that prints the token s to stdout as soon as entire words are formed tip warning true the api for the streamer classes is still under development and may change in the future tip parameters tokenizer autotokenizer the tokenized used to decode the tokens skip_prompt bool optional defaults to false whether to skip the prompt to generate or not useful e g for chatbots decode_kwargs dict optional additional keyword arguments to pass to the tokenizer s decode method examples python from transformers import automodelforcausallm autotokenizer textstreamer tok autotokenizer from_pretrained gpt2 model automodelforcausallm from_pretrained gpt2 inputs tok an increasing sequence one return_tensors pt streamer textstreamer tok despite returning the usual output the streamer will also print the generated text to stdout _ model generate inputs streamer streamer max_new_tokens 20 an increasing sequence one two three four five six seven eight nine ten eleven variables used in the streaming process receives tokens decodes them and prints them to stdout as soon as they form entire words add the new token to the cache and decodes the entire thing after the symbol for a new line we flush the cache if the last token is a cjk character we print the characters otherwise prints until the last space char simple heuristic to avoid printing incomplete words which may change with the subsequent token there are probably smarter ways to do this flushes any remaining cache and prints a newline to stdout flush the cache if it exists prints the new text to stdout if the stream is ending also prints a newline checks whether cp is the codepoint of a cjk character this defines a chinese character as anything in the cjk unicode block https en wikipedia org wiki cjk_unified_ideographs_ unicode_block note that the cjk unicode block is not all japanese and korean characters despite its name the modern korean hangul alphabet is a different block as is japanese hiragana and katakana those alphabets are used to write space separated words so they are not treated specially and handled like the all of the other languages streamer that stores print ready text in a queue to be used by a downstream application as an iterator this is useful for applications that benefit from acessing the generated text in a non blocking way e g in an interactive gradio demo tip warning true the api for the streamer classes is still under development and may change in the future tip parameters tokenizer autotokenizer the tokenized used to decode the tokens skip_prompt bool optional defaults to false whether to skip the prompt to generate or not useful e g for chatbots timeout float optional the timeout for the text queue if none the queue will block indefinitely useful to handle exceptions in generate when it is called in a separate thread decode_kwargs dict optional additional keyword arguments to pass to the tokenizer s decode method examples python from transformers import automodelforcausallm autotokenizer textiteratorstreamer from threading import thread tok autotokenizer from_pretrained gpt2 model automodelforcausallm from_pretrained gpt2 inputs tok an increasing sequence one return_tensors pt streamer textiteratorstreamer tok run the generation in a separate thread so that we can fetch the generated text in a non blocking way generation_kwargs dict inputs streamer streamer max_new_tokens 20 thread thread target model generate kwargs generation_kwargs thread start generated_text for new_text in streamer generated_text new_text generated_text an increasing sequence one two three four five six seven eight nine ten eleven put the new text in the queue if the stream is ending also put a stop signal in the queue
from queue import Queue
from typing import TYPE_CHECKING, Optional


if TYPE_CHECKING:
    from ..models.auto import AutoTokenizer


class BaseStreamer:
    

    def put(self, value):
        
        raise NotImplementedError()

    def end(self):
        
        raise NotImplementedError()


class TextStreamer(BaseStreamer):
    

    def __init__(self, tokenizer: "AutoTokenizer", skip_prompt: bool = False, **decode_kwargs):
        self.tokenizer = tokenizer
        self.skip_prompt = skip_prompt
        self.decode_kwargs = decode_kwargs

        
        self.token_cache = []
        self.print_len = 0
        self.next_tokens_are_prompt = True

    def put(self, value):
        
        if len(value.shape) > 1 and value.shape[0] > 1:
            raise ValueError("TextStreamer only supports batch size 1")
        elif len(value.shape) > 1:
            value = value[0]

        if self.skip_prompt and self.next_tokens_are_prompt:
            self.next_tokens_are_prompt = False
            return

        
        self.token_cache.extend(value.tolist())
        text = self.tokenizer.decode(self.token_cache, **self.decode_kwargs)

        
        if text.endswith("\n"):
            printable_text = text[self.print_len :]
            self.token_cache = []
            self.print_len = 0
        
        elif len(text) > 0 and self._is_chinese_char(ord(text[-1])):
            printable_text = text[self.print_len :]
            self.print_len += len(printable_text)
        
        
        else:
            printable_text = text[self.print_len : text.rfind(" ") + 1]
            self.print_len += len(printable_text)

        self.on_finalized_text(printable_text)

    def end(self):
        
        
        if len(self.token_cache) > 0:
            text = self.tokenizer.decode(self.token_cache, **self.decode_kwargs)
            printable_text = text[self.print_len :]
            self.token_cache = []
            self.print_len = 0
        else:
            printable_text = ""

        self.next_tokens_are_prompt = True
        self.on_finalized_text(printable_text, stream_end=True)

    def on_finalized_text(self, text: str, stream_end: bool = False):
        
        print(text, flush=True, end="" if not stream_end else None)

    def _is_chinese_char(self, cp):
        
        
        
        
        
        
        
        
        
        if (
            (cp >= 0x4E00 and cp <= 0x9FFF)
            or (cp >= 0x3400 and cp <= 0x4DBF)  
            or (cp >= 0x20000 and cp <= 0x2A6DF)  
            or (cp >= 0x2A700 and cp <= 0x2B73F)  
            or (cp >= 0x2B740 and cp <= 0x2B81F)  
            or (cp >= 0x2B820 and cp <= 0x2CEAF)  
            or (cp >= 0xF900 and cp <= 0xFAFF)
            or (cp >= 0x2F800 and cp <= 0x2FA1F)  
        ):  
            return True

        return False


class TextIteratorStreamer(TextStreamer):
    

    def __init__(
        self, tokenizer: "AutoTokenizer", skip_prompt: bool = False, timeout: Optional[float] = None, **decode_kwargs
    ):
        super().__init__(tokenizer, skip_prompt, **decode_kwargs)
        self.text_queue = Queue()
        self.stop_signal = None
        self.timeout = timeout

    def on_finalized_text(self, text: str, stream_end: bool = False):
        
        self.text_queue.put(text, timeout=self.timeout)
        if stream_end:
            self.text_queue.put(self.stop_signal, timeout=self.timeout)

    def __iter__(self):
        return self

    def __next__(self):
        value = self.text_queue.get(timeout=self.timeout)
        if value == self.stop_signal:
            raise StopIteration()
        else:
            return value