
'''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 		_a     :







   """simple docstring"""




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






   def 					__a			(       self      ):
    SCREAMING_SNAKE_CASE  :  Any								=		floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

    SCREAMING_SNAKE_CASE  :  str								=		None
    if self.use_labels:
     SCREAMING_SNAKE_CASE  :  List[str]								=		ids_tensor([self.batch_size]    ,self.num_labels      )

    SCREAMING_SNAKE_CASE  :  Tuple								=		self.get_config()

    return config, pixel_values, labels






   def 					__a			(       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 					__a			(       self    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE      ):
    SCREAMING_SNAKE_CASE  :  str								=		NatModel(config=lowercase_      )
    model.to(lowercase_      )
    model.eval()
    SCREAMING_SNAKE_CASE  :  Union[str, Any]								=		model(lowercase_      )

    SCREAMING_SNAKE_CASE  :  Optional[int]								=		(config.image_size // config.patch_size) // (2 ** (len(config.depths      ) - 1))
    SCREAMING_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 					__a			(       self    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE      ):
    SCREAMING_SNAKE_CASE  :  int								=		NatForImageClassification(lowercase_      )
    model.to(lowercase_      )
    model.eval()
    SCREAMING_SNAKE_CASE  :  List[Any]								=		model(lowercase_    ,labels=lowercase_      )
    self.parent.assertEqual(result.logits.shape    ,(self.batch_size, self.num_labels)      )

    # test greyscale images
    SCREAMING_SNAKE_CASE  :  Any								=		1
    SCREAMING_SNAKE_CASE  :  Optional[int]								=		NatForImageClassification(lowercase_      )
    model.to(lowercase_      )
    model.eval()

    SCREAMING_SNAKE_CASE  :  Tuple								=		floats_tensor([self.batch_size, 1, self.image_size, self.image_size]      )
    SCREAMING_SNAKE_CASE  :  Tuple								=		model(lowercase_      )
    self.parent.assertEqual(result.logits.shape    ,(self.batch_size, self.num_labels)      )






   def 					__a			(       self    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE      ):
    SCREAMING_SNAKE_CASE  :  int								=		NatBackbone(config=lowercase_      )
    model.to(lowercase_      )
    model.eval()
    SCREAMING_SNAKE_CASE  :  List[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
    SCREAMING_SNAKE_CASE  :  Optional[Any]								=		None
    SCREAMING_SNAKE_CASE  :  str								=		NatBackbone(config=lowercase_      )
    model.to(lowercase_      )
    model.eval()
    SCREAMING_SNAKE_CASE  :  List[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 					__a			(       self      ):
    SCREAMING_SNAKE_CASE  :  Union[str, Any]								=		self.prepare_config_and_inputs()
    SCREAMING_SNAKE_CASE  :  int								=		config_and_inputs
    SCREAMING_SNAKE_CASE  :  Dict								=		{"pixel_values": pixel_values}
    return config, inputs_dict



@require_natten
@require_torch
class 		_a     (							__A			,    __A			,    unittest.TestCase  ):







   """simple docstring"""




   A		= (
       (
           NatModel,
           NatForImageClassification,
           NatBackbone,
       )
       if is_torch_available()
       else ()
   )
   A		= (
       {"feature-extraction": NatModel, "image-classification": NatForImageClassification}
       if is_torch_available()
       else {}
   )
   A		= False

   A		= False
   A		= False
   A		= False
   A		= False






   def 					__a			(       self      ):
    SCREAMING_SNAKE_CASE  :  Any								=		NatModelTester(self      )
    SCREAMING_SNAKE_CASE  :  Union[str, Any]								=		ConfigTester(self    ,config_class=lowercase_    ,embed_dim=37      )






   def 					__a			(       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 					__a			(       self      ):
    return






   def 					__a			(       self      ):
    SCREAMING_SNAKE_CASE  :  List[Any]								=		self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_model(*lowercase_      )






   def 					__a			(       self      ):
    SCREAMING_SNAKE_CASE  :  List[Any]								=		self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_image_classification(*lowercase_      )






   def 					__a			(       self      ):
    SCREAMING_SNAKE_CASE  :  List[str]								=		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 					__a			(       self      ):
    pass






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






   def 					__a			(       self      ):
    SCREAMING_SNAKE_CASE  :  Any								=		self.model_tester.prepare_config_and_inputs_for_common()

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






   def 					__a			(       self      ):
    SCREAMING_SNAKE_CASE  :  Optional[int]								=		self.model_tester.prepare_config_and_inputs_for_common()

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

     SCREAMING_SNAKE_CASE  :  Dict								=		["pixel_values"]
     self.assertListEqual(arg_names[:1]    ,lowercase_      )






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






   def 					__a			(       self    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE      ):
    SCREAMING_SNAKE_CASE  :  List[Any]								=		model_class(lowercase_      )
    model.to(lowercase_      )
    model.eval()

    with torch.no_grad():
     SCREAMING_SNAKE_CASE  :  List[str]								=		model(**self._prepare_for_class(lowercase_    ,lowercase_      )      )

    SCREAMING_SNAKE_CASE  :  Tuple								=		outputs.hidden_states

    SCREAMING_SNAKE_CASE  :  List[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
    SCREAMING_SNAKE_CASE  :  int								=		(
        config.patch_size
        if isinstance(config.patch_size    ,collections.abc.Iterable      )
        else (config.patch_size, config.patch_size)
    )

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

     SCREAMING_SNAKE_CASE  :  Union[str, Any]								=		reshaped_hidden_states[0].shape
     SCREAMING_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 					__a			(       self      ):
    SCREAMING_SNAKE_CASE  :  List[Any]								=		self.model_tester.prepare_config_and_inputs_for_common()

    SCREAMING_SNAKE_CASE  :  int								=		(
        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:
     SCREAMING_SNAKE_CASE  :  str								=		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"]
     SCREAMING_SNAKE_CASE  :  Optional[Any]								=		True

     self.check_hidden_states_output(lowercase_    ,lowercase_    ,lowercase_    ,lowercase_      )






   @slow
   def 					__a			(       self      ):
    for model_name in NAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     SCREAMING_SNAKE_CASE  :  Optional[Any]								=		NatModel.from_pretrained(lowercase_      )
     self.assertIsNotNone(lowercase_      )






   def 					__a			(       self      ):
    SCREAMING_SNAKE_CASE  :  Any								=		self.model_tester.prepare_config_and_inputs_for_common()

    SCREAMING_SNAKE_CASE  :  Optional[Any]								=		_config_zero_init(lowercase_      )
    for model_class in self.all_model_classes:
     SCREAMING_SNAKE_CASE  :  List[Any]								=		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 		_a     (							unittest.TestCase  ):







   """simple docstring"""




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






   @slow
   def 					__a			(       self      ):
    SCREAMING_SNAKE_CASE  :  Optional[Any]								=		NatForImageClassification.from_pretrained('shi-labs/nat-mini-in1k-224'      ).to(lowercase_      )
    SCREAMING_SNAKE_CASE  :  Optional[int]								=		self.default_image_processor

    SCREAMING_SNAKE_CASE  :  str								=		Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'      )
    SCREAMING_SNAKE_CASE  :  Dict								=		image_processor(images=lowercase_    ,return_tensors='pt'      ).to(lowercase_      )

    # forward pass
    with torch.no_grad():
     SCREAMING_SNAKE_CASE  :  Union[str, Any]								=		model(**lowercase_      )

    # verify the logits
    SCREAMING_SNAKE_CASE  :  Dict								=		torch.Size((1, 1000)      )
    self.assertEqual(outputs.logits.shape    ,lowercase_      )
    SCREAMING_SNAKE_CASE  :  Dict								=		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 		_a     (							unittest.TestCase			,    __A  ):







   """simple docstring"""




   A		= (NatBackbone,) if is_torch_available() else ()
   A		= NatConfig






   def 					__a			(       self      ):
    SCREAMING_SNAKE_CASE  :  Dict								=		NatModelTester(self      )



def      lowercase_							(     _lowerCamelCase				:			str):
    lowercase__					:  Union[str, Any]          =							0
    for plain_chr in plain_text:
        lowercase__					:  Union[str, Any]          =							(
            ord(_lowerCamelCase) + (hash_value << 6) + (hash_value << 16) - hash_value
        )
    return hash_value



