From 8a16edce671e28a20f5469be66c7e097fb4fce33 Mon Sep 17 00:00:00 2001 From: Andreas Abdi Date: Fri, 7 Mar 2025 03:16:09 -0800 Subject: [PATCH] Export base streamer. (#36500) * Export base streamer. Previously, the base streamer class was not exported so the set of available streamers was fixed to 3 streamer classes. This change makes it so that customers may extend the default base streamer class. * make fixup --------- Co-authored-by: Joao Gante Co-authored-by: Joao Gante --- src/transformers/generation/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/transformers/generation/__init__.py b/src/transformers/generation/__init__.py index ea39e8a10b8..cf1fa3661e0 100644 --- a/src/transformers/generation/__init__.py +++ b/src/transformers/generation/__init__.py @@ -26,7 +26,7 @@ _import_structure = { "SynthIDTextWatermarkingConfig", "WatermarkingConfig", ], - "streamers": ["AsyncTextIteratorStreamer", "TextIteratorStreamer", "TextStreamer"], + "streamers": ["AsyncTextIteratorStreamer", "BaseStreamer", "TextIteratorStreamer", "TextStreamer"], } try: @@ -197,7 +197,7 @@ if TYPE_CHECKING: SynthIDTextWatermarkingConfig, WatermarkingConfig, ) - from .streamers import AsyncTextIteratorStreamer, TextIteratorStreamer, TextStreamer + from .streamers import AsyncTextIteratorStreamer, BaseStreamer, TextIteratorStreamer, TextStreamer try: if not is_torch_available():