





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







       """simple docstring"""





       SCREAMING_SNAKE_CASE_      : List[str]      =  MODEL_FOR_CAUSAL_LM_MAPPING
       SCREAMING_SNAKE_CASE_      : List[Any]      =  TF_MODEL_FOR_CAUSAL_LM_MAPPING







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

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

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







       @require_tf
       def    __UpperCAmelCase (    self	:	int  )						->					int:
         _lowercase							= pipeline(task='text-generation'						,model='sshleifer/tiny-ctrl'						,framework='tf'  )

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

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







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

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







       def    __UpperCAmelCase (    self	:	str						,__A	:	List[str]						,__A	:	int  )						->					Dict:
         _lowercase							= text_generator.model
         _lowercase							= text_generator.tokenizer

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

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

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

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

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

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

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

         # 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__
         ):
           _lowercase							= text_generator(''  )
           self.assertEqual(__A						,[{'generated_text': ANY(__A  )}]  )
         else:
           with self.assertRaises((ValueError, AssertionError)  ):
             _lowercase							= 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.
         _lowercase							= ['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  )

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

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

         _lowercase							= 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	:	Any  )						->					Any:
         import torch

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







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

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

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






snake_case 				=       """0.18.2"""

from .configuration_utils import ConfigMixin
from .utils import (
    OptionalDependencyNotAvailable,
    is_flax_available,
    is_inflect_available,
    is_invisible_watermark_available,
    is_k_diffusion_available,
    is_k_diffusion_version,
    is_librosa_available,
    is_note_seq_available,
    is_onnx_available,
    is_scipy_available,
    is_torch_available,
    is_torchsde_available,
    is_transformers_available,
    is_transformers_version,
    is_unidecode_available,
    logging,
)


try:
      if not is_onnx_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      from .utils.dummy_onnx_objects import *  # noqa F403
else:
      from .pipelines import OnnxRuntimeModel

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      from .utils.dummy_pt_objects import *  # noqa F403
else:
      from .models import (
          AutoencoderKL,
          ControlNetModel,
          ModelMixin,
          PriorTransformer,
          TaFilmDecoder,
          TransformeraDModel,
          UNetaDModel,
          UNetaDConditionModel,
          UNetaDModel,
          UNetaDConditionModel,
          VQModel,
      )
      from .optimization import (
          get_constant_schedule,
          get_constant_schedule_with_warmup,
          get_cosine_schedule_with_warmup,
          get_cosine_with_hard_restarts_schedule_with_warmup,
          get_linear_schedule_with_warmup,
          get_polynomial_decay_schedule_with_warmup,
          get_scheduler,
      )
      from .pipelines import (
          AudioPipelineOutput,
          ConsistencyModelPipeline,
          DanceDiffusionPipeline,
          DDIMPipeline,
          DDPMPipeline,
          DiffusionPipeline,
          DiTPipeline,
          ImagePipelineOutput,
          KarrasVePipeline,
          LDMPipeline,
          LDMSuperResolutionPipeline,
          PNDMPipeline,
          RePaintPipeline,
          ScoreSdeVePipeline,
      )
      from .schedulers import (
          CMStochasticIterativeScheduler,
          DDIMInverseScheduler,
          DDIMParallelScheduler,
          DDIMScheduler,
          DDPMParallelScheduler,
          DDPMScheduler,
          DEISMultistepScheduler,
          DPMSolverMultistepInverseScheduler,
          DPMSolverMultistepScheduler,
          DPMSolverSinglestepScheduler,
          EulerAncestralDiscreteScheduler,
          EulerDiscreteScheduler,
          HeunDiscreteScheduler,
          IPNDMScheduler,
          KarrasVeScheduler,
          KDPMaAncestralDiscreteScheduler,
          KDPMaDiscreteScheduler,
          PNDMScheduler,
          RePaintScheduler,
          SchedulerMixin,
          ScoreSdeVeScheduler,
          UnCLIPScheduler,
          UniPCMultistepScheduler,
          VQDiffusionScheduler,
      )
      from .training_utils import EMAModel

try:
      if not (is_torch_available() and is_scipy_available()):
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      from .utils.dummy_torch_and_scipy_objects import *  # noqa F403
else:
      from .schedulers import LMSDiscreteScheduler

try:
      if not (is_torch_available() and is_torchsde_available()):
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      from .utils.dummy_torch_and_torchsde_objects import *  # noqa F403
else:
      from .schedulers import DPMSolverSDEScheduler

