

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     :			Dict					,					UpperCAmelCase__     :			Tuple					,					UpperCAmelCase__     :			Any=1_3					,					UpperCAmelCase__     :			Dict=6_4					,					UpperCAmelCase__     :			Dict=4					,					UpperCAmelCase__     :			Dict=3					,					UpperCAmelCase__     :			int=1_6					,					UpperCAmelCase__     :			List[str]=[1, 2, 1]					,					UpperCAmelCase__     :			List[Any]=[2, 4, 8]					,					UpperCAmelCase__     :			List[Any]=3					,					UpperCAmelCase__     :			Optional[int]=2.0					,					UpperCAmelCase__     :			Tuple=True					,					UpperCAmelCase__     :			Any=0.0					,					UpperCAmelCase__     :			Union[str, Any]=0.0					,					UpperCAmelCase__     :			int=0.1					,					UpperCAmelCase__     :			Optional[Any]="gelu"					,					UpperCAmelCase__     :			List[str]=True					,					UpperCAmelCase__     :			List[Any]=0.02					,					UpperCAmelCase__     :			int=1e-5					,					UpperCAmelCase__     :			str=True					,					UpperCAmelCase__     :			List[str]=None					,					UpperCAmelCase__     :			Any=True					,					UpperCAmelCase__     :			Dict=1_0					,					UpperCAmelCase__     :			List[str]=["stage1", "stage2"]					,					UpperCAmelCase__     :			List[str]=[1, 2]					,					)       ->Dict:
        """simple docstring"""







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


  def 			_lowercase   (      self     :			Optional[int]		)       ->Union[str, Any]:
        """simple docstring"""







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

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

        SCREAMING_SNAKE_CASE			:   Dict	  = self.get_config()

        return config, pixel_values, labels


  def 			_lowercase   (      self     :			Optional[Any]		)       ->Dict:
        """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 			_lowercase   (      self     :			List[Any]					,					UpperCAmelCase__     :			Union[str, Any]					,					UpperCAmelCase__     :			Optional[Any]					,					UpperCAmelCase__     :			List[str]		)       ->Tuple:
        """simple docstring"""







        SCREAMING_SNAKE_CASE			:   Optional[Any]	  = NatModel(config=UpperCAmelCase__		)
        model.to(UpperCAmelCase__		)
        model.eval()
        SCREAMING_SNAKE_CASE			:   int	  = model(UpperCAmelCase__		)

        SCREAMING_SNAKE_CASE			:   Any	  = (config.image_size // config.patch_size) // (2 ** (len(config.depths		) - 1))
        SCREAMING_SNAKE_CASE			:   Union[str, Any]	  = 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 			_lowercase   (      self     :			Tuple					,					UpperCAmelCase__     :			Tuple					,					UpperCAmelCase__     :			str					,					UpperCAmelCase__     :			int		)       ->Any:
        """simple docstring"""







        SCREAMING_SNAKE_CASE			:   List[str]	  = NatForImageClassification(UpperCAmelCase__		)
        model.to(UpperCAmelCase__		)
        model.eval()
        SCREAMING_SNAKE_CASE			:   Union[str, Any]	  = model(UpperCAmelCase__					,					labels=UpperCAmelCase__		)
        self.parent.assertEqual(result.logits.shape					,					(self.batch_size, self.num_labels)		)

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

        SCREAMING_SNAKE_CASE			:   List[Any]	  = floats_tensor([self.batch_size, 1, self.image_size, self.image_size]		)
        SCREAMING_SNAKE_CASE			:   Optional[Any]	  = model(UpperCAmelCase__		)
        self.parent.assertEqual(result.logits.shape					,					(self.batch_size, self.num_labels)		)


  def 			_lowercase   (      self     :			List[str]					,					UpperCAmelCase__     :			Optional[Any]					,					UpperCAmelCase__     :			str					,					UpperCAmelCase__     :			Any		)       ->Union[str, Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE			:   List[Any]	  = NatBackbone(config=UpperCAmelCase__		)
        model.to(UpperCAmelCase__		)
        model.eval()
        SCREAMING_SNAKE_CASE			:   Optional[int]	  = model(UpperCAmelCase__		)

        # 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], 1_6, 1_6]		)

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

        # verify backbone works with out_features=None
        SCREAMING_SNAKE_CASE			:   List[str]	  = None
        SCREAMING_SNAKE_CASE			:   int	  = NatBackbone(config=UpperCAmelCase__		)
        model.to(UpperCAmelCase__		)
        model.eval()
        SCREAMING_SNAKE_CASE			:   Optional[Any]	  = model(UpperCAmelCase__		)

        # 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 			_lowercase   (      self     :			Tuple		)       ->Dict:
        """simple docstring"""







        SCREAMING_SNAKE_CASE			:   Union[str, Any]	  = self.prepare_config_and_inputs()
        SCREAMING_SNAKE_CASE  ,						SCREAMING_SNAKE_CASE  ,						SCREAMING_SNAKE_CASE			:   Any	  = config_and_inputs
        SCREAMING_SNAKE_CASE			:   int	  = {"""pixel_values""": pixel_values}
        return config, inputs_dict







@require_natten
@require_torch
class 							a__							( UpperCAmelCase						,      UpperCAmelCase						,      unittest.TestCase							):







  """simple docstring"""




  UpperCAmelCase__   :			Any      							=(
      (
          NatModel,
          NatForImageClassification,
          NatBackbone,
      )
      if is_torch_available()
      else ()
  )
  UpperCAmelCase__   :			Dict      							=(
      {"""feature-extraction""": NatModel, """image-classification""": NatForImageClassification}
      if is_torch_available()
      else {}
  )
  UpperCAmelCase__   :			List[Any]      							=False

  UpperCAmelCase__   :			Optional[int]      							=False
  UpperCAmelCase__   :			Union[str, Any]      							=False
  UpperCAmelCase__   :			Union[str, Any]      							=False
  UpperCAmelCase__   :			Any      							=False


  def 			_lowercase   (      self     :			Optional[int]		)       ->Any:
        """simple docstring"""







        SCREAMING_SNAKE_CASE			:   Union[str, Any]	  = NatModelTester(self		)
        SCREAMING_SNAKE_CASE			:   Dict	  = ConfigTester(self					,					config_class=UpperCAmelCase__					,					embed_dim=3_7		)


  def 			_lowercase   (      self     :			Optional[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 			_lowercase   (      self     :			List[str]		)       ->Optional[Any]:
        """simple docstring"""







        return


  def 			_lowercase   (      self     :			int		)       ->Dict:
        """simple docstring"""







        SCREAMING_SNAKE_CASE			:   List[str]	  = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*UpperCAmelCase__		)


  def 			_lowercase   (      self     :			Optional[int]		)       ->Dict:
        """simple docstring"""







        SCREAMING_SNAKE_CASE			:   Union[str, Any]	  = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*UpperCAmelCase__		)


  def 			_lowercase   (      self     :			List[str]		)       ->int:
        """simple docstring"""







        SCREAMING_SNAKE_CASE			:   Optional[int]	  = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_backbone(*UpperCAmelCase__		)


  @unittest.skip(reason="""Nat does not use inputs_embeds"""		)
  def 			_lowercase   (      self     :			List[Any]		)       ->int:
        """simple docstring"""







        pass


  @unittest.skip(reason="""Nat does not use feedforward chunking"""		)
  def 			_lowercase   (      self     :			Optional[int]		)       ->str:
        """simple docstring"""







        pass


  def 			_lowercase   (      self     :			List[Any]		)       ->Optional[int]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE  ,						SCREAMING_SNAKE_CASE			:   List[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(UpperCAmelCase__		)
              self.assertIsInstance(model.get_input_embeddings()					,					(nn.Module)		)
              SCREAMING_SNAKE_CASE			:   Union[str, Any]	  = model.get_output_embeddings()
              self.assertTrue(x is None or isinstance(UpperCAmelCase__					,					nn.Linear		)		)


  def 			_lowercase   (      self     :			Union[str, Any]		)       ->str:
        """simple docstring"""







        SCREAMING_SNAKE_CASE  ,						SCREAMING_SNAKE_CASE			:   Union[str, Any]	  = self.model_tester.prepare_config_and_inputs_for_common()

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

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


  def 			_lowercase   (      self     :			int		)       ->List[str]:
        """simple docstring"""







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


  def 			_lowercase   (      self     :			Optional[Any]					,					UpperCAmelCase__     :			str					,					UpperCAmelCase__     :			Optional[Any]					,					UpperCAmelCase__     :			int					,					UpperCAmelCase__     :			str		)       ->int:
        """simple docstring"""







        SCREAMING_SNAKE_CASE			:   str	  = model_class(UpperCAmelCase__		)
        model.to(UpperCAmelCase__		)
        model.eval()

        with torch.no_grad():
              SCREAMING_SNAKE_CASE			:   Any	  = model(**self._prepare_for_class(UpperCAmelCase__					,					UpperCAmelCase__		)		)

        SCREAMING_SNAKE_CASE			:   Any	  = outputs.hidden_states

        SCREAMING_SNAKE_CASE			:   str	  = getattr(
            self.model_tester					,					"""expected_num_hidden_layers"""					,					len(self.model_tester.depths		) + 1		)
        self.assertEqual(len(UpperCAmelCase__		)					,					UpperCAmelCase__		)

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

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

              SCREAMING_SNAKE_CASE  ,						SCREAMING_SNAKE_CASE  ,						SCREAMING_SNAKE_CASE  ,						SCREAMING_SNAKE_CASE			:   Optional[int]	  = reshaped_hidden_states[0].shape
              SCREAMING_SNAKE_CASE			:   str	  = (
                  reshaped_hidden_states[0].view(UpperCAmelCase__					,					UpperCAmelCase__					,					UpperCAmelCase__					,					UpperCAmelCase__		).permute(0					,					2					,					3					,					1		)
              )
              self.assertListEqual(
                  list(reshaped_hidden_states.shape[-3:]		)					,					[height, width, self.model_tester.embed_dim]					,					)


  def 			_lowercase   (      self     :			Union[str, Any]		)       ->Optional[int]:
        """simple docstring"""







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

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

              self.check_hidden_states_output(UpperCAmelCase__					,					UpperCAmelCase__					,					UpperCAmelCase__					,					UpperCAmelCase__		)


  @slow
  def 			_lowercase   (      self     :			List[Any]		)       ->Union[str, Any]:
        """simple docstring"""







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




  def 			_lowercase   (      self     :			List[str]		)       ->Union[str, Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE  ,						SCREAMING_SNAKE_CASE			:   int	  = self.model_tester.prepare_config_and_inputs_for_common()

        SCREAMING_SNAKE_CASE			:   Tuple	  = _config_zero_init(UpperCAmelCase__		)
        for model_class in self.all_model_classes:
              SCREAMING_SNAKE_CASE			:   List[Any]	  = model_class(config=UpperCAmelCase__		)
              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 			_lowercase   (      self     :			List[str]		)       ->str:
        """simple docstring"""







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




  @slow
  def 			_lowercase   (      self     :			Optional[Any]		)       ->Any:
        """simple docstring"""







        SCREAMING_SNAKE_CASE			:   Optional[int]	  = NatForImageClassification.from_pretrained("""shi-labs/nat-mini-in1k-224"""		).to(UpperCAmelCase__		)
        SCREAMING_SNAKE_CASE			:   Union[str, Any]	  = self.default_image_processor

        SCREAMING_SNAKE_CASE			:   List[str]	  = Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""		)
        SCREAMING_SNAKE_CASE			:   List[str]	  = image_processor(images=UpperCAmelCase__					,					return_tensors="""pt"""		).to(UpperCAmelCase__		)

        # forward pass
        with torch.no_grad():
              SCREAMING_SNAKE_CASE			:   Optional[Any]	  = model(**UpperCAmelCase__		)

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







@require_torch
@require_natten
class 							a__							( unittest.TestCase						,      UpperCAmelCase							):







  """simple docstring"""




  UpperCAmelCase__   :			Tuple      							=(NatBackbone,) if is_torch_available() else ()
  UpperCAmelCase__   :			Optional[int]      							=NatConfig


  def 			_lowercase   (      self     :			Union[str, Any]		)       ->Dict:
        """simple docstring"""







        SCREAMING_SNAKE_CASE			:   List[str]	  = NatModelTester(self		)







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


UpperCAmelCase__			:  Dict					     =     {"""configuration_hubert""": ["""HUBERT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """HubertConfig"""]}

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase__			:  Union[str, Any]					     =     [
      """HUBERT_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """HubertForCTC""",
      """HubertForSequenceClassification""",
      """HubertModel""",
      """HubertPreTrainedModel""",
  ]


try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase__			:  List[Any]					     =     [
      """TF_HUBERT_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """TFHubertForCTC""",
      """TFHubertModel""",
      """TFHubertPreTrainedModel""",
  ]

if TYPE_CHECKING:
  from .configuration_hubert import HUBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, HubertConfig

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_hubert import (
        HUBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        HubertForCTC,
        HubertForSequenceClassification,
        HubertModel,
        HubertPreTrainedModel,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_hubert import (
        TF_HUBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFHubertForCTC,
        TFHubertModel,
        TFHubertPreTrainedModel,
    )


else:
  import sys

  UpperCAmelCase__			:  str					     =     _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)




