saymrwulf/transformers
1
0
Fork 0
mirror of https://github.com/saymrwulf/transformers.git synced 2026-05-15 21:01:19 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Generate: Enable easier TextStreamer customization (#22516)

Browse source
This commit is contained in:
Vladimir Blagojevic 2023-04-03 19:49:38 +02:00 committed by GitHub
parent 80d1319e1b
commit a17841ac49
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 6 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
src/transformers/generation/streamers.py
View file

@ -88,7 +88,7 @@ class TextStreamer(BaseStreamer):
printable_text = text[self.print_len : text.rfind(" ") + 1]
self.print_len += len(printable_text)
print(printable_text, flush=True, end="")
self.on_finalized_text(printable_text)
def end(self):
"""Flushes any remaining cache and prints a newline to stdout."""
@ -102,7 +102,11 @@ class TextStreamer(BaseStreamer):
printable_text = ""
# Print a newline (and the remaining text, if any)
print(printable_text, flush=True)
self.on_finalized_text(printable_text, stream_end=True)
def on_finalized_text(self, token: str, stream_end: bool = False):
"""Prints the new text to stdout."""
print(token, flush=True, end="" if not stream_end else None)
class TextIteratorStreamer(BaseStreamer):