






"""simple docstring"""







import collections
import inspect
import unittest

from transformers import NatConfig
from transformers.testing_utils import require_natten, require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch
  from torch import nn

  from transformers import NatBackbone, NatForImageClassification, NatModel
  from transformers.models.nat.modeling_nat import NAT_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
  from PIL import Image

  from transformers import AutoImageProcessor







class 					lowerCamelCase__					:






    def __init__( self	,A	,A=13	,A=64	,A=4	,A=3	,A=16	,A=[1, 2, 1]	,A=[2, 4, 8]	,A=3	,A=2.0	,A=True	,A=0.0	,A=0.0	,A=0.1	,A="gelu"	,A=True	,A=0.02	,A=1e-5	,A=True	,A=None	,A=True	,A=10	,A=["stage1", "stage2"]	,A=[1, 2]	,):
          UpperCAmelCase							=       parent
          UpperCAmelCase							=       batch_size
          UpperCAmelCase							=       image_size
          UpperCAmelCase							=       patch_size
          UpperCAmelCase							=       num_channels
          UpperCAmelCase							=       embed_dim
          UpperCAmelCase							=       depths
          UpperCAmelCase							=       num_heads
          UpperCAmelCase							=       kernel_size
          UpperCAmelCase							=       mlp_ratio
          UpperCAmelCase							=       qkv_bias
          UpperCAmelCase							=       hidden_dropout_prob
          UpperCAmelCase							=       attention_probs_dropout_prob
          UpperCAmelCase							=       drop_path_rate
          UpperCAmelCase							=       hidden_act
          UpperCAmelCase							=       patch_norm
          UpperCAmelCase							=       layer_norm_eps
          UpperCAmelCase							=       initializer_range
          UpperCAmelCase							=       is_training
          UpperCAmelCase							=       scope
          UpperCAmelCase							=       use_labels
          UpperCAmelCase							=       num_labels
          UpperCAmelCase							=       out_features
          UpperCAmelCase							=       out_indices






    def 				_UpperCamelCase							( self					):
          UpperCAmelCase							=       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

          UpperCAmelCase							=       None
          if self.use_labels:
                UpperCAmelCase							=       ids_tensor([self.batch_size]	,self.num_labels					)

          UpperCAmelCase							=       self.get_config()

          return config, pixel_values, labels






    def 				_UpperCamelCase							( self					):
          return NatConfig(
              num_labels=self.num_labels	,image_size=self.image_size	,patch_size=self.patch_size	,num_channels=self.num_channels	,embed_dim=self.embed_dim	,depths=self.depths	,num_heads=self.num_heads	,kernel_size=self.kernel_size	,mlp_ratio=self.mlp_ratio	,qkv_bias=self.qkv_bias	,hidden_dropout_prob=self.hidden_dropout_prob	,attention_probs_dropout_prob=self.attention_probs_dropout_prob	,drop_path_rate=self.drop_path_rate	,hidden_act=self.hidden_act	,patch_norm=self.patch_norm	,layer_norm_eps=self.layer_norm_eps	,initializer_range=self.initializer_range	,out_features=self.out_features	,out_indices=self.out_indices	,)






    def 				_UpperCamelCase							( self	,A	,A	,A					):
          UpperCAmelCase							=       NatModel(config=snake_case__					)
          model.to(snake_case__					)
          model.eval()
          UpperCAmelCase							=       model(snake_case__					)

          UpperCAmelCase							=       UpperCAmelCase							=       (config.image_size // config.patch_size) // (2 ** (len(config.depths					) - 1))
          UpperCAmelCase							=       int(config.embed_dim * 2 ** (len(config.depths					) - 1)					)

          self.parent.assertEqual(
              result.last_hidden_state.shape	,(self.batch_size, expected_height, expected_width, expected_dim)					)






    def 				_UpperCamelCase							( self	,A	,A	,A					):
          UpperCAmelCase							=       NatForImageClassification(snake_case__					)
          model.to(snake_case__					)
          model.eval()
          UpperCAmelCase							=       model(snake_case__	,labels=snake_case__					)
          self.parent.assertEqual(result.logits.shape	,(self.batch_size, self.num_labels)					)

          # test greyscale images
          UpperCAmelCase							=       1
          UpperCAmelCase							=       NatForImageClassification(snake_case__					)
          model.to(snake_case__					)
          model.eval()

          UpperCAmelCase							=       floats_tensor([self.batch_size, 1, self.image_size, self.image_size]					)
          UpperCAmelCase							=       model(snake_case__					)
          self.parent.assertEqual(result.logits.shape	,(self.batch_size, self.num_labels)					)






    def 				_UpperCamelCase							( self	,A	,A	,A					):
          UpperCAmelCase							=       NatBackbone(config=snake_case__					)
          model.to(snake_case__					)
          model.eval()
          UpperCAmelCase							=       model(snake_case__					)

          # verify hidden states
          self.parent.assertEqual(len(result.feature_maps					)	,len(config.out_features					)					)
          self.parent.assertListEqual(list(result.feature_maps[0].shape					)	,[self.batch_size, model.channels[0], 16, 16]					)

          # verify channels
          self.parent.assertEqual(len(model.channels					)	,len(config.out_features					)					)

          # verify backbone works with out_features=None
          UpperCAmelCase							=       None
          UpperCAmelCase							=       NatBackbone(config=snake_case__					)
          model.to(snake_case__					)
          model.eval()
          UpperCAmelCase							=       model(snake_case__					)

          # verify feature maps
          self.parent.assertEqual(len(result.feature_maps					)	,1					)
          self.parent.assertListEqual(list(result.feature_maps[0].shape					)	,[self.batch_size, model.channels[-1], 4, 4]					)

          # verify channels
          self.parent.assertEqual(len(model.channels					)	,1					)






    def 				_UpperCamelCase							( self					):
          UpperCAmelCase							=       self.prepare_config_and_inputs()
          UpperCAmelCase			,					UpperCAmelCase			,					UpperCAmelCase							=       config_and_inputs
          UpperCAmelCase							=       {"""pixel_values""": pixel_values}
          return config, inputs_dict







@require_natten
@require_torch
class 					lowerCamelCase__					(       a_    ,     a_    ,     unittest.TestCase ):
    SCREAMING_SNAKE_CASE		=			(
        (
            NatModel,
            NatForImageClassification,
            NatBackbone,
        )
        if is_torch_available()
        else ()
    )
    SCREAMING_SNAKE_CASE		=			(
        {'feature-extraction': NatModel, 'image-classification': NatForImageClassification}
        if is_torch_available()
        else {}
    )
    SCREAMING_SNAKE_CASE		=			False

    SCREAMING_SNAKE_CASE		=			False
    SCREAMING_SNAKE_CASE		=			False
    SCREAMING_SNAKE_CASE		=			False
    SCREAMING_SNAKE_CASE		=			False






    def 				_UpperCamelCase							( self					):
          UpperCAmelCase							=       NatModelTester(self					)
          UpperCAmelCase							=       ConfigTester(self	,config_class=snake_case__	,embed_dim=37					)






    def 				_UpperCamelCase							( self					):
          self.create_and_test_config_common_properties()
          self.config_tester.create_and_test_config_to_json_string()
          self.config_tester.create_and_test_config_to_json_file()
          self.config_tester.create_and_test_config_from_and_save_pretrained()
          self.config_tester.create_and_test_config_with_num_labels()
          self.config_tester.check_config_can_be_init_without_params()
          self.config_tester.check_config_arguments_init()






    def 				_UpperCamelCase							( self					):
          return






    def 				_UpperCamelCase							( self					):
          UpperCAmelCase							=       self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*snake_case__					)






    def 				_UpperCamelCase							( self					):
          UpperCAmelCase							=       self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_image_classification(*snake_case__					)






    def 				_UpperCamelCase							( self					):
          UpperCAmelCase							=       self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_backbone(*snake_case__					)






    @unittest.skip(reason="""Nat does not use inputs_embeds"""					)
    def 				_UpperCamelCase							( self					):
          pass






    @unittest.skip(reason="""Nat does not use feedforward chunking"""					)
    def 				_UpperCamelCase							( self					):
          pass






    def 				_UpperCamelCase							( self					):
          UpperCAmelCase			,					UpperCAmelCase							=       self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
                UpperCAmelCase							=       model_class(snake_case__					)
                self.assertIsInstance(model.get_input_embeddings()	,(nn.Module)					)
                UpperCAmelCase							=       model.get_output_embeddings()
                self.assertTrue(x is None or isinstance(snake_case__	,nn.Linear					)					)






    def 				_UpperCamelCase							( self					):
          UpperCAmelCase			,					UpperCAmelCase							=       self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
                UpperCAmelCase							=       model_class(snake_case__					)
                UpperCAmelCase							=       inspect.signature(model.forward					)
                # signature.parameters is an OrderedDict => so arg_names order is deterministic
                UpperCAmelCase							=       [*signature.parameters.keys()]

                UpperCAmelCase							=       ["""pixel_values"""]
                self.assertListEqual(arg_names[:1]	,snake_case__					)






    def 				_UpperCamelCase							( self					):
          self.skipTest("""Nat's attention operation is handled entirely by NATTEN."""					)






    def 				_UpperCamelCase							( self	,A	,A	,A	,A					):
          UpperCAmelCase							=       model_class(snake_case__					)
          model.to(snake_case__					)
          model.eval()

          with torch.no_grad():
                UpperCAmelCase							=       model(**self._prepare_for_class(snake_case__	,snake_case__					)					)

          UpperCAmelCase							=       outputs.hidden_states

          UpperCAmelCase							=       getattr(
              self.model_tester	,"""expected_num_hidden_layers"""	,len(self.model_tester.depths					) + 1					)
          self.assertEqual(len(snake_case__					)	,snake_case__					)

          # Nat has a different seq_length
          UpperCAmelCase							=       (
              config.patch_size
              if isinstance(config.patch_size	,collections.abc.Iterable					)
              else (config.patch_size, config.patch_size)
          )

          UpperCAmelCase							=       image_size[0] // patch_size[0]
          UpperCAmelCase							=       image_size[1] // patch_size[1]

          self.assertListEqual(
              list(hidden_states[0].shape[-3:]					)	,[height, width, self.model_tester.embed_dim]	,)

          if model_class.__name__ != "NatBackbone":
                UpperCAmelCase							=       outputs.reshaped_hidden_states
                self.assertEqual(len(snake_case__					)	,snake_case__					)

                UpperCAmelCase			,					UpperCAmelCase			,					UpperCAmelCase			,					UpperCAmelCase							=       reshaped_hidden_states[0].shape
                UpperCAmelCase							=       (
                    reshaped_hidden_states[0].view(snake_case__	,snake_case__	,snake_case__	,snake_case__					).permute(0	,2	,3	,1					)
                )
                self.assertListEqual(
                    list(reshaped_hidden_states.shape[-3:]					)	,[height, width, self.model_tester.embed_dim]	,)






    def 				_UpperCamelCase							( self					):
          UpperCAmelCase			,					UpperCAmelCase							=       self.model_tester.prepare_config_and_inputs_for_common()

          UpperCAmelCase							=       (
              self.model_tester.image_size
              if isinstance(self.model_tester.image_size	,collections.abc.Iterable					)
              else (self.model_tester.image_size, self.model_tester.image_size)
          )

          for model_class in self.all_model_classes:
                UpperCAmelCase							=       True
                self.check_hidden_states_output(snake_case__	,snake_case__	,snake_case__	,snake_case__					)

                # check that output_hidden_states also work using config
                del inputs_dict["output_hidden_states"]
                UpperCAmelCase							=       True

                self.check_hidden_states_output(snake_case__	,snake_case__	,snake_case__	,snake_case__					)






    @slow
    def 				_UpperCamelCase							( self					):





          for model_name in NAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                UpperCAmelCase							=       NatModel.from_pretrained(snake_case__					)
                self.assertIsNotNone(snake_case__					)






    def 				_UpperCamelCase							( self					):
          UpperCAmelCase			,					UpperCAmelCase							=       self.model_tester.prepare_config_and_inputs_for_common()

          UpperCAmelCase							=       _config_zero_init(snake_case__					)
          for model_class in self.all_model_classes:
                UpperCAmelCase							=       model_class(config=snake_case__					)
                for name, param in model.named_parameters():
                      if "embeddings" not in name and param.requires_grad:
                            self.assertIn(
                                ((param.data.mean() * 1e9).round() / 1e9).item()	,[0.0, 1.0]	,msg=F'''Parameter {name} of model {model_class} seems not properly initialized'''	,)







