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