



import unittest
from queue import Empty
from threading import Thread

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


if is_torch_available():
      import torch

      from transformers import AutoModelForCausalLM

@require_torch
class       _SCREAMING_SNAKE_CASE	(unittest.TestCase ):





     def 	__snake_case    (       self :							List[str]				)->List[str]:
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2"				)
      __SCREAMING_SNAKE_CASE       :    int       =	AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2"				).to(UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	-1

      __SCREAMING_SNAKE_CASE       :    Any       =	ids_tensor((1, 5)     ,  vocab_size=model.config.vocab_size				).to(UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	model.generate(UpperCamelCase     ,  max_new_tokens=1_0     ,  do_sample=UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    Any       =	tokenizer.decode(greedy_ids[0]				)

      with CaptureStdout() as cs:
       __SCREAMING_SNAKE_CASE       :    int       =	TextStreamer(UpperCamelCase				)
       model.generate(UpperCamelCase     ,  max_new_tokens=1_0     ,  do_sample=UpperCamelCase     ,  streamer=UpperCamelCase				)
      # The greedy text should be printed to stdout, except for the final "\n" in the streamer
      __SCREAMING_SNAKE_CASE       :    str       =	cs.out[:-1]

      self.assertEqual(UpperCamelCase     ,  UpperCamelCase				)





     def 	__snake_case    (       self :							Dict				)->Dict:
      __SCREAMING_SNAKE_CASE       :    str       =	AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2"				)
      __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2"				).to(UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	-1

      __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	ids_tensor((1, 5)     ,  vocab_size=model.config.vocab_size				).to(UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	model.generate(UpperCamelCase     ,  max_new_tokens=1_0     ,  do_sample=UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	tokenizer.decode(greedy_ids[0]				)

      __SCREAMING_SNAKE_CASE       :    Tuple       =	TextIteratorStreamer(UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	{"input_ids": input_ids, "max_new_tokens": 1_0, "do_sample": False, "streamer": streamer}
      __SCREAMING_SNAKE_CASE       :    str       =	Thread(target=model.generate     ,  kwargs=UpperCamelCase				)
      thread.start()
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	""
      for new_text in streamer:
       streamer_text += new_text

      self.assertEqual(UpperCamelCase     ,  UpperCamelCase				)





     def 	__snake_case    (       self :							Optional[Any]				)->Dict:
      __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2"				)
      __SCREAMING_SNAKE_CASE       :    Any       =	AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2"				).to(UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    Tuple       =	-1

      __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	ids_tensor((1, 5)     ,  vocab_size=model.config.vocab_size				).to(UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	model.generate(UpperCamelCase     ,  max_new_tokens=1_0     ,  do_sample=UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    Any       =	greedy_ids[:, input_ids.shape[1] :]
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	tokenizer.decode(new_greedy_ids[0]				)

      with CaptureStdout() as cs:
       __SCREAMING_SNAKE_CASE       :    List[str]       =	TextStreamer(UpperCamelCase     ,  skip_prompt=UpperCamelCase				)
       model.generate(UpperCamelCase     ,  max_new_tokens=1_0     ,  do_sample=UpperCamelCase     ,  streamer=UpperCamelCase				)
      # The greedy text should be printed to stdout, except for the final "\n" in the streamer
      __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	cs.out[:-1]

      self.assertEqual(UpperCamelCase     ,  UpperCamelCase				)





     def 	__snake_case    (       self :							Optional[Any]				)->Dict:
      # Tests that we can pass `decode_kwargs` to the streamer to control how the tokens are decoded. Must be tested
      # with actual models -- the dummy models' tokenizers are not aligned with their models, and
      # `skip_special_tokens=True` has no effect on them
      __SCREAMING_SNAKE_CASE       :    int       =	AutoTokenizer.from_pretrained("distilgpt2"				)
      __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	AutoModelForCausalLM.from_pretrained("distilgpt2"				).to(UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	-1

      __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	torch.ones((1, 5)     ,  device=UpperCamelCase				).long() * model.config.bos_token_id
      with CaptureStdout() as cs:
       __SCREAMING_SNAKE_CASE       :    List[Any]       =	TextStreamer(UpperCamelCase     ,  skip_special_tokens=UpperCamelCase				)
       model.generate(UpperCamelCase     ,  max_new_tokens=1     ,  do_sample=UpperCamelCase     ,  streamer=UpperCamelCase				)

      # The prompt contains a special token, so the streamer should not print it. As such, the output text, when
      # re-tokenized, must only contain one token
      __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	cs.out[:-1]  # Remove the final "\n"
      __SCREAMING_SNAKE_CASE       :    Any       =	tokenizer(UpperCamelCase     ,  return_tensors="pt"				)
      self.assertEqual(streamer_text_tokenized.input_ids.shape     ,  (1, 1)				)






     def 	__snake_case    (       self :							Optional[Any]				)->Union[str, Any]:
      __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2"				)
      __SCREAMING_SNAKE_CASE       :    Tuple       =	AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2"				).to(UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    List[str]       =	-1

      __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	ids_tensor((1, 5)     ,  vocab_size=model.config.vocab_size				).to(UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    str       =	TextIteratorStreamer(UpperCamelCase     ,  timeout=0.0_0_1				)
      __SCREAMING_SNAKE_CASE       :    Tuple       =	{"input_ids": input_ids, "max_new_tokens": 1_0, "do_sample": False, "streamer": streamer}
      __SCREAMING_SNAKE_CASE       :    List[str]       =	Thread(target=model.generate     ,  kwargs=UpperCamelCase				)
      thread.start()

      # The streamer will timeout after 0.001 seconds, so an exception will be raised
      with self.assertRaises(UpperCamelCase				):
       __SCREAMING_SNAKE_CASE       :    List[str]       =	""
       for new_text in streamer:
        streamer_text += new_text











# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from ..models.auto import AutoModelForSeqaSeqLM, AutoTokenizer
from .base import PipelineTool

class       _SCREAMING_SNAKE_CASE	(UpperCamelCase ):
     lowerCAmelCase					=			"""philschmid/bart-large-cnn-samsum"""
     lowerCAmelCase					=			(
         """This is a tool that summarizes an English text. It takes an input `text` containing the text to summarize, """
         """and returns a summary of the text."""
     )
     lowerCAmelCase					=			"""summarizer"""
     lowerCAmelCase					=			AutoTokenizer
     lowerCAmelCase					=			AutoModelForSeqaSeqLM

     lowerCAmelCase					=			["""text"""]
     lowerCAmelCase					=			["""text"""]





     def 	__snake_case    (       self :							Union[str, Any]     ,  UpperCamelCase :							Tuple				)->Tuple:
      return self.pre_processor(UpperCamelCase     ,  return_tensors="pt"     ,  truncation=UpperCamelCase				)





     def 	__snake_case    (       self :							Dict     ,  UpperCamelCase :							List[Any]				)->Optional[Any]:
      return self.model.generate(**UpperCamelCase				)[0]






     def 	__snake_case    (       self :							Tuple     ,  UpperCamelCase :							Dict				)->Optional[int]:
      return self.pre_processor.decode(UpperCamelCase     ,  skip_special_tokens=UpperCamelCase     ,  clean_up_tokenization_spaces=UpperCamelCase				)






