
import unittest

from transformers import (
    MODEL_FOR_CAUSAL_LM_MAPPING,
    TF_MODEL_FOR_CAUSAL_LM_MAPPING,
    TextGenerationPipeline,
    logging,
    pipeline,
)
from transformers.testing_utils import (
    CaptureLogger,
    is_pipeline_test,
    require_accelerate,
    require_tf,
    require_torch,
    require_torch_gpu,
    require_torch_or_tf,
)

from .test_pipelines_common import ANY





@is_pipeline_test
@require_torch_or_tf
class     __UpperCamelCase (							unittest.TestCase       ):






     """simple docstring"""







     lowerCAmelCase_							   =			MODEL_FOR_CAUSAL_LM_MAPPING
     lowerCAmelCase_							   =			TF_MODEL_FOR_CAUSAL_LM_MAPPING



     @require_torch
     def  UpperCAmelCase__							(				self							:							Optional[int] ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Optional[int]					=		pipeline(task='''text-generation'''     ,						model='''sshleifer/tiny-ctrl'''     ,						framework='''pt''' )
       # Using `do_sample=False` to force deterministic output
       __SCREAMING_SNAKE_CASE      :	List[str]					=		text_generator('''This is a test'''     ,						do_sample=snake_case__ )
       self.assertEqual(
           snake_case__     ,						[
               {
                   '''generated_text''': (
                       '''This is a test ☃ ☃ segmental segmental segmental 议议eski eski flutter flutter Lacy oscope.'''
                       ''' oscope. FiliFili@@'''
                   )
               }
           ]     ,						)

       __SCREAMING_SNAKE_CASE      :	List[Any]					=		text_generator(['''This is a test''', '''This is a second test'''] )
       self.assertEqual(
           snake_case__     ,						[
               [
                   {
                       '''generated_text''': (
                           '''This is a test ☃ ☃ segmental segmental segmental 议议eski eski flutter flutter Lacy oscope.'''
                           ''' oscope. FiliFili@@'''
                       )
                   }
               ],
               [
                   {
                       '''generated_text''': (
                           '''This is a second test ☃ segmental segmental segmental 议议eski eski flutter flutter Lacy'''
                           ''' oscope. oscope. FiliFili@@'''
                       )
                   }
               ],
           ]     ,						)

       __SCREAMING_SNAKE_CASE      :	List[str]					=		text_generator('''This is a test'''     ,						do_sample=snake_case__     ,						num_return_sequences=2     ,						return_tensors=snake_case__ )
       self.assertEqual(
           snake_case__     ,						[
               {'''generated_token_ids''': ANY(snake_case__ )},
               {'''generated_token_ids''': ANY(snake_case__ )},
           ]     ,						)
       __SCREAMING_SNAKE_CASE      :	Dict					=		text_generator.model.config.eos_token_id
       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		'''<pad>'''
       __SCREAMING_SNAKE_CASE      :	Any					=		text_generator(
           ['''This is a test''', '''This is a second test''']     ,						do_sample=snake_case__     ,						num_return_sequences=2     ,						batch_size=2     ,						return_tensors=snake_case__     ,						)
       self.assertEqual(
           snake_case__     ,						[
               [
                   {'''generated_token_ids''': ANY(snake_case__ )},
                   {'''generated_token_ids''': ANY(snake_case__ )},
               ],
               [
                   {'''generated_token_ids''': ANY(snake_case__ )},
                   {'''generated_token_ids''': ANY(snake_case__ )},
               ],
           ]     ,						)



     @require_tf
     def  UpperCAmelCase__							(				self							:							str ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	int					=		pipeline(task='''text-generation'''     ,						model='''sshleifer/tiny-ctrl'''     ,						framework='''tf''' )

       # Using `do_sample=False` to force deterministic output
       __SCREAMING_SNAKE_CASE      :	Tuple					=		text_generator('''This is a test'''     ,						do_sample=snake_case__ )
       self.assertEqual(
           snake_case__     ,						[
               {
                   '''generated_text''': (
                       '''This is a test FeyFeyFey(Croatis.), s.), Cannes Cannes Cannes 閲閲Cannes Cannes Cannes 攵'''
                       ''' please,'''
                   )
               }
           ]     ,						)

       __SCREAMING_SNAKE_CASE      :	Dict					=		text_generator(['''This is a test''', '''This is a second test''']     ,						do_sample=snake_case__ )
       self.assertEqual(
           snake_case__     ,						[
               [
                   {
                       '''generated_text''': (
                           '''This is a test FeyFeyFey(Croatis.), s.), Cannes Cannes Cannes 閲閲Cannes Cannes Cannes 攵'''
                           ''' please,'''
                       )
                   }
               ],
               [
                   {
                       '''generated_text''': (
                           '''This is a second test Chieftain Chieftain prefecture prefecture prefecture Cannes Cannes'''
                           ''' Cannes 閲閲Cannes Cannes Cannes 攵 please,'''
                       )
                   }
               ],
           ]     ,						)



     def  UpperCAmelCase__							(				self							:							Any     ,						_A							:							Optional[Any]     ,						_A							:							Union[str, Any]     ,						_A							:							List[str] ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Any					=		TextGenerationPipeline(model=snake_case__     ,						tokenizer=snake_case__ )
       return text_generator, ["This is a test", "Another test"]



     def  UpperCAmelCase__							(				self							:							Any ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	int					=		'''Hello I believe in'''
       __SCREAMING_SNAKE_CASE      :	List[Any]					=		pipeline('''text-generation'''     ,						model='''hf-internal-testing/tiny-random-gpt2''' )
       __SCREAMING_SNAKE_CASE      :	List[Any]					=		text_generator(snake_case__ )
       self.assertEqual(
           snake_case__     ,						[{'''generated_text''': '''Hello I believe in fe fe fe fe fe fe fe fe fe fe fe fe'''}]     ,						)

       __SCREAMING_SNAKE_CASE      :	str					=		text_generator(snake_case__     ,						stop_sequence=''' fe''' )
       self.assertEqual(snake_case__     ,						[{'''generated_text''': '''Hello I believe in fe'''}] )



     def  UpperCAmelCase__							(				self							:							str     ,						_A							:							str     ,						_A							:							Tuple ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Dict					=		text_generator.model
       __SCREAMING_SNAKE_CASE      :	int					=		text_generator.tokenizer

       __SCREAMING_SNAKE_CASE      :	Tuple					=		text_generator('''This is a test''' )
       self.assertEqual(snake_case__     ,						[{'''generated_text''': ANY(snake_case__ )}] )
       self.assertTrue(outputs[0]['''generated_text'''].startswith('''This is a test''' ) )

       __SCREAMING_SNAKE_CASE      :	Tuple					=		text_generator('''This is a test'''     ,						return_full_text=snake_case__ )
       self.assertEqual(snake_case__     ,						[{'''generated_text''': ANY(snake_case__ )}] )
       self.assertNotIn('''This is a test'''     ,						outputs[0]['''generated_text'''] )

       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		pipeline(task='''text-generation'''     ,						model=snake_case__     ,						tokenizer=snake_case__     ,						return_full_text=snake_case__ )
       __SCREAMING_SNAKE_CASE      :	List[Any]					=		text_generator('''This is a test''' )
       self.assertEqual(snake_case__     ,						[{'''generated_text''': ANY(snake_case__ )}] )
       self.assertNotIn('''This is a test'''     ,						outputs[0]['''generated_text'''] )

       __SCREAMING_SNAKE_CASE      :	str					=		text_generator('''This is a test'''     ,						return_full_text=snake_case__ )
       self.assertEqual(snake_case__     ,						[{'''generated_text''': ANY(snake_case__ )}] )
       self.assertTrue(outputs[0]['''generated_text'''].startswith('''This is a test''' ) )

       __SCREAMING_SNAKE_CASE      :	Any					=		text_generator(['''This is great !''', '''Something else''']     ,						num_return_sequences=2     ,						do_sample=snake_case__ )
       self.assertEqual(
           snake_case__     ,						[
               [{'''generated_text''': ANY(snake_case__ )}, {'''generated_text''': ANY(snake_case__ )}],
               [{'''generated_text''': ANY(snake_case__ )}, {'''generated_text''': ANY(snake_case__ )}],
           ]     ,						)

       if text_generator.tokenizer.pad_token is not None:
         __SCREAMING_SNAKE_CASE      :	str					=		text_generator(
             ['''This is great !''', '''Something else''']     ,						num_return_sequences=2     ,						batch_size=2     ,						do_sample=snake_case__ )
         self.assertEqual(
             snake_case__     ,						[
                 [{'''generated_text''': ANY(snake_case__ )}, {'''generated_text''': ANY(snake_case__ )}],
                 [{'''generated_text''': ANY(snake_case__ )}, {'''generated_text''': ANY(snake_case__ )}],
             ]     ,						)

       with self.assertRaises(snake_case__ ):
         __SCREAMING_SNAKE_CASE      :	Tuple					=		text_generator('''test'''     ,						return_full_text=snake_case__     ,						return_text=snake_case__ )
       with self.assertRaises(snake_case__ ):
         __SCREAMING_SNAKE_CASE      :	Dict					=		text_generator('''test'''     ,						return_full_text=snake_case__     ,						return_tensors=snake_case__ )
       with self.assertRaises(snake_case__ ):
         __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		text_generator('''test'''     ,						return_text=snake_case__     ,						return_tensors=snake_case__ )

       # Empty prompt is slighly special
       # it requires BOS token to exist.
       # Special case for Pegasus which will always append EOS so will
       # work even without BOS.
       if (
           text_generator.tokenizer.bos_token_id is not None
           or "Pegasus" in tokenizer.__class__.__name__
           or "Git" in model.__class__.__name__
       ):
         __SCREAMING_SNAKE_CASE      :	List[str]					=		text_generator('''''' )
         self.assertEqual(snake_case__     ,						[{'''generated_text''': ANY(snake_case__ )}] )
       else:
         with self.assertRaises((ValueError, AssertionError) ):
           __SCREAMING_SNAKE_CASE      :	List[str]					=		text_generator('''''' )

       if text_generator.framework == "tf":
         # TF generation does not support max_new_tokens, and it's impossible
         # to control long generation with only max_length without
         # fancy calculation, dismissing tests for now.
         return
       # We don't care about infinite range models.
       # They already work.
       # Skip this test for XGLM, since it uses sinusoidal positional embeddings which are resized on-the-fly.
       __SCREAMING_SNAKE_CASE      :	Dict					=		['''RwkvForCausalLM''', '''XGLMForCausalLM''', '''GPTNeoXForCausalLM''']
       if (
           tokenizer.model_max_length < 1_0000
           and text_generator.model.__class__.__name__ not in EXTRA_MODELS_CAN_HANDLE_LONG_INPUTS
       ):
         # Handling of large generations
         with self.assertRaises((RuntimeError, IndexError, ValueError, AssertionError) ):
           text_generator('''This is a test''' * 500     ,						max_new_tokens=20 )

         __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		text_generator('''This is a test''' * 500     ,						handle_long_generation='''hole'''     ,						max_new_tokens=20 )
         # Hole strategy cannot work
         with self.assertRaises(snake_case__ ):
           text_generator(
               '''This is a test''' * 500     ,						handle_long_generation='''hole'''     ,						max_new_tokens=tokenizer.model_max_length + 10     ,						)



     @require_torch
     @require_accelerate
     @require_torch_gpu
     def  UpperCAmelCase__							(				self							:							List[str] ):







       """simple docstring"""




       import torch

       # Classic `model_kwargs`
       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		pipeline(
           model='''hf-internal-testing/tiny-random-bloom'''     ,						model_kwargs={'''device_map''': '''auto''', '''torch_dtype''': torch.bfloataa}     ,						)
       self.assertEqual(pipe.model.device     ,						torch.device(0 ) )
       self.assertEqual(pipe.model.lm_head.weight.dtype     ,						torch.bfloataa )
       __SCREAMING_SNAKE_CASE      :	int					=		pipe('''This is a test''' )
       self.assertEqual(
           snake_case__     ,						[
               {
                   '''generated_text''': (
                       '''This is a test test test test test test test test test test test test test test test test'''
                       ''' test'''
                   )
               }
           ]     ,						)

       # Upgraded those two to real pipeline arguments (they just get sent for the model as they're unlikely to mean anything else.)
       __SCREAMING_SNAKE_CASE      :	int					=		pipeline(model='''hf-internal-testing/tiny-random-bloom'''     ,						device_map='''auto'''     ,						torch_dtype=torch.bfloataa )
       self.assertEqual(pipe.model.device     ,						torch.device(0 ) )
       self.assertEqual(pipe.model.lm_head.weight.dtype     ,						torch.bfloataa )
       __SCREAMING_SNAKE_CASE      :	Dict					=		pipe('''This is a test''' )
       self.assertEqual(
           snake_case__     ,						[
               {
                   '''generated_text''': (
                       '''This is a test test test test test test test test test test test test test test test test'''
                       ''' test'''
                   )
               }
           ]     ,						)

       # torch_dtype will be automatically set to float32 if not provided - check: https://github.com/huggingface/transformers/pull/20602
       __SCREAMING_SNAKE_CASE      :	str					=		pipeline(model='''hf-internal-testing/tiny-random-bloom'''     ,						device_map='''auto''' )
       self.assertEqual(pipe.model.device     ,						torch.device(0 ) )
       self.assertEqual(pipe.model.lm_head.weight.dtype     ,						torch.floataa )
       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		pipe('''This is a test''' )
       self.assertEqual(
           snake_case__     ,						[
               {
                   '''generated_text''': (
                       '''This is a test test test test test test test test test test test test test test test test'''
                       ''' test'''
                   )
               }
           ]     ,						)



     @require_torch
     @require_torch_gpu
     def  UpperCAmelCase__							(				self							:							Optional[Any] ):







       """simple docstring"""




       import torch

       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		pipeline(model='''hf-internal-testing/tiny-random-bloom'''     ,						device=0     ,						torch_dtype=torch.floataa )
       pipe('''This is a test''' )



     @require_torch
     @require_accelerate
     @require_torch_gpu
     def  UpperCAmelCase__							(				self							:							List[str] ):







       """simple docstring"""




       import torch

       __SCREAMING_SNAKE_CASE      :	List[str]					=		pipeline(model='''hf-internal-testing/tiny-random-bloom'''     ,						device_map='''auto'''     ,						torch_dtype=torch.floataa )
       pipe('''This is a test'''     ,						do_sample=snake_case__     ,						top_p=0.5 )




     def  UpperCAmelCase__							(				self							:							str ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		'''Hello world'''
       __SCREAMING_SNAKE_CASE      :	Dict					=		pipeline('''text-generation'''     ,						model='''hf-internal-testing/tiny-random-gpt2''' )
       if text_generator.model.framework == "tf":
         __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		logging.get_logger('''transformers.generation.tf_utils''' )
       else:
         __SCREAMING_SNAKE_CASE      :	Optional[int]					=		logging.get_logger('''transformers.generation.utils''' )
       __SCREAMING_SNAKE_CASE      :	Dict					=		'''Both `max_new_tokens`'''  # The beggining of the message to be checked in this test

       # Both are set by the user -> log warning
       with CaptureLogger(snake_case__ ) as cl:
         __SCREAMING_SNAKE_CASE      :	Any					=		text_generator(snake_case__     ,						max_length=10     ,						max_new_tokens=1 )
       self.assertIn(snake_case__     ,						cl.out )

       # The user only sets one -> no warning
       with CaptureLogger(snake_case__ ) as cl:
         __SCREAMING_SNAKE_CASE      :	List[str]					=		text_generator(snake_case__     ,						max_new_tokens=1 )
       self.assertNotIn(snake_case__     ,						cl.out )

       with CaptureLogger(snake_case__ ) as cl:
         __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		text_generator(snake_case__     ,						max_length=10 )
       self.assertNotIn(snake_case__     ,						cl.out )

from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import convert_to_rgb, normalize, rescale, resize, to_channel_dimension_format
from ...image_utils import (
    OPENAI_CLIP_MEAN,
    OPENAI_CLIP_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
				import PIL


A						=			logging.get_logger(__name__)





class     lowerCAmelCase__     (				SCREAMING_SNAKE_CASE_			):


		'''simple docstring'''






		lowerCAmelCase_     =   ['pixel_values']
		def __init__(       self     :     List[str]     ,  snake_case__     :     bool = True     ,  snake_case__     :     Dict[str, int] = None     ,  snake_case__     :     PILImageResampling = PILImageResampling.BICUBIC     ,  snake_case__     :     bool = True     ,  snake_case__     :     Union[int, float] = 1 / 2_5_5     ,  snake_case__     :     bool = True     ,  snake_case__     :     Optional[Union[float, List[float]]] = None     ,  snake_case__     :     Optional[Union[float, List[float]]] = None     ,  snake_case__     :     bool = True     ,  **snake_case__     :     List[Any]     ,  )      ->	None:
				super().__init__(**snake_case__		)
				_lowerCamelCase				   =  size if size is not None else {'height': 3_8_4, 'width': 3_8_4}
				_lowerCamelCase				   =  get_size_dict(snake_case__     ,  default_to_square=snake_case__		)

				_lowerCamelCase				   =  do_resize
				_lowerCamelCase				   =  size
				_lowerCamelCase				   =  resample
				_lowerCamelCase				   =  do_rescale
				_lowerCamelCase				   =  rescale_factor
				_lowerCamelCase				   =  do_normalize
				_lowerCamelCase				   =  image_mean if image_mean is not None else OPENAI_CLIP_MEAN
				_lowerCamelCase				   =  image_std if image_std is not None else OPENAI_CLIP_STD
				_lowerCamelCase				   =  do_convert_rgb
		def   _snake_case       (       self     :     Union[str, Any]     ,  snake_case__     :     np.ndarray     ,  snake_case__     :     Dict[str, int]     ,  snake_case__     :     PILImageResampling = PILImageResampling.BICUBIC     ,  snake_case__     :     Optional[Union[str, ChannelDimension]] = None     ,  **snake_case__     :     Dict     ,  )      ->	np.ndarray:
				_lowerCamelCase				   =  get_size_dict(snake_case__     ,  default_to_square=snake_case__		)
				if "height" not in size or "width" not in size:
						raise ValueError(f"""The `size` dictionary must contain the keys `height` and `width`. Got {size.keys()}"""		)
				_lowerCamelCase				   =  (size['height'], size['width'])
				return resize(snake_case__     ,  size=snake_case__     ,  resample=snake_case__     ,  data_format=snake_case__     ,  **snake_case__		)
		def   _snake_case       (       self     :     Dict     ,  snake_case__     :     np.ndarray     ,  snake_case__     :     Union[int, float]     ,  snake_case__     :     Optional[Union[str, ChannelDimension]] = None     ,  **snake_case__     :     Optional[int]     ,  )      ->	int:
				return rescale(snake_case__     ,  scale=snake_case__     ,  data_format=snake_case__     ,  **snake_case__		)
		def   _snake_case       (       self     :     Optional[Any]     ,  snake_case__     :     np.ndarray     ,  snake_case__     :     Union[float, List[float]]     ,  snake_case__     :     Union[float, List[float]]     ,  snake_case__     :     Optional[Union[str, ChannelDimension]] = None     ,  **snake_case__     :     Optional[int]     ,  )      ->	np.ndarray:
				return normalize(snake_case__     ,  mean=snake_case__     ,  std=snake_case__     ,  data_format=snake_case__     ,  **snake_case__		)



		def   _snake_case       (       self     :     List[str]     ,  snake_case__     :     ImageInput     ,  snake_case__     :     Optional[bool] = None     ,  snake_case__     :     Optional[Dict[str, int]] = None     ,  snake_case__     :     PILImageResampling = None     ,  snake_case__     :     Optional[bool] = None     ,  snake_case__     :     Optional[float] = None     ,  snake_case__     :     Optional[bool] = None     ,  snake_case__     :     Optional[Union[float, List[float]]] = None     ,  snake_case__     :     Optional[Union[float, List[float]]] = None     ,  snake_case__     :     Optional[Union[str, TensorType]] = None     ,  snake_case__     :     bool = None     ,  snake_case__     :     ChannelDimension = ChannelDimension.FIRST     ,  **snake_case__     :     str     ,  )      ->	PIL.Image.Image:
				_lowerCamelCase				   =  do_resize if do_resize is not None else self.do_resize
				_lowerCamelCase				   =  resample if resample is not None else self.resample
				_lowerCamelCase				   =  do_rescale if do_rescale is not None else self.do_rescale
				_lowerCamelCase				   =  rescale_factor if rescale_factor is not None else self.rescale_factor
				_lowerCamelCase				   =  do_normalize if do_normalize is not None else self.do_normalize
				_lowerCamelCase				   =  image_mean if image_mean is not None else self.image_mean
				_lowerCamelCase				   =  image_std if image_std is not None else self.image_std
				_lowerCamelCase				   =  do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

				_lowerCamelCase				   =  size if size is not None else self.size
				_lowerCamelCase				   =  get_size_dict(snake_case__     ,  default_to_square=snake_case__		)

				_lowerCamelCase				   =  make_list_of_images(snake_case__		)

				if not valid_images(snake_case__		):
						raise ValueError(
						    'Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '
						    'torch.Tensor, tf.Tensor or jax.ndarray.'		)

				if do_resize and size is None or resample is None:
						raise ValueError('Size and resample must be specified if do_resize is True.'		)

				if do_rescale and rescale_factor is None:
						raise ValueError('Rescale factor must be specified if do_rescale is True.'		)

				if do_normalize and (image_mean is None or image_std is None):
						raise ValueError('Image mean and std must be specified if do_normalize is True.'		)

				# PIL RGBA images are converted to RGB
				if do_convert_rgb:
						_lowerCamelCase				   =  [convert_to_rgb(snake_case__		) for image in images]

				# All transformations expect numpy arrays.
				_lowerCamelCase				   =  [to_numpy_array(snake_case__		) for image in images]

				if do_resize:
						_lowerCamelCase				   =  [self.resize(image=snake_case__     ,  size=snake_case__     ,  resample=snake_case__		) for image in images]

				if do_rescale:
						_lowerCamelCase				   =  [self.rescale(image=snake_case__     ,  scale=snake_case__		) for image in images]

				if do_normalize:
						_lowerCamelCase				   =  [self.normalize(image=snake_case__     ,  mean=snake_case__     ,  std=snake_case__		) for image in images]

				_lowerCamelCase				   =  [to_channel_dimension_format(snake_case__     ,  snake_case__		) for image in images]

				_lowerCamelCase				   =  BatchFeature(data={'pixel_values': images}     ,  tensor_type=snake_case__		)

				return encoded_outputs