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  _UpperCAmelCase       :



 '''simple docstring'''




 def __init__(						self							:					Dict				,      lowercase_							:					Optional[Any]				,      lowercase_							:					Optional[int]=13				,      lowercase_							:					List[Any]=64				,      lowercase_							:					List[str]=4				,      lowercase_							:					Any=3				,      lowercase_							:					Any=16				,      lowercase_							:					Tuple=[1, 2, 1]				,      lowercase_							:					Optional[Any]=[2, 4, 8]				,      lowercase_							:					str=3				,      lowercase_							:					int=2.0				,      lowercase_							:					Optional[Any]=True				,      lowercase_							:					Union[str, Any]=0.0				,      lowercase_							:					int=0.0				,      lowercase_							:					str=0.1				,      lowercase_							:					int="gelu"				,      lowercase_							:					List[Any]=True				,      lowercase_							:					str=0.02				,      lowercase_							:					Optional[int]=1e-5				,      lowercase_							:					Dict=True				,      lowercase_							:					List[str]=None				,      lowercase_							:					Optional[Any]=True				,      lowercase_							:					Union[str, Any]=10				,      lowercase_							:					int=["stage1", "stage2"]				,      lowercase_							:					Dict=[1, 2]				,      )		->     str:


   """simple docstring"""







   _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							:					Optional[int])		->     int:


   """simple docstring"""







   _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							:					Any)		->     List[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							:					Any				,      lowercase_							:					Tuple				,      lowercase_							:					int				,      lowercase_							:					Union[str, Any])		->     Dict:


   """simple docstring"""







   _UpperCamelCase				=  NatModel(config=lowercase_)
   model.to(lowercase_)
   model.eval()
   _UpperCamelCase				=  model(lowercase_)

   _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							:					Dict				,      lowercase_							:					Optional[int]				,      lowercase_							:					Tuple				,      lowercase_							:					List[str])		->     List[str]:


   """simple docstring"""







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

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

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







 def 						__UpperCAmelCase							(						self							:					int				,      lowercase_							:					Optional[Any]				,      lowercase_							:					Optional[int]				,      lowercase_							:					Union[str, Any])		->     str:


   """simple docstring"""







   _UpperCamelCase				=  NatBackbone(config=lowercase_)
   model.to(lowercase_)
   model.eval()
   _UpperCamelCase				=  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
   _UpperCamelCase				=  None
   _UpperCamelCase				=  NatBackbone(config=lowercase_)
   model.to(lowercase_)
   model.eval()
   _UpperCamelCase				=  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							:					Dict)		->     List[Any]:


   """simple docstring"""







   _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  _UpperCAmelCase       (			lowerCAmelCase,       lowerCAmelCase,       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 						__UpperCAmelCase							(						self							:					int)		->     List[Any]:


   """simple docstring"""







   _UpperCamelCase				=  NatModelTester(self)
   _UpperCamelCase				=  ConfigTester(self				,      config_class=lowercase_				,      embed_dim=37)







 def 						__UpperCAmelCase							(						self							:					Any)		->     List[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							:					List[str])		->     Dict:


   """simple docstring"""







   return







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


   """simple docstring"""







   _UpperCamelCase				=  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*lowercase_)







 def 						__UpperCAmelCase							(						self							:					List[str])		->     Tuple:


   """simple docstring"""







   _UpperCamelCase				=  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_image_classification(*lowercase_)







 def 						__UpperCAmelCase							(						self							:					Dict)		->     Dict:


   """simple docstring"""







   _UpperCamelCase				=  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							:					List[str])		->     List[Any]:


   """simple docstring"""







   pass







 @unittest.skip(reason="Nat does not use feedforward chunking")
 def 						__UpperCAmelCase							(						self							:					int)		->     Dict:


   """simple docstring"""







   pass







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


   """simple docstring"""







   _UpperCamelCase	,		_UpperCamelCase				=  self.model_tester.prepare_config_and_inputs_for_common()

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







 def 						__UpperCAmelCase							(						self							:					Optional[int])		->     Optional[Any]:


   """simple docstring"""







   _UpperCamelCase	,		_UpperCamelCase				=  self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
     _UpperCamelCase				=  model_class(lowercase_)
     _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]				,      lowercase_)







 def 						__UpperCAmelCase							(						self							:					Optional[int])		->     Dict:


   """simple docstring"""







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







 def 						__UpperCAmelCase							(						self							:					List[Any]				,      lowercase_							:					Union[str, Any]				,      lowercase_							:					int				,      lowercase_							:					List[str]				,      lowercase_							:					List[str])		->     List[Any]:


   """simple docstring"""







   _UpperCamelCase				=  model_class(lowercase_)
   model.to(lowercase_)
   model.eval()

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

   _UpperCamelCase				=  outputs.hidden_states

   _UpperCamelCase				=  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
   _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(lowercase_)				,      lowercase_)

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


   """simple docstring"""







   _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(lowercase_				,      lowercase_				,      lowercase_				,      lowercase_)

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

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







 @slow
 def 						__UpperCAmelCase							(						self							:					str)		->     Dict:


   """simple docstring"""







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







 def 						__UpperCAmelCase							(						self							:					str)		->     str:


   """simple docstring"""







   _UpperCamelCase	,		_UpperCamelCase				=  self.model_tester.prepare_config_and_inputs_for_common()

   _UpperCamelCase				=  _config_zero_init(lowercase_)
   for model_class in self.all_model_classes:
     _UpperCamelCase				=  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  _UpperCAmelCase       (			unittest.TestCase   ):



 '''simple docstring'''




 @cached_property
 def 						__UpperCAmelCase							(						self							:					Any)		->     Tuple:


   """simple docstring"""







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







 @slow
 def 						__UpperCAmelCase							(						self							:					Optional[Any])		->     List[str]:


   """simple docstring"""







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

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

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

   # verify the logits
   _UpperCamelCase				=  torch.Size((1, 1000))
   self.assertEqual(outputs.logits.shape				,      lowercase_)
   _UpperCamelCase				=  torch.tensor([0.38_05, -0.86_76, -0.39_12]).to(lowercase_)
   self.assertTrue(torch.allclose(outputs.logits[0, :3]				,      lowercase_				,      atol=1e-4))




