'''simple docstring'''




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 	_lowercase(  unittest.TestCase					):



   """simple docstring"""





   __lowerCamelCase     =	MODEL_FOR_CAUSAL_LM_MAPPING
   __lowerCamelCase     =	TF_MODEL_FOR_CAUSAL_LM_MAPPING





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

    __UpperCAmelCase								=       text_generator(['This is a test', 'This is a second test']						)
    self.assertEqual(
        _lowercase      ,[
            [
                {
                    '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@@'
                    )
                }
            ],
        ]      ,)

    __UpperCAmelCase								=       text_generator('This is a test'      ,do_sample=_lowercase      ,num_return_sequences=2      ,return_tensors=_lowercase						)
    self.assertEqual(
        _lowercase      ,[
            {'generated_token_ids': ANY(_lowercase						)},
            {'generated_token_ids': ANY(_lowercase						)},
        ]      ,)
    __UpperCAmelCase								=       text_generator.model.config.eos_token_id
    __UpperCAmelCase								=       "<pad>"
    __UpperCAmelCase								=       text_generator(
        ['This is a test', 'This is a second test']      ,do_sample=_lowercase      ,num_return_sequences=2      ,batch_size=2      ,return_tensors=_lowercase      ,)
    self.assertEqual(
        _lowercase      ,[
            [
                {'generated_token_ids': ANY(_lowercase						)},
                {'generated_token_ids': ANY(_lowercase						)},
            ],
            [
                {'generated_token_ids': ANY(_lowercase						)},
                {'generated_token_ids': ANY(_lowercase						)},
            ],
        ]      ,)





   @require_tf
   def       snake_case						(				self:     List[str]						):
    __UpperCAmelCase								=       pipeline(task='text-generation'      ,model='sshleifer/tiny-ctrl'      ,framework='tf'						)

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

    __UpperCAmelCase								=       text_generator(['This is a test', 'This is a second test']      ,do_sample=_lowercase						)
    self.assertEqual(
        _lowercase      ,[
            [
                {
                    '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       snake_case						(				self:     int      ,a:     Tuple      ,a:     int      ,a:     int						):
    __UpperCAmelCase								=       TextGenerationPipeline(model=_lowercase      ,tokenizer=_lowercase						)
    return text_generator, ["This is a test", "Another test"]





   def       snake_case						(				self:     Any						):
    __UpperCAmelCase								=       "Hello I believe in"
    __UpperCAmelCase								=       pipeline('text-generation'      ,model='hf-internal-testing/tiny-random-gpt2'						)
    __UpperCAmelCase								=       text_generator(_lowercase						)
    self.assertEqual(
        _lowercase      ,[{'generated_text': 'Hello I believe in fe fe fe fe fe fe fe fe fe fe fe fe'}]      ,)

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





   def       snake_case						(				self:     Optional[Any]      ,a:     Union[str, Any]      ,a:     int						):
    __UpperCAmelCase								=       text_generator.model
    __UpperCAmelCase								=       text_generator.tokenizer

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

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

    __UpperCAmelCase								=       pipeline(task='text-generation'      ,model=_lowercase      ,tokenizer=_lowercase      ,return_full_text=_lowercase						)
    __UpperCAmelCase								=       text_generator('This is a test'						)
    self.assertEqual(_lowercase      ,[{'generated_text': ANY(_lowercase						)}]						)
    self.assertNotIn('This is a test'      ,outputs[0]['generated_text']						)

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

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

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

    with self.assertRaises(_lowercase						):
     __UpperCAmelCase								=       text_generator('test'      ,return_full_text=_lowercase      ,return_text=_lowercase						)
    with self.assertRaises(_lowercase						):
     __UpperCAmelCase								=       text_generator('test'      ,return_full_text=_lowercase      ,return_tensors=_lowercase						)
    with self.assertRaises(_lowercase						):
     __UpperCAmelCase								=       text_generator('test'      ,return_text=_lowercase      ,return_tensors=_lowercase						)

    # 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__
    ):
     __UpperCAmelCase								=       text_generator(''						)
     self.assertEqual(_lowercase      ,[{'generated_text': ANY(_lowercase						)}]						)
    else:
     with self.assertRaises((ValueError, AssertionError)						):
      __UpperCAmelCase								=       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.
    __UpperCAmelCase								=       ["RwkvForCausalLM", "XGLMForCausalLM", "GPTNeoXForCausalLM"]
    if (
        tokenizer.model_max_length < 10000
        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						)

     __UpperCAmelCase								=       text_generator('This is a test' * 500      ,handle_long_generation='hole'      ,max_new_tokens=20						)
     # Hole strategy cannot work
     with self.assertRaises(_lowercase						):
      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       snake_case						(				self:     Any						):
    import torch

    # Classic `model_kwargs`
    __UpperCAmelCase								=       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						)
    __UpperCAmelCase								=       pipe('This is a test'						)
    self.assertEqual(
        _lowercase      ,[
            {
                '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.)
    __UpperCAmelCase								=       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						)
    __UpperCAmelCase								=       pipe('This is a test'						)
    self.assertEqual(
        _lowercase      ,[
            {
                '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
    __UpperCAmelCase								=       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						)
    __UpperCAmelCase								=       pipe('This is a test'						)
    self.assertEqual(
        _lowercase      ,[
            {
                '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       snake_case						(				self:     List[str]						):
    import torch

    __UpperCAmelCase								=       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       snake_case						(				self:     Optional[Any]						):
    import torch

    __UpperCAmelCase								=       pipeline(model='hf-internal-testing/tiny-random-bloom'      ,device_map='auto'      ,torch_dtype=torch.floataa						)
    pipe('This is a test'      ,do_sample=_lowercase      ,top_p=0.5						)





   def       snake_case						(				self:     Optional[Any]						):
    __UpperCAmelCase								=       "Hello world"
    __UpperCAmelCase								=       pipeline('text-generation'      ,model='hf-internal-testing/tiny-random-gpt2'						)
    if text_generator.model.framework == "tf":
     __UpperCAmelCase								=       logging.get_logger('transformers.generation.tf_utils'						)
    else:
     __UpperCAmelCase								=       logging.get_logger('transformers.generation.utils'						)
    __UpperCAmelCase								=       "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(_lowercase						) as cl:
     __UpperCAmelCase								=       text_generator(_lowercase      ,max_length=10      ,max_new_tokens=1						)
    self.assertIn(_lowercase      ,cl.out						)

    # The user only sets one -> no warning
    with CaptureLogger(_lowercase						) as cl:
     __UpperCAmelCase								=       text_generator(_lowercase      ,max_new_tokens=1						)
    self.assertNotIn(_lowercase      ,cl.out						)

    with CaptureLogger(_lowercase						) as cl:
     __UpperCAmelCase								=       text_generator(_lowercase      ,max_length=10						)
    self.assertNotIn(_lowercase      ,cl.out						)








"""simple docstring"""







import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


A		=						False

class       _a  ( unittest.TestCase):



			def 				__lowercase					(			self							:				str		)							->					Optional[int]:
									# clean up the VRAM after each test
									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()



			@property
			def 				__lowercase					(			self							:				List[str]		)							->					Any:
									return 12



			@property
			def 				__lowercase					(			self							:				List[str]		)							->					Dict:
									return 12



			@property
			def 				__lowercase					(			self							:				Optional[int]		)							->					Union[str, Any]:
									return 32



			@property
			def 				__lowercase					(			self							:				int		)							->					Dict:
									torch.manual_seed(0		)
									snake_case   :			List[str]						    = VQModel(
									    block_out_channels=[32, 64]   ,	in_channels=3   ,	out_channels=3   ,	down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]   ,	up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]   ,	latent_channels=3   ,	num_vq_embeddings=self.num_embed   ,	vq_embed_dim=3   ,	)
									return model



			@property
			def 				__lowercase					(			self							:				str		)							->					List[Any]:
									snake_case   :			Tuple						    = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		)
									return tokenizer



			@property
			def 				__lowercase					(			self							:				Union[str, Any]		)							->					Any:
									torch.manual_seed(0		)
									snake_case   :			Optional[int]						    = CLIPTextConfig(
									    bos_token_id=0   ,	eos_token_id=2   ,	hidden_size=self.text_embedder_hidden_size   ,	intermediate_size=37   ,	layer_norm_eps=1E-05   ,	num_attention_heads=4   ,	num_hidden_layers=5   ,	pad_token_id=1   ,	vocab_size=1000   ,	)
									return CLIPTextModel(_lowercase		)



			@property
			def 				__lowercase					(			self							:				Optional[int]		)							->					Optional[int]:
									torch.manual_seed(0		)

									snake_case   :			List[Any]						    = 12
									snake_case   :			Dict						    = 12

									snake_case   :			Tuple						    = {
									    "attention_bias": True,
									    "cross_attention_dim": 32,
									    "attention_head_dim": height * width,
									    "num_attention_heads": 1,
									    "num_vector_embeds": self.num_embed,
									    "num_embeds_ada_norm": self.num_embeds_ada_norm,
									    "norm_num_groups": 32,
									    "sample_size": width,
									    "activation_fn": "geglu-approximate",
									}

									snake_case   :			Dict						    = TransformeraDModel(**_lowercase		)
									return model



			def 				__lowercase					(			self							:				Optional[int]		)							->					Tuple:
									snake_case   :			Optional[Any]						    = "cpu"

									snake_case   :			Optional[int]						    = self.dummy_vqvae
									snake_case   :			Dict						    = self.dummy_text_encoder
									snake_case   :			Tuple						    = self.dummy_tokenizer
									snake_case   :			List[Any]						    = self.dummy_transformer
									snake_case   :			List[Any]						    = VQDiffusionScheduler(self.num_embed		)
									snake_case   :			List[Any]						    = LearnedClassifierFreeSamplingEmbeddings(learnable=_lowercase		)

									snake_case   :			Dict						    = VQDiffusionPipeline(
									    vqvae=_lowercase   ,	text_encoder=_lowercase   ,	tokenizer=_lowercase   ,	transformer=_lowercase   ,	scheduler=_lowercase   ,	learned_classifier_free_sampling_embeddings=_lowercase   ,	)
									snake_case   :			Any						    = pipe.to(_lowercase		)
									pipe.set_progress_bar_config(disable=_lowercase		)

									snake_case   :			Optional[Any]						    = "teddy bear playing in the pool"

									snake_case   :			Tuple						    = torch.Generator(device=_lowercase		).manual_seed(0		)
									snake_case   :			Tuple						    = pipe([prompt]   ,	generator=_lowercase   ,	num_inference_steps=2   ,	output_type="np"		)
									snake_case   :			Optional[int]						    = output.images

									snake_case   :			int						    = torch.Generator(device=_lowercase		).manual_seed(0		)
									snake_case   :			List[str]						    = pipe(
									    [prompt]   ,	generator=_lowercase   ,	output_type="np"   ,	return_dict=_lowercase   ,	num_inference_steps=2		)[0]

									snake_case   :			List[Any]						    = image[0, -3:, -3:, -1]
									snake_case   :			Tuple						    = image_from_tuple[0, -3:, -3:, -1]

									assert image.shape == (1, 24, 24, 3)

									snake_case   :			str						    = np.array([0.6551, 0.6168, 0.5008, 0.5676, 0.5659, 0.4295, 0.6073, 0.5599, 0.4992]		)

									assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-2
									assert np.abs(image_from_tuple_slice.flatten() - expected_slice		).max() < 1E-2



			def 				__lowercase					(			self							:				Union[str, Any]		)							->					Optional[int]:
									snake_case   :			List[str]						    = "cpu"

									snake_case   :			Dict						    = self.dummy_vqvae
									snake_case   :			List[Any]						    = self.dummy_text_encoder
									snake_case   :			Optional[Any]						    = self.dummy_tokenizer
									snake_case   :			int						    = self.dummy_transformer
									snake_case   :			str						    = VQDiffusionScheduler(self.num_embed		)
									snake_case   :			Optional[int]						    = LearnedClassifierFreeSamplingEmbeddings(
									    learnable=_lowercase   ,	hidden_size=self.text_embedder_hidden_size   ,	length=tokenizer.model_max_length		)

									snake_case   :			List[str]						    = VQDiffusionPipeline(
									    vqvae=_lowercase   ,	text_encoder=_lowercase   ,	tokenizer=_lowercase   ,	transformer=_lowercase   ,	scheduler=_lowercase   ,	learned_classifier_free_sampling_embeddings=_lowercase   ,	)
									snake_case   :			Optional[Any]						    = pipe.to(_lowercase		)
									pipe.set_progress_bar_config(disable=_lowercase		)

									snake_case   :			Dict						    = "teddy bear playing in the pool"

									snake_case   :			Union[str, Any]						    = torch.Generator(device=_lowercase		).manual_seed(0		)
									snake_case   :			Union[str, Any]						    = pipe([prompt]   ,	generator=_lowercase   ,	num_inference_steps=2   ,	output_type="np"		)
									snake_case   :			Optional[Any]						    = output.images

									snake_case   :			Dict						    = torch.Generator(device=_lowercase		).manual_seed(0		)
									snake_case   :			Union[str, Any]						    = pipe(
									    [prompt]   ,	generator=_lowercase   ,	output_type="np"   ,	return_dict=_lowercase   ,	num_inference_steps=2		)[0]

									snake_case   :			Any						    = image[0, -3:, -3:, -1]
									snake_case   :			Any						    = image_from_tuple[0, -3:, -3:, -1]

									assert image.shape == (1, 24, 24, 3)

									snake_case   :			Optional[Any]						    = np.array([0.6693, 0.6075, 0.4959, 0.5701, 0.5583, 0.4333, 0.6171, 0.5684, 0.4988]		)

									assert np.abs(image_slice.flatten() - expected_slice		).max() < 2.0
									assert np.abs(image_from_tuple_slice.flatten() - expected_slice		).max() < 1E-2



@slow
@require_torch_gpu
class       _a  ( unittest.TestCase):



			def 				__lowercase					(			self							:				Optional[int]		)							->					Tuple:
									# clean up the VRAM after each test
									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()



			def 				__lowercase					(			self							:				Dict		)							->					Tuple:
									snake_case   :			Union[str, Any]						    = load_numpy(
									    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
									    "/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy"		)

									snake_case   :			Tuple						    = VQDiffusionPipeline.from_pretrained("microsoft/vq-diffusion-ithq"		)
									snake_case   :			Union[str, Any]						    = pipeline.to(_lowercase		)
									pipeline.set_progress_bar_config(disable=_lowercase		)

									# requires GPU generator for gumbel softmax
									# don't use GPU generator in tests though
									snake_case   :			Tuple						    = torch.Generator(device=_lowercase		).manual_seed(0		)
									snake_case   :			Optional[int]						    = pipeline(
									    "teddy bear playing in the pool"   ,	num_images_per_prompt=1   ,	generator=_lowercase   ,	output_type="np"   ,	)

									snake_case   :			List[str]						    = output.images[0]

									assert image.shape == (256, 256, 3)
									assert np.abs(expected_image - image		).max() < 2.0