@require_natten
@require_vision
@require_torch
class 					lowerCamelCase__					(       unittest.TestCase ):






    @cached_property
    def 				_UpperCamelCase							( self					):
          return AutoImageProcessor.from_pretrained("""shi-labs/nat-mini-in1k-224"""					) if is_vision_available() else None






    @slow
    def 				_UpperCamelCase							( self					):
          UpperCAmelCase							=       NatForImageClassification.from_pretrained("""shi-labs/nat-mini-in1k-224"""					).to(snake_case__					)
          UpperCAmelCase							=       self.default_image_processor

          UpperCAmelCase							=       Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""					)
          UpperCAmelCase							=       image_processor(images=snake_case__	,return_tensors="""pt"""					).to(snake_case__					)

          # forward pass
          with torch.no_grad():
                UpperCAmelCase							=       model(**snake_case__					)

          # verify the logits
          UpperCAmelCase							=       torch.Size((1, 1_000)					)
          self.assertEqual(outputs.logits.shape	,snake_case__					)
          UpperCAmelCase							=       torch.tensor([0.3805, -0.8676, -0.3912]					).to(snake_case__					)
          self.assertTrue(torch.allclose(outputs.logits[0, :3]	,snake_case__	,atol=1e-4					)					)







@require_torch
@require_natten
class 					lowerCamelCase__					(       unittest.TestCase    ,     a_ ):
    SCREAMING_SNAKE_CASE		=			(NatBackbone,) if is_torch_available() else ()
    SCREAMING_SNAKE_CASE		=			NatConfig






    def 				_UpperCamelCase							( self					):
          UpperCAmelCase							=       NatModelTester(self					)