@require_torch
@require_natten
class  _UpperCAmelCase       (			unittest.TestCase,       lowerCAmelCase   ):



 '''simple docstring'''




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







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


   """simple docstring"""







   _UpperCamelCase				=  NatModelTester(self)





from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_sentencepiece_available, is_torch_available


lowerCamelCase__   = {
    '''configuration_ernie_m''': ['''ERNIE_M_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''ErnieMConfig'''],
}

try:
   if not is_sentencepiece_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   lowerCamelCase__   = ['''ErnieMTokenizer''']

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   lowerCamelCase__   = [
       '''ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST''',
       '''ErnieMForMultipleChoice''',
       '''ErnieMForQuestionAnswering''',
       '''ErnieMForSequenceClassification''',
       '''ErnieMForTokenClassification''',
       '''ErnieMModel''',
       '''ErnieMPreTrainedModel''',
       '''ErnieMForInformationExtraction''',
   ]


if TYPE_CHECKING:
   from .configuration_ernie_m import ERNIE_M_PRETRAINED_CONFIG_ARCHIVE_MAP, ErnieMConfig

   try:
      if not is_sentencepiece_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .tokenization_ernie_m import ErnieMTokenizer

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_ernie_m import (
          ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST,
          ErnieMForInformationExtraction,
          ErnieMForMultipleChoice,
          ErnieMForQuestionAnswering,
          ErnieMForSequenceClassification,
          ErnieMForTokenClassification,
          ErnieMModel,
          ErnieMPreTrainedModel,
      )


else:
   import sys

   lowerCamelCase__   = _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)




