



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



						"""simple docstring"""






						def __init__(   self					:List[str]  ,    lowercase_					:Tuple  ,    lowercase_					:Union[str, Any]=13  ,    lowercase_					:int=7  ,    lowercase_					:Optional[int]=True  ,    lowercase_					:Any=True  ,    lowercase_					:Tuple=True  ,    lowercase_					:List[str]=True  ,    lowercase_					:Any=99  ,    lowercase_					:Optional[int]=24  ,    lowercase_					:Optional[int]=32  ,    lowercase_					:int=5  ,    lowercase_					:str=4  ,    lowercase_					:Tuple=37  ,    lowercase_					:Optional[int]="gelu"  ,    lowercase_					:Union[str, Any]=0.1  ,    lowercase_					:Dict=0.1  ,    lowercase_					:Optional[int]=5_12  ,    lowercase_					:Optional[int]=16  ,    lowercase_					:List[Any]=2  ,    lowercase_					:str=0.02  ,    lowercase_					:Optional[Any]=4  ,    )							->     Optional[int]:
												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[int]   )							->     List[str]:
												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					:Tuple   )							->     Union[str, Any]:
												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  A_     (	SCREAMING_SNAKE_CASE_      ,  unittest.TestCase							):



						"""simple docstring"""
						__UpperCamelCase													=				True

						__UpperCamelCase													=				(
						    (
						        FlaxElectraModel,
						        FlaxElectraForCausalLM,
						        FlaxElectraForMaskedLM,
						        FlaxElectraForPreTraining,
						        FlaxElectraForTokenClassification,
						        FlaxElectraForQuestionAnswering,
						        FlaxElectraForMultipleChoice,
						        FlaxElectraForSequenceClassification,
						    )
						    if is_flax_available()
						    else ()
						)






						def       UpperCAmelCase__	(   self					:List[str]   )							->     str:
												UpperCAmelCase			 =					FlaxElectraModelTester(self   )






						@slow
						def       UpperCAmelCase__	(   self					:Any   )							->     Optional[Any]:
												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(lowercase_   )












'''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  ,  lowercase__  ,  lowercase__=13  ,  lowercase__=64  ,  lowercase__=4  ,  lowercase__=3  ,  lowercase__=16  ,  lowercase__=[1, 2, 1]  ,  lowercase__=[2, 4, 8]  ,  lowercase__=3  ,  lowercase__=2.0  ,  lowercase__=True  ,  lowercase__=0.0  ,  lowercase__=0.0  ,  lowercase__=0.1  ,  lowercase__="gelu"  ,  lowercase__=True  ,  lowercase__=0.02  ,  lowercase__=1E-5  ,  lowercase__=True  ,  lowercase__=None  ,  lowercase__=True  ,  lowercase__=10  ,  lowercase__=["stage1", "stage2"]  ,  lowercase__=[1, 2]  ,  )		->    List[str]:






        """simple docstring"""





        _snake_case		:						Optional[int]         =  parent
        _snake_case		:						List[str]         =  batch_size
        _snake_case		:						int         =  image_size
        _snake_case		:						Union[str, Any]         =  patch_size
        _snake_case		:						List[Any]         =  num_channels
        _snake_case		:						Tuple         =  embed_dim
        _snake_case		:						Tuple         =  depths
        _snake_case		:						Optional[Any]         =  num_heads
        _snake_case		:						Union[str, Any]         =  kernel_size
        _snake_case		:						Optional[Any]         =  mlp_ratio
        _snake_case		:						Optional[int]         =  qkv_bias
        _snake_case		:						List[Any]         =  hidden_dropout_prob
        _snake_case		:						Optional[int]         =  attention_probs_dropout_prob
        _snake_case		:						List[str]         =  drop_path_rate
        _snake_case		:						Optional[Any]         =  hidden_act
        _snake_case		:						Any         =  patch_norm
        _snake_case		:						Any         =  layer_norm_eps
        _snake_case		:						Union[str, Any]         =  initializer_range
        _snake_case		:						List[Any]         =  is_training
        _snake_case		:						List[Any]         =  scope
        _snake_case		:						Dict         =  use_labels
        _snake_case		:						Tuple         =  num_labels
        _snake_case		:						Union[str, Any]         =  out_features
        _snake_case		:						Dict         =  out_indices







  def 					UpperCAmelCase_					(		self		)		->    int:






        """simple docstring"""





        _snake_case		:						List[str]         =  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)

        _snake_case		:						Union[str, Any]         =  None
        if self.use_labels:
              _snake_case		:						int         =  ids_tensor([self.batch_size]  ,  self.num_labels		)

        _snake_case		:						Union[str, Any]         =  self.get_config()

        return config, pixel_values, labels







  def 					UpperCAmelCase_					(		self		)		->    Optional[int]:






        """simple docstring"""





        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  ,  lowercase__  ,  lowercase__  ,  lowercase__		)		->    str:






        """simple docstring"""





        _snake_case		:						int         =  NatModel(config=lowercase__		)
        model.to(lowercase__		)
        model.eval()
        _snake_case		:						Optional[Any]         =  model(lowercase__		)

        _snake_case		:						Optional[int]         =  (config.image_size // config.patch_size) // (2 ** (len(config.depths		) - 1))
        _snake_case		:						str         =  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  ,  lowercase__  ,  lowercase__  ,  lowercase__		)		->    Tuple:






        """simple docstring"""





        _snake_case		:						int         =  NatForImageClassification(lowercase__		)
        model.to(lowercase__		)
        model.eval()
        _snake_case		:						List[str]         =  model(lowercase__  ,  labels=lowercase__		)
        self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.num_labels)		)

        # test greyscale images
        _snake_case		:						str         =  1
        _snake_case		:						Dict         =  NatForImageClassification(lowercase__		)
        model.to(lowercase__		)
        model.eval()

        _snake_case		:						int         =  floats_tensor([self.batch_size, 1, self.image_size, self.image_size]		)
        _snake_case		:						Optional[Any]         =  model(lowercase__		)
        self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.num_labels)		)







  def 					UpperCAmelCase_					(		self  ,  lowercase__  ,  lowercase__  ,  lowercase__		)		->    Union[str, Any]:






        """simple docstring"""





        _snake_case		:						int         =  NatBackbone(config=lowercase__		)
        model.to(lowercase__		)
        model.eval()
        _snake_case		:						Any         =  model(lowercase__		)

        # 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
        _snake_case		:						Dict         =  None
        _snake_case		:						Any         =  NatBackbone(config=lowercase__		)
        model.to(lowercase__		)
        model.eval()
        _snake_case		:						str         =  model(lowercase__		)

        # 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		)		->    Any:






        """simple docstring"""





        _snake_case		:						Union[str, Any]         =  self.prepare_config_and_inputs()
        _snake_case  ,				_snake_case  ,				_snake_case		:						Optional[Any]         =  config_and_inputs
        _snake_case		:						List[Any]         =  {'''pixel_values''': pixel_values}
        return config, inputs_dict





@require_natten
@require_torch
class lowerCamelCase    (a__    ,   a__    ,   unittest.TestCase       ):
  _lowercase	:						int							       =					(
      (
          NatModel,
          NatForImageClassification,
          NatBackbone,
      )
      if is_torch_available()
      else ()
  )
  _lowercase	:						Optional[Any]							       =					(
      {"""feature-extraction""": NatModel, """image-classification""": NatForImageClassification}
      if is_torch_available()
      else {}
  )
  _lowercase	:						List[str]							       =					False

  _lowercase	:						str							       =					False
  _lowercase	:						Dict							       =					False
  _lowercase	:						Tuple							       =					False
  _lowercase	:						List[str]							       =					False







  def 					UpperCAmelCase_					(		self		)		->    List[Any]:






        """simple docstring"""





        _snake_case		:						List[str]         =  NatModelTester(self		)
        _snake_case		:						List[str]         =  ConfigTester(self  ,  config_class=lowercase__  ,  embed_dim=37		)







  def 					UpperCAmelCase_					(		self		)		->    Union[str, Any]:






        """simple docstring"""





        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		)		->    int:






        """simple docstring"""





        return







  def 					UpperCAmelCase_					(		self		)		->    Optional[int]:






        """simple docstring"""





        _snake_case		:						Tuple         =  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*lowercase__		)







  def 					UpperCAmelCase_					(		self		)		->    Dict:






        """simple docstring"""





        _snake_case		:						Optional[int]         =  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*lowercase__		)







  def 					UpperCAmelCase_					(		self		)		->    Any:






        """simple docstring"""





        _snake_case		:						List[Any]         =  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_backbone(*lowercase__		)







  @unittest.skip(reason='''Nat does not use inputs_embeds'''		)
  def 					UpperCAmelCase_					(		self		)		->    Optional[int]:






        """simple docstring"""





        pass







  @unittest.skip(reason='''Nat does not use feedforward chunking'''		)
  def 					UpperCAmelCase_					(		self		)		->    List[Any]:






        """simple docstring"""





        pass







  def 					UpperCAmelCase_					(		self		)		->    str:






        """simple docstring"""





        _snake_case  ,				_snake_case		:						Tuple         =  self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
              _snake_case		:						List[Any]         =  model_class(lowercase__		)
              self.assertIsInstance(model.get_input_embeddings()  ,  (nn.Module)		)
              _snake_case		:						str         =  model.get_output_embeddings()
              self.assertTrue(x is None or isinstance(lowercase__  ,  nn.Linear		)		)







  def 					UpperCAmelCase_					(		self		)		->    int:






        """simple docstring"""





        _snake_case  ,				_snake_case		:						Dict         =  self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
              _snake_case		:						Union[str, Any]         =  model_class(lowercase__		)
              _snake_case		:						Tuple         =  inspect.signature(model.forward		)
              # signature.parameters is an OrderedDict => so arg_names order is deterministic
              _snake_case		:						List[Any]         =  [*signature.parameters.keys()]

              _snake_case		:						str         =  ['''pixel_values''']
              self.assertListEqual(arg_names[:1]  ,  lowercase__		)







  def 					UpperCAmelCase_					(		self		)		->    List[Any]:






        """simple docstring"""





        self.skipTest('''Nat\'s attention operation is handled entirely by NATTEN.'''		)







  def 					UpperCAmelCase_					(		self  ,  lowercase__  ,  lowercase__  ,  lowercase__  ,  lowercase__		)		->    Any:






        """simple docstring"""





        _snake_case		:						str         =  model_class(lowercase__		)
        model.to(lowercase__		)
        model.eval()

        with torch.no_grad():
              _snake_case		:						Optional[Any]         =  model(**self._prepare_for_class(lowercase__  ,  lowercase__		)		)

        _snake_case		:						Any         =  outputs.hidden_states

        _snake_case		:						Optional[Any]         =  getattr(
            self.model_tester  ,  '''expected_num_hidden_layers'''  ,  len(self.model_tester.depths		) + 1		)
        self.assertEqual(len(lowercase__		)  ,  lowercase__		)

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

        _snake_case		:						int         =  image_size[0] // patch_size[0]
        _snake_case		:						Optional[Any]         =  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":
              _snake_case		:						Dict         =  outputs.reshaped_hidden_states
              self.assertEqual(len(lowercase__		)  ,  lowercase__		)

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







  def 					UpperCAmelCase_					(		self		)		->    Union[str, Any]:






        """simple docstring"""





        _snake_case  ,				_snake_case		:						str         =  self.model_tester.prepare_config_and_inputs_for_common()

        _snake_case		:						Dict         =  (
            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:
              _snake_case		:						Optional[Any]         =  True
              self.check_hidden_states_output(lowercase__  ,  lowercase__  ,  lowercase__  ,  lowercase__		)

              # check that output_hidden_states also work using config
              del inputs_dict["output_hidden_states"]
              _snake_case		:						str         =  True

              self.check_hidden_states_output(lowercase__  ,  lowercase__  ,  lowercase__  ,  lowercase__		)







  @slow
  def 					UpperCAmelCase_					(		self		)		->    Dict:






        """simple docstring"""





        for model_name in NAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              _snake_case		:						Optional[int]         =  NatModel.from_pretrained(lowercase__		)
              self.assertIsNotNone(lowercase__		)







  def 					UpperCAmelCase_					(		self		)		->    Dict:






        """simple docstring"""





        _snake_case  ,				_snake_case		:						Tuple         =  self.model_tester.prepare_config_and_inputs_for_common()

        _snake_case		:						str         =  _config_zero_init(lowercase__		)
        for model_class in self.all_model_classes:
              _snake_case		:						int         =  model_class(config=lowercase__		)
              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		)		->    int:






        """simple docstring"""





        return AutoImageProcessor.from_pretrained('''shi-labs/nat-mini-in1k-224'''		) if is_vision_available() else None







  @slow
  def 					UpperCAmelCase_					(		self		)		->    Dict:






        """simple docstring"""





        _snake_case		:						Tuple         =  NatForImageClassification.from_pretrained('''shi-labs/nat-mini-in1k-224'''		).to(lowercase__		)
        _snake_case		:						Optional[int]         =  self.default_image_processor

        _snake_case		:						List[str]         =  Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''		)
        _snake_case		:						List[Any]         =  image_processor(images=lowercase__  ,  return_tensors='''pt'''		).to(lowercase__		)

        # forward pass
        with torch.no_grad():
              _snake_case		:						Tuple         =  model(**lowercase__		)

        # verify the logits
        _snake_case		:						List[str]         =  torch.Size((1, 1_000)		)
        self.assertEqual(outputs.logits.shape  ,  lowercase__		)
        _snake_case		:						List[str]         =  torch.tensor([0.3_805, -0.8_676, -0.3_912]		).to(lowercase__		)
        self.assertTrue(torch.allclose(outputs.logits[0, :3]  ,  lowercase__  ,  atol=1E-4		)		)





@require_torch
@require_natten
class lowerCamelCase    (unittest.TestCase    ,   a__       ):
  _lowercase	:						str							       =					(NatBackbone,) if is_torch_available() else ()
  _lowercase	:						List[Any]							       =					NatConfig







  def 					UpperCAmelCase_					(		self		)		->    Dict:






        """simple docstring"""





        _snake_case		:						Any         =  NatModelTester(self		)


