




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  _lowerCAmelCase     (  unittest.TestCase	):
      A__		     =		MODEL_FOR_CAUSAL_LM_MAPPING
      A__		     =		TF_MODEL_FOR_CAUSAL_LM_MAPPING

      @require_torch
      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  List[str]		       =     pipeline(task='''text-generation'''      ,				model='''sshleifer/tiny-ctrl'''      ,				framework='''pt'''		)
        # Using `do_sample=False` to force deterministic output
        lowerCAmelCase__				:  List[Any]		       =     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@@'''
                    )
                }
            ]      ,				)

        lowerCAmelCase__				:  int		       =     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@@'''
                        )
                    }
                ],
            ]      ,				)

        lowerCAmelCase__				:  str		       =     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_		)},
            ]      ,				)
        lowerCAmelCase__				:  Union[str, Any]		       =     text_generator.model.config.eos_token_id
        lowerCAmelCase__				:  int		       =     '''<pad>'''
        lowerCAmelCase__				:  int		       =     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 	__magic_name__(   self		):
        lowerCAmelCase__				:  Tuple		       =     pipeline(task='''text-generation'''      ,				model='''sshleifer/tiny-ctrl'''      ,				framework='''tf'''		)

        # Using `do_sample=False` to force deterministic output
        lowerCAmelCase__				:  Dict		       =     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,'''
                    )
                }
            ]      ,				)

        lowerCAmelCase__				:  Dict		       =     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 	__magic_name__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase		):
        lowerCAmelCase__				:  Optional[Any]		       =     TextGenerationPipeline(model=A_      ,				tokenizer=A_		)
        return text_generator, ["This is a test", "Another test"]

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  List[str]		       =     '''Hello I believe in'''
        lowerCAmelCase__				:  List[str]		       =     pipeline('''text-generation'''      ,				model='''hf-internal-testing/tiny-random-gpt2'''		)
        lowerCAmelCase__				:  str		       =     text_generator(A_		)
        self.assertEqual(
            A_      ,				[{'''generated_text''': '''Hello I believe in fe fe fe fe fe fe fe fe fe fe fe fe'''}]      ,				)

        lowerCAmelCase__				:  Optional[Any]		       =     text_generator(A_      ,				stop_sequence=''' fe'''		)
        self.assertEqual(A_      ,				[{'''generated_text''': '''Hello I believe in fe'''}]		)

      def 	__magic_name__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase		):
        lowerCAmelCase__				:  Union[str, Any]		       =     text_generator.model
        lowerCAmelCase__				:  Dict		       =     text_generator.tokenizer

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

        lowerCAmelCase__				:  str		       =     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''']		)

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

        lowerCAmelCase__				:  Tuple		       =     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'''		)		)

        lowerCAmelCase__				:  List[Any]		       =     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:
          lowerCAmelCase__				:  List[str]		       =     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_		):
          lowerCAmelCase__				:  Union[str, Any]		       =     text_generator('''test'''      ,				return_full_text=A_      ,				return_text=A_		)
        with self.assertRaises(A_		):
          lowerCAmelCase__				:  Optional[int]		       =     text_generator('''test'''      ,				return_full_text=A_      ,				return_tensors=A_		)
        with self.assertRaises(A_		):
          lowerCAmelCase__				:  Tuple		       =     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__
        ):
          lowerCAmelCase__				:  Optional[Any]		       =     text_generator(''''''		)
          self.assertEqual(A_      ,				[{'''generated_text''': ANY(A_		)}]		)
        else:
          with self.assertRaises((ValueError, AssertionError)		):
            lowerCAmelCase__				:  List[Any]		       =     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.
        lowerCAmelCase__				:  List[str]		       =     ['''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		)

          lowerCAmelCase__				:  Dict		       =     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 	__magic_name__(   self		):
        import torch

        # Classic `model_kwargs`
        lowerCAmelCase__				:  Dict		       =     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		)
        lowerCAmelCase__				:  Optional[Any]		       =     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.)
        lowerCAmelCase__				:  str		       =     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		)
        lowerCAmelCase__				:  Any		       =     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
        lowerCAmelCase__				:  int		       =     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		)
        lowerCAmelCase__				:  Any		       =     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 	__magic_name__(   self		):
        import torch

        lowerCAmelCase__				:  str		       =     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 	__magic_name__(   self		):
        import torch

        lowerCAmelCase__				:  Optional[Any]		       =     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 	__magic_name__(   self		):
        lowerCAmelCase__				:  Dict		       =     '''Hello world'''
        lowerCAmelCase__				:  List[str]		       =     pipeline('''text-generation'''      ,				model='''hf-internal-testing/tiny-random-gpt2'''		)
        if text_generator.model.framework == "tf":
          lowerCAmelCase__				:  List[Any]		       =     logging.get_logger('''transformers.generation.tf_utils'''		)
        else:
          lowerCAmelCase__				:  Tuple		       =     logging.get_logger('''transformers.generation.utils'''		)
        lowerCAmelCase__				:  Union[str, Any]		       =     '''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:
          lowerCAmelCase__				:  str		       =     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:
          lowerCAmelCase__				:  List[Any]		       =     text_generator(A_      ,				max_new_tokens=1		)
        self.assertNotIn(A_      ,				cl.out		)

        with CaptureLogger(A_		) as cl:
          lowerCAmelCase__				:  Optional[int]		       =     text_generator(A_      ,				max_length=10		)
        self.assertNotIn(A_      ,				cl.out		)







import unittest

import numpy as np

from diffusers import LMSDiscreteScheduler, OnnxStableDiffusionInpaintPipeline
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
     import onnxruntime as ort






class  _lowerCAmelCase     (  _lowercase    ,			unittest.TestCase	):
      # FIXME: add fast tests
      pass






@nightly
@require_onnxruntime
@require_torch_gpu
class  _lowerCAmelCase     (  unittest.TestCase	):

      @property
      def 	__magic_name__(   self		):
        return (
            "CUDAExecutionProvider",
            {
                "gpu_mem_limit": "15000000000",  # 15GB
                "arena_extend_strategy": "kSameAsRequested",
            },
        )

      @property
      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  str		       =     ort.SessionOptions()
        lowerCAmelCase__				:  List[Any]		       =     False
        return options

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  Tuple		       =     load_image(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
            '''/in_paint/overture-creations-5sI6fQgYIuo.png'''		)
        lowerCAmelCase__				:  Optional[int]		       =     load_image(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
            '''/in_paint/overture-creations-5sI6fQgYIuo_mask.png'''		)
        lowerCAmelCase__				:  List[Any]		       =     OnnxStableDiffusionInpaintPipeline.from_pretrained(
            '''runwayml/stable-diffusion-inpainting'''      ,				revision='''onnx'''      ,				safety_checker=__UpperCAmelCase      ,				feature_extractor=__UpperCAmelCase      ,				provider=self.gpu_provider      ,				sess_options=self.gpu_options      ,				)
        pipe.set_progress_bar_config(disable=__UpperCAmelCase		)

        lowerCAmelCase__				:  Optional[int]		       =     '''A red cat sitting on a park bench'''

        lowerCAmelCase__				:  List[Any]		       =     np.random.RandomState(0		)
        lowerCAmelCase__				:  int		       =     pipe(
            prompt=__UpperCAmelCase      ,				image=__UpperCAmelCase      ,				mask_image=__UpperCAmelCase      ,				guidance_scale=7.5      ,				num_inference_steps=10      ,				generator=__UpperCAmelCase      ,				output_type='''np'''      ,				)
        lowerCAmelCase__				:  Any		       =     output.images
        lowerCAmelCase__				:  Any		       =     images[0, 255:258, 255:258, -1]

        assert images.shape == (1, 512, 512, 3)
        lowerCAmelCase__				:  Any		       =     np.array([0.2514, 0.3007, 0.3517, 0.1790, 0.2382, 0.3167, 0.1944, 0.2273, 0.2464]		)

        assert np.abs(image_slice.flatten() - expected_slice		).max() < 1e-3







      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  List[str]		       =     load_image(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
            '''/in_paint/overture-creations-5sI6fQgYIuo.png'''		)
        lowerCAmelCase__				:  str		       =     load_image(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
            '''/in_paint/overture-creations-5sI6fQgYIuo_mask.png'''		)
        lowerCAmelCase__				:  Union[str, Any]		       =     LMSDiscreteScheduler.from_pretrained(
            '''runwayml/stable-diffusion-inpainting'''      ,				subfolder='''scheduler'''      ,				revision='''onnx'''		)
        lowerCAmelCase__				:  Union[str, Any]		       =     OnnxStableDiffusionInpaintPipeline.from_pretrained(
            '''runwayml/stable-diffusion-inpainting'''      ,				revision='''onnx'''      ,				scheduler=__UpperCAmelCase      ,				safety_checker=__UpperCAmelCase      ,				feature_extractor=__UpperCAmelCase      ,				provider=self.gpu_provider      ,				sess_options=self.gpu_options      ,				)
        pipe.set_progress_bar_config(disable=__UpperCAmelCase		)

        lowerCAmelCase__				:  str		       =     '''A red cat sitting on a park bench'''

        lowerCAmelCase__				:  Union[str, Any]		       =     np.random.RandomState(0		)
        lowerCAmelCase__				:  Optional[Any]		       =     pipe(
            prompt=__UpperCAmelCase      ,				image=__UpperCAmelCase      ,				mask_image=__UpperCAmelCase      ,				guidance_scale=7.5      ,				num_inference_steps=20      ,				generator=__UpperCAmelCase      ,				output_type='''np'''      ,				)
        lowerCAmelCase__				:  Tuple		       =     output.images
        lowerCAmelCase__				:  Optional[Any]		       =     images[0, 255:258, 255:258, -1]

        assert images.shape == (1, 512, 512, 3)
        lowerCAmelCase__				:  Optional[int]		       =     np.array([0.0086, 0.0077, 0.0083, 0.0093, 0.0107, 0.0139, 0.0094, 0.0097, 0.0125]		)

        assert np.abs(image_slice.flatten() - expected_slice		).max() < 1e-3