"""simple docstring"""







import unittest

import numpy as np

from transformers import ElectraConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
 from transformers.models.electra.modeling_flax_electra import (
     FlaxElectraForCausalLM,
     FlaxElectraForMaskedLM,
     FlaxElectraForMultipleChoice,
     FlaxElectraForPreTraining,
     FlaxElectraForQuestionAnswering,
     FlaxElectraForSequenceClassification,
     FlaxElectraForTokenClassification,
     FlaxElectraModel,
 )





class      UpperCamelCase_       (				unittest.TestCase     ):



  def __init__(					self       ,						snake_case__       ,						snake_case__=13       ,						snake_case__=7       ,						snake_case__=True       ,						snake_case__=True       ,						snake_case__=True       ,						snake_case__=True       ,						snake_case__=99       ,						snake_case__=24       ,						snake_case__=32       ,						snake_case__=5       ,						snake_case__=4       ,						snake_case__=37       ,						snake_case__="gelu"       ,						snake_case__=0.1       ,						snake_case__=0.1       ,						snake_case__=5_12       ,						snake_case__=16       ,						snake_case__=2       ,						snake_case__=0.02       ,						snake_case__=4       ,						)	->	int:



         """simple docstring"""





         UpperCAmelCase					=	parent
         UpperCAmelCase					=	batch_size
         UpperCAmelCase					=	seq_length
         UpperCAmelCase					=	is_training
         UpperCAmelCase					=	use_attention_mask
         UpperCAmelCase					=	use_token_type_ids
         UpperCAmelCase					=	use_labels
         UpperCAmelCase					=	vocab_size
         UpperCAmelCase					=	hidden_size
         UpperCAmelCase					=	embedding_size
         UpperCAmelCase					=	num_hidden_layers
         UpperCAmelCase					=	num_attention_heads
         UpperCAmelCase					=	intermediate_size
         UpperCAmelCase					=	hidden_act
         UpperCAmelCase					=	hidden_dropout_prob
         UpperCAmelCase					=	attention_probs_dropout_prob
         UpperCAmelCase					=	max_position_embeddings
         UpperCAmelCase					=	type_vocab_size
         UpperCAmelCase					=	type_sequence_label_size
         UpperCAmelCase					=	initializer_range
         UpperCAmelCase					=	num_choices



  def 						UpperCamelCase_ (					self )	->	Optional[Any]:



         """simple docstring"""





         UpperCAmelCase					=	ids_tensor([self.batch_size, self.seq_length]       ,						self.vocab_size )

         UpperCAmelCase					=	None
         if self.use_attention_mask:
                UpperCAmelCase					=	random_attention_mask([self.batch_size, self.seq_length] )

         UpperCAmelCase					=	None
         if self.use_token_type_ids:
                UpperCAmelCase					=	ids_tensor([self.batch_size, self.seq_length]       ,						self.type_vocab_size )

         UpperCAmelCase					=	ElectraConfig(
             vocab_size=self.vocab_size       ,						hidden_size=self.hidden_size       ,						embedding_size=self.embedding_size       ,						num_hidden_layers=self.num_hidden_layers       ,						num_attention_heads=self.num_attention_heads       ,						intermediate_size=self.intermediate_size       ,						hidden_act=self.hidden_act       ,						hidden_dropout_prob=self.hidden_dropout_prob       ,						attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,						max_position_embeddings=self.max_position_embeddings       ,						type_vocab_size=self.type_vocab_size       ,						initializer_range=self.initializer_range       ,						)

         return config, input_ids, token_type_ids, attention_mask



  def 						UpperCamelCase_ (					self )	->	List[str]:



         """simple docstring"""





         UpperCAmelCase					=	self.prepare_config_and_inputs()
         UpperCAmelCase			,				UpperCAmelCase			,				UpperCAmelCase			,				UpperCAmelCase					=	config_and_inputs
         UpperCAmelCase					=	{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": attention_mask}
         return config, inputs_dict





@require_flax
class      UpperCamelCase_       (				a_      ,    unittest.TestCase     ):
  _A				:    str =		True

  _A				:    List[Any] =		(
      (
          FlaxElectraModel,
          FlaxElectraForCausalLM,
          FlaxElectraForMaskedLM,
          FlaxElectraForPreTraining,
          FlaxElectraForTokenClassification,
          FlaxElectraForQuestionAnswering,
          FlaxElectraForMultipleChoice,
          FlaxElectraForSequenceClassification,
      )
      if is_flax_available()
      else ()
  )



  def 						UpperCamelCase_ (					self )	->	str:



         """simple docstring"""





         UpperCAmelCase					=	FlaxElectraModelTester(self )



  @slow
  def 						UpperCamelCase_ (					self )	->	Tuple:



         """simple docstring"""





         for model_class_name in self.all_model_classes:
                if model_class_name == FlaxElectraForMaskedLM:
                       UpperCAmelCase					=	model_class_name.from_pretrained("""google/electra-small-generator""" )
                else:
                       UpperCAmelCase					=	model_class_name.from_pretrained("""google/electra-small-discriminator""" )
                UpperCAmelCase					=	model(np.ones((1, 1) ) )
                self.assertIsNotNone(snake_case__ )
