



'''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        UpperCAmelCase_			(			unittest.TestCase      ):
    """simple docstring"""







    lowercase   =						MODEL_FOR_CAUSAL_LM_MAPPING
    lowercase   =						TF_MODEL_FOR_CAUSAL_LM_MAPPING




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

          snake_case__    :       List[str]       		= 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@@"""
                          )
                      }
                  ],
              ]					,	)

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

          # Using `do_sample=False` to force deterministic output
          snake_case__    :       Optional[int]       		= 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,"""
                      )
                  }
              ]					,	)

          snake_case__    :       Tuple       		= 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 lowerCamelCase			(       self   :	Union[str, Any]					,	snake_case_   :	int					,	snake_case_   :	Dict					,	snake_case_   :	Optional[Any]		):
          snake_case__    :       int       		= TextGenerationPipeline(model=snake_case_					,	tokenizer=snake_case_		)
          return text_generator, ["This is a test", "Another test"]




    def lowerCamelCase			(       self   :	Union[str, Any]		):
          snake_case__    :       List[Any]       		= """Hello I believe in"""
          snake_case__    :       int       		= pipeline("""text-generation"""					,	model="""hf-internal-testing/tiny-random-gpt2"""		)
          snake_case__    :       Dict       		= 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"""}]					,	)

          snake_case__    :       Dict       		= text_generator(snake_case_					,	stop_sequence=""" fe"""		)
          self.assertEqual(snake_case_					,	[{"""generated_text""": """Hello I believe in fe"""}]		)




    def lowerCamelCase			(       self   :	Optional[int]					,	snake_case_   :	Optional[Any]					,	snake_case_   :	str		):
          snake_case__    :       Optional[int]       		= text_generator.model
          snake_case__    :       Any       		= text_generator.tokenizer

          snake_case__    :       Dict       		= 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"""		)		)

          snake_case__    :       Any       		= 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"""]		)

          snake_case__    :       Dict       		= pipeline(task="""text-generation"""					,	model=snake_case_					,	tokenizer=snake_case_					,	return_full_text=snake_case_		)
          snake_case__    :       Tuple       		= 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"""]		)

          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"""		)		)

          snake_case__    :       Union[str, 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:
                snake_case__    :       List[Any]       		= 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_		):
                snake_case__    :       Optional[int]       		= text_generator("""test"""					,	return_full_text=snake_case_					,	return_text=snake_case_		)
          with self.assertRaises(snake_case_		):
                snake_case__    :       List[str]       		= text_generator("""test"""					,	return_full_text=snake_case_					,	return_tensors=snake_case_		)
          with self.assertRaises(snake_case_		):
                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__
          ):
                snake_case__    :       Optional[int]       		= text_generator(""""""		)
                self.assertEqual(snake_case_					,	[{"""generated_text""": ANY(snake_case_		)}]		)
          else:
                with self.assertRaises((ValueError, AssertionError)		):
                      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.
          snake_case__    :       Optional[int]       		= ["""RwkvForCausalLM""", """XGLMForCausalLM""", """GPTNeoXForCausalLM"""]
          if (
              tokenizer.model_max_length < 10_000
              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		)

                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 lowerCamelCase			(       self   :	Union[str, Any]		):
          import torch

          # Classic `model_kwargs`
          snake_case__    :       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		)
          snake_case__    :       Optional[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.)
          snake_case__    :       Union[str, Any]       		= 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		)
          snake_case__    :       Tuple       		= 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
          snake_case__    :       Tuple       		= 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		)
          snake_case__    :       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 lowerCamelCase			(       self   :	Optional[int]		):
          import torch

          snake_case__    :       Optional[int]       		= 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 lowerCamelCase			(       self   :	Dict		):
          import torch

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

          with CaptureLogger(snake_case_		) as cl:
                snake_case__    :       str       		= text_generator(snake_case_					,	max_length=10		)
          self.assertNotIn(snake_case_					,	cl.out		)





'''simple docstring'''





from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


__a  			=						logging.get_logger(__name__)

__a  			=						{
    "shi-labs/dinat-mini-in1k-224": "https://huggingface.co/shi-labs/dinat-mini-in1k-224/resolve/main/config.json",
    # See all Dinat models at https://huggingface.co/models?filter=dinat
}







class        UpperCAmelCase_			(			_a						,						_a      ):
    """simple docstring"""







    lowercase   =						"dinat"

    lowercase   =						{
        "num_attention_heads": "num_heads",
        "num_hidden_layers": "num_layers",
    }




    def __init__(       self   :	Optional[int]					,	snake_case_   :	Optional[Any]=4					,	snake_case_   :	List[str]=3					,	snake_case_   :	Tuple=64					,	snake_case_   :	List[Any]=[3, 4, 6, 5]					,	snake_case_   :	List[Any]=[2, 4, 8, 16]					,	snake_case_   :	List[str]=7					,	snake_case_   :	List[str]=[[1, 8, 1], [1, 4, 1, 4], [1, 2, 1, 2, 1, 2], [1, 1, 1, 1, 1]]					,	snake_case_   :	List[Any]=3.0					,	snake_case_   :	Dict=True					,	snake_case_   :	int=0.0					,	snake_case_   :	Optional[int]=0.0					,	snake_case_   :	str=0.1					,	snake_case_   :	Optional[Any]="gelu"					,	snake_case_   :	Union[str, Any]=0.02					,	snake_case_   :	Optional[int]=1E-5					,	snake_case_   :	str=0.0					,	snake_case_   :	Dict=None					,	snake_case_   :	Dict=None					,	**snake_case_   :	Dict					,	):
          super().__init__(**snake_case_		)

          snake_case__    :       Optional[Any]       		= patch_size
          snake_case__    :       Union[str, Any]       		= num_channels
          snake_case__    :       int       		= embed_dim
          snake_case__    :       Dict       		= depths
          snake_case__    :       Dict       		= len(snake_case_		)
          snake_case__    :       Optional[int]       		= num_heads
          snake_case__    :       Any       		= kernel_size
          snake_case__    :       str       		= dilations
          snake_case__    :       Any       		= mlp_ratio
          snake_case__    :       int       		= qkv_bias
          snake_case__    :       int       		= hidden_dropout_prob
          snake_case__    :       int       		= attention_probs_dropout_prob
          snake_case__    :       Optional[Any]       		= drop_path_rate
          snake_case__    :       int       		= hidden_act
          snake_case__    :       Tuple       		= layer_norm_eps
          snake_case__    :       Any       		= initializer_range
          # we set the hidden_size attribute in order to make Dinat work with VisionEncoderDecoderModel
          # this indicates the channel dimension after the last stage of the model
          snake_case__    :       Optional[Any]       		= int(embed_dim * 2 ** (len(snake_case_		) - 1)		)
          snake_case__    :       int       		= layer_scale_init_value
          snake_case__    :       Dict       		= ["""stem"""] + [f"stage{idx}" for idx in range(1					,	len(snake_case_		) + 1		)]
          snake_case__		,							snake_case__    :       Tuple       		= get_aligned_output_features_output_indices(
              out_features=snake_case_					,	out_indices=snake_case_					,	stage_names=self.stage_names		)
