|
from transformers import AutoTokenizer |
|
|
|
|
|
class BaseStreamer: |
|
""" |
|
Base class from which `.generate()` streamers should inherit. |
|
""" |
|
|
|
def put(self, value): |
|
"""Function that is called by `.generate()` to push new tokens""" |
|
raise NotImplementedError() |
|
|
|
def end(self): |
|
"""Function that is called by `.generate()` to signal the end of generation""" |
|
raise NotImplementedError() |
|
|
|
|
|
class ByteStreamer(BaseStreamer): |
|
""" |
|
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, |
|
``` |
|
""" |
|
|
|
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): |
|
""" |
|
Receives tokens, decodes them, and prints them to stdout as soon as they form entire words. |
|
""" |
|
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 |
|
else: |
|
printable_text = text[self.print_len : self.print_len + 1] |
|
self.print_len += len(printable_text) |
|
|
|
self.on_finalized_text(printable_text) |
|
|
|
def end(self): |
|
"""Flushes any remaining cache and prints a newline to stdout.""" |
|
|
|
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): |
|
"""Prints the new text to stdout. If the stream is ending, also prints a newline.""" |
|
print(text, flush=True, end="" if not stream_end else None) |
|
|