try:
      if not (is_torch_available() and is_transformers_available()):
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      from .utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
      from .pipelines import (
          AltDiffusionImgaImgPipeline,
          AltDiffusionPipeline,
          AudioLDMPipeline,
          CycleDiffusionPipeline,
          IFImgaImgPipeline,
          IFImgaImgSuperResolutionPipeline,
          IFInpaintingPipeline,
          IFInpaintingSuperResolutionPipeline,
          IFPipeline,
          IFSuperResolutionPipeline,
          ImageTextPipelineOutput,
          KandinskyImgaImgPipeline,
          KandinskyInpaintPipeline,
          KandinskyPipeline,
          KandinskyPriorPipeline,
          KandinskyVaaControlnetImgaImgPipeline,
          KandinskyVaaControlnetPipeline,
          KandinskyVaaImgaImgPipeline,
          KandinskyVaaInpaintPipeline,
          KandinskyVaaPipeline,
          KandinskyVaaPriorEmbaEmbPipeline,
          KandinskyVaaPriorPipeline,
          LDMTextToImagePipeline,
          PaintByExamplePipeline,
          SemanticStableDiffusionPipeline,
          ShapEImgaImgPipeline,
          ShapEPipeline,
          StableDiffusionAttendAndExcitePipeline,
          StableDiffusionControlNetImgaImgPipeline,
          StableDiffusionControlNetInpaintPipeline,
          StableDiffusionControlNetPipeline,
          StableDiffusionDepthaImgPipeline,
          StableDiffusionDiffEditPipeline,
          StableDiffusionImageVariationPipeline,
          StableDiffusionImgaImgPipeline,
          StableDiffusionInpaintPipeline,
          StableDiffusionInpaintPipelineLegacy,
          StableDiffusionInstructPixaPixPipeline,
          StableDiffusionLatentUpscalePipeline,
          StableDiffusionLDMaDPipeline,
          StableDiffusionModelEditingPipeline,
          StableDiffusionPanoramaPipeline,
          StableDiffusionParadigmsPipeline,
          StableDiffusionPipeline,
          StableDiffusionPipelineSafe,
          StableDiffusionPixaPixZeroPipeline,
          StableDiffusionSAGPipeline,
          StableDiffusionUpscalePipeline,
          StableUnCLIPImgaImgPipeline,
          StableUnCLIPPipeline,
          TextToVideoSDPipeline,
          TextToVideoZeroPipeline,
          UnCLIPImageVariationPipeline,
          UnCLIPPipeline,
          UniDiffuserModel,
          UniDiffuserPipeline,
          UniDiffuserTextDecoder,
          VersatileDiffusionDualGuidedPipeline,
          VersatileDiffusionImageVariationPipeline,
          VersatileDiffusionPipeline,
          VersatileDiffusionTextToImagePipeline,
          VideoToVideoSDPipeline,
          VQDiffusionPipeline,
      )

try:
      if not (is_torch_available() and is_transformers_available() and is_invisible_watermark_available()):
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      from .utils.dummy_torch_and_transformers_and_invisible_watermark_objects import *  # noqa F403
else:
      from .pipelines import StableDiffusionXLImgaImgPipeline, StableDiffusionXLPipeline

try:
      if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()):
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      from .utils.dummy_torch_and_transformers_and_k_diffusion_objects import *  # noqa F403
else:
      from .pipelines import StableDiffusionKDiffusionPipeline

try:
      if not (is_torch_available() and is_transformers_available() and is_onnx_available()):
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      from .utils.dummy_torch_and_transformers_and_onnx_objects import *  # noqa F403
else:
      from .pipelines import (
          OnnxStableDiffusionImgaImgPipeline,
          OnnxStableDiffusionInpaintPipeline,
          OnnxStableDiffusionInpaintPipelineLegacy,
          OnnxStableDiffusionPipeline,
          OnnxStableDiffusionUpscalePipeline,
          StableDiffusionOnnxPipeline,
      )

try:
      if not (is_torch_available() and is_librosa_available()):
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      from .utils.dummy_torch_and_librosa_objects import *  # noqa F403
else:
      from .pipelines import AudioDiffusionPipeline, Mel

try:
      if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      from .utils.dummy_transformers_and_torch_and_note_seq_objects import *  # noqa F403
else:
      from .pipelines import SpectrogramDiffusionPipeline

try:
      if not is_flax_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      from .utils.dummy_flax_objects import *  # noqa F403
else:
      from .models.controlnet_flax import FlaxControlNetModel
      from .models.modeling_flax_utils import FlaxModelMixin
      from .models.unet_ad_condition_flax import FlaxUNetaDConditionModel
      from .models.vae_flax import FlaxAutoencoderKL
      from .pipelines import FlaxDiffusionPipeline
      from .schedulers import (
          FlaxDDIMScheduler,
          FlaxDDPMScheduler,
          FlaxDPMSolverMultistepScheduler,
          FlaxKarrasVeScheduler,
          FlaxLMSDiscreteScheduler,
          FlaxPNDMScheduler,
          FlaxSchedulerMixin,
          FlaxScoreSdeVeScheduler,
      )


try:
      if not (is_flax_available() and is_transformers_available()):
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      from .utils.dummy_flax_and_transformers_objects import *  # noqa F403
else:
      from .pipelines import (
          FlaxStableDiffusionControlNetPipeline,
          FlaxStableDiffusionImgaImgPipeline,
          FlaxStableDiffusionInpaintPipeline,
          FlaxStableDiffusionPipeline,
      )

try:
      if not (is_note_seq_available()):
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      from .utils.dummy_note_seq_objects import *  # noqa F403
else:
      from .pipelines import MidiProcessor