





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  :Optional[Any]    ,							_lowercase  :str    ,							_lowercase  :Tuple=13    ,							_lowercase  :Tuple=64    ,							_lowercase  :Tuple=4    ,							_lowercase  :List[str]=3    ,							_lowercase  :Any=16    ,							_lowercase  :Optional[Any]=[1, 2, 1]    ,							_lowercase  :Tuple=[2, 4, 8]    ,							_lowercase  :int=3    ,							_lowercase  :Optional[Any]=2.0    ,							_lowercase  :Optional[Any]=True    ,							_lowercase  :Any=0.0    ,							_lowercase  :List[Any]=0.0    ,							_lowercase  :List[str]=0.1    ,							_lowercase  :Optional[int]="gelu"    ,							_lowercase  :Optional[int]=True    ,							_lowercase  :str=0.02    ,							_lowercase  :Union[str, Any]=1e-5    ,							_lowercase  :Optional[Any]=True    ,							_lowercase  :Dict=None    ,							_lowercase  :str=True    ,							_lowercase  :List[str]=10    ,							_lowercase  :str=["stage1", "stage2"]    ,							_lowercase  :Optional[Any]=[1, 2]    ,							):






     '''simple docstring'''





     lowercase__    =  parent
     lowercase__    =  batch_size
     lowercase__    =  image_size
     lowercase__    =  patch_size
     lowercase__    =  num_channels
     lowercase__    =  embed_dim
     lowercase__    =  depths
     lowercase__    =  num_heads
     lowercase__    =  kernel_size
     lowercase__    =  mlp_ratio
     lowercase__    =  qkv_bias
     lowercase__    =  hidden_dropout_prob
     lowercase__    =  attention_probs_dropout_prob
     lowercase__    =  drop_path_rate
     lowercase__    =  hidden_act
     lowercase__    =  patch_norm
     lowercase__    =  layer_norm_eps
     lowercase__    =  initializer_range
     lowercase__    =  is_training
     lowercase__    =  scope
     lowercase__    =  use_labels
     lowercase__    =  num_labels
     lowercase__    =  out_features
     lowercase__    =  out_indices







  def   UpperCAmelCase     (			self  :str					):






     '''simple docstring'''





     lowercase__    =  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

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

     lowercase__    =  self.get_config()

     return config, pixel_values, labels







  def   UpperCAmelCase     (			self  :str					):






     '''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  :List[str]    ,							_lowercase  :Dict    ,							_lowercase  :Dict    ,							_lowercase  :Any					):






     '''simple docstring'''





     lowercase__    =  NatModel(config=lowercase__					)
     model.to(lowercase__					)
     model.eval()
     lowercase__    =  model(lowercase__					)

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






     '''simple docstring'''





     lowercase__    =  NatForImageClassification(lowercase__					)
     model.to(lowercase__					)
     model.eval()
     lowercase__    =  model(lowercase__    ,							labels=lowercase__					)
     self.parent.assertEqual(result.logits.shape    ,							(self.batch_size, self.num_labels)					)

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

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







  def   UpperCAmelCase     (			self  :Dict    ,							_lowercase  :Optional[Any]    ,							_lowercase  :Optional[int]    ,							_lowercase  :List[str]					):






     '''simple docstring'''





     lowercase__    =  NatBackbone(config=lowercase__					)
     model.to(lowercase__					)
     model.eval()
     lowercase__    =  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
     lowercase__    =  None
     lowercase__    =  NatBackbone(config=lowercase__					)
     model.to(lowercase__					)
     model.eval()
     lowercase__    =  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  :Optional[Any]					):






     '''simple docstring'''





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





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

  __lowerCamelCase		=			False
  __lowerCamelCase		=			False
  __lowerCamelCase		=			False
  __lowerCamelCase		=			False







  def   UpperCAmelCase     (			self  :Any					):






     '''simple docstring'''





     lowercase__    =  NatModelTester(self					)
     lowercase__    =  ConfigTester(self    ,							config_class=lowercase__    ,							embed_dim=37					)







  def   UpperCAmelCase     (			self  :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  :Optional[Any]					):






     '''simple docstring'''





     return







  def   UpperCAmelCase     (			self  :Optional[Any]					):






     '''simple docstring'''





     lowercase__    =  self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*lowercase__					)







  def   UpperCAmelCase     (			self  :str					):






     '''simple docstring'''





     lowercase__    =  self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_classification(*lowercase__					)







  def   UpperCAmelCase     (			self  :Union[str, Any]					):






     '''simple docstring'''





     lowercase__    =  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[Any]					):






     '''simple docstring'''





     pass







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






     '''simple docstring'''





     pass







  def   UpperCAmelCase     (			self  :Optional[int]					):






     '''simple docstring'''





     lowercase__   ,							lowercase__    =  self.model_tester.prepare_config_and_inputs_for_common()

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







  def   UpperCAmelCase     (			self  :List[str]					):






     '''simple docstring'''





     lowercase__   ,							lowercase__    =  self.model_tester.prepare_config_and_inputs_for_common()

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

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







  def   UpperCAmelCase     (			self  :Dict					):






     '''simple docstring'''





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







  def   UpperCAmelCase     (			self  :Any    ,							_lowercase  :Tuple    ,							_lowercase  :Tuple    ,							_lowercase  :Tuple    ,							_lowercase  :int					):






     '''simple docstring'''





     lowercase__    =  model_class(lowercase__					)
     model.to(lowercase__					)
     model.eval()

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

     lowercase__    =  outputs.hidden_states

     lowercase__    =  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
     lowercase__    =  (
         config.patch_size
         if isinstance(config.patch_size    ,							collections.abc.Iterable					)
         else (config.patch_size, config.patch_size)
     )

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

        lowercase__   ,							lowercase__   ,							lowercase__   ,							lowercase__    =  reshaped_hidden_states[0].shape
        lowercase__    =  (
            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  :Any					):






     '''simple docstring'''





     lowercase__   ,							lowercase__    =  self.model_tester.prepare_config_and_inputs_for_common()

     lowercase__    =  (
         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:
        lowercase__    =  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"]
        lowercase__    =  True

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







  @slow
  def   UpperCAmelCase     (			self  :List[Any]					):






     '''simple docstring'''





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







  def   UpperCAmelCase     (			self  :Dict					):






     '''simple docstring'''





     lowercase__   ,							lowercase__    =  self.model_tester.prepare_config_and_inputs_for_common()

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






     '''simple docstring'''





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







  @slow
  def   UpperCAmelCase     (			self  :int					):






     '''simple docstring'''





     lowercase__    =  NatForImageClassification.from_pretrained("shi-labs/nat-mini-in1k-224"					).to(lowercase__					)
     lowercase__    =  self.default_image_processor

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

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

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





@require_torch
@require_natten
class 					lowerCAmelCase							(   unittest.TestCase							,							lowercase_					):
  __lowerCamelCase		=			(NatBackbone,) if is_torch_available() else ()
  __lowerCamelCase		=			NatConfig







  def   UpperCAmelCase     (			self  :List[str]					):






     '''simple docstring'''





     lowercase__    =  NatModelTester(self					)







from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCAmelCase__							= logging.get_logger(__name__)

lowerCAmelCase__							= {
    'dandelin/vilt-b32-mlm': 'https://huggingface.co/dandelin/vilt-b32-mlm/blob/main/config.json'
}
class 				a_ (						SCREAMING_SNAKE_CASE						):






 '''simple docstring'''






 UpperCAmelCase_	 =					'vilt'




 def __init__(			self					:  Optional[Any]			,    lowercase__					:  Dict=30_522			,    lowercase__					:  Tuple=2			,    lowercase__					:  int=2			,    lowercase__					:  Optional[int]=40			,    lowercase__					:  Optional[int]=768			,    lowercase__					:  Tuple=12			,    lowercase__					:  Tuple=12			,    lowercase__					:  Union[str, Any]=3_072			,    lowercase__					:  int="gelu"			,    lowercase__					:  Tuple=0.0			,    lowercase__					:  Dict=0.0			,    lowercase__					:  Optional[int]=0.02			,    lowercase__					:  Tuple=1e-12			,    lowercase__					:  Dict=384			,    lowercase__					:  int=32			,    lowercase__					:  List[str]=3			,    lowercase__					:  int=True			,    lowercase__					:  List[str]=-1			,    lowercase__					:  Union[str, Any]=False			,    lowercase__					:  List[str]=-1			,    **lowercase__					:  int			,    ):







        '''simple docstring'''



        super().__init__(tie_word_embeddings=lowercase__			,    **lowercase__)

        lowerCAmelCase__     	=					vocab_size
        lowerCAmelCase__     	=					type_vocab_size
        lowerCAmelCase__     	=					modality_type_vocab_size
        lowerCAmelCase__     	=					max_position_embeddings

        lowerCAmelCase__     	=					hidden_size
        lowerCAmelCase__     	=					num_hidden_layers
        lowerCAmelCase__     	=					num_attention_heads
        lowerCAmelCase__     	=					intermediate_size
        lowerCAmelCase__     	=					hidden_act
        lowerCAmelCase__     	=					hidden_dropout_prob
        lowerCAmelCase__     	=					attention_probs_dropout_prob
        lowerCAmelCase__     	=					initializer_range
        lowerCAmelCase__     	=					layer_norm_eps

        lowerCAmelCase__     	=					image_size
        lowerCAmelCase__     	=					patch_size
        lowerCAmelCase__     	=					num_channels
        lowerCAmelCase__     	=					qkv_bias
        lowerCAmelCase__     	=					max_image_length
        lowerCAmelCase__     	=					num_images
