SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Implement AsyncTextIteratorStreamer for asynchronous streaming (#34931)

Browse Source 
* Add AsyncTextIteratorStreamer class

* export AsyncTextIteratorStreamer

* export AsyncTextIteratorStreamer

* improve docs

* missing import

* missing import

* doc example fix

* doc example output fix

* add pytest-asyncio

* first attempt at tests

* missing import

* add pytest-asyncio

* fallback to wait_for and raise TimeoutError on timeout

* check for TimeoutError

* autodoc

* reorder imports

* fix style

---------

Co-authored-by: Arthur Zucker <arthur.zucker@gmail.com>
Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>
This commit is contained in:
Sigbjørn Skjæret
2024-12-20 12:08:12 +01:00
committed by GitHub
parent b5a557e5fe
commit eafbb0eca7
7 changed files with 154 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
tests/generation/test_streamers.py
View File

@@ -17,7 +17,15 @@ import unittest
from queue import Empty
from threading import Thread
from transformers import AutoTokenizer, TextIteratorStreamer, TextStreamer, is_torch_available
import pytest
from transformers import (
AsyncTextIteratorStreamer,
AutoTokenizer,
TextIteratorStreamer,
TextStreamer,
is_torch_available,
)
from transformers.testing_utils import CaptureStdout, require_torch, torch_device
from ..test_modeling_common import ids_tensor
@@ -120,3 +128,43 @@ class StreamerTester(unittest.TestCase):
streamer_text = ""
for new_text in streamer:
streamer_text += new_text
@require_torch
@pytest.mark.asyncio(loop_scope="class")
class AsyncStreamerTester(unittest.IsolatedAsyncioTestCase):
async def test_async_iterator_streamer_matches_non_streaming(self):
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(torch_device)
model.config.eos_token_id = -1
input_ids = ids_tensor((1, 5), vocab_size=model.config.vocab_size).to(torch_device)
greedy_ids = model.generate(input_ids, max_new_tokens=10, do_sample=False)
greedy_text = tokenizer.decode(greedy_ids[0])
streamer = AsyncTextIteratorStreamer(tokenizer)
generation_kwargs = {"input_ids": input_ids, "max_new_tokens": 10, "do_sample": False, "streamer": streamer}
thread = Thread(target=model.generate, kwargs=generation_kwargs)
thread.start()
streamer_text = ""
async for new_text in streamer:
streamer_text += new_text
self.assertEqual(streamer_text, greedy_text)
async def test_async_iterator_streamer_timeout(self):
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(torch_device)
model.config.eos_token_id = -1
input_ids = ids_tensor((1, 5), vocab_size=model.config.vocab_size).to(torch_device)
streamer = AsyncTextIteratorStreamer(tokenizer, timeout=0.001)
generation_kwargs = {"input_ids": input_ids, "max_new_tokens": 10, "do_sample": False, "streamer": streamer}
thread = Thread(target=model.generate, kwargs=generation_kwargs)
thread.start()
# The streamer will timeout after 0.001 seconds, so TimeoutError will be raised
with self.assertRaises(TimeoutError):
streamer_text = ""
async for new_text in streamer:
streamer_text += new_text