Generate: TextIteratorStreamer (streamer for gradio) (#22501)

* haha text go brrr (but in gradio)

tests/generation/test_streamers.py

@@ -14,8 +14,9 @@
 # limitations under the License.
 
 import unittest
+from threading import Thread
 
-from transformers import AutoTokenizer, TextStreamer, is_torch_available
+from transformers import AutoTokenizer, TextIteratorStreamer, TextStreamer, is_torch_available
 from transformers.testing_utils import CaptureStdout, require_torch, torch_device
 
 from ..test_modeling_common import ids_tensor
@@ -27,7 +28,7 @@ if is_torch_available():
 
 @require_torch
 class StreamerTester(unittest.TestCase):
-    def test_text_streamer_stdout(self):
+    def test_text_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
@@ -39,6 +40,26 @@ class StreamerTester(unittest.TestCase):
         with CaptureStdout() as cs:
             streamer = TextStreamer(tokenizer)
             model.generate(input_ids, max_new_tokens=10, do_sample=False, streamer=streamer)
-
         # The greedy text should be printed to stdout, except for the final "\n" in the streamer
-        self.assertEqual(cs.out[:-1], greedy_text)
+        streamer_text = cs.out[:-1]
+
+        self.assertEqual(streamer_text, greedy_text)
+
+    def test_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 = TextIteratorStreamer(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 = ""
+        for new_text in streamer:
+            streamer_text += new_text
+
+        self.assertEqual(streamer_text, greedy_text)