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






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






   def  __UpperCamelCase	( self							:	List[Any]    )   ->		Any:
       lowercase__					:  Any          =							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

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

       lowercase__					:  Tuple          =							self.get_config()

       return config, pixel_values, labels






   def  __UpperCamelCase	( self							:	int    )   ->		str:
       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_							:	str      ,		lowercase_							:	Tuple    )   ->		List[str]:
       lowercase__					:  str          =							NatModel(config=lowercase_    )
       model.to(lowercase_    )
       model.eval()
       lowercase__					:  Union[str, Any]          =							model(lowercase_    )

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

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

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






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



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

   __A      :						Union[str, Any]   =							False
   __A      :						Optional[int]   =							False
   __A      :						Optional[Any]   =							False
   __A      :						List[Any]   =							False






   def  __UpperCamelCase	( self							:	Optional[Any]    )   ->		Optional[int]:
       lowercase__					:  Any          =							NatModelTester(self    )
       lowercase__					:  Union[str, Any]          =							ConfigTester(self      ,		config_class=lowercase_      ,		embed_dim=37    )






   def  __UpperCamelCase	( self							:	List[str]    )   ->		Union[str, Any]:
       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[Any]    )   ->		Tuple:
       return






   def  __UpperCamelCase	( self							:	Union[str, Any]    )   ->		List[str]:
       lowercase__					:  List[Any]          =							self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*lowercase_    )






   def  __UpperCamelCase	( self							:	List[Any]    )   ->		Tuple:
       lowercase__					:  List[Any]          =							self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_classification(*lowercase_    )






   def  __UpperCamelCase	( self							:	int    )   ->		Optional[Any]:
       lowercase__					:  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  __UpperCamelCase	( self							:	str    )   ->		int:
       pass






   @unittest.skip(reason="Nat does not use feedforward chunking"    )
   def  __UpperCamelCase	( self							:	Optional[Any]    )   ->		Union[str, Any]:
       pass






   def  __UpperCamelCase	( self							:	int    )   ->		Tuple:
       lowercase__   ,  lowercase__					:  Any          =							self.model_tester.prepare_config_and_inputs_for_common()

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






   def  __UpperCamelCase	( self							:	Optional[int]    )   ->		int:
       lowercase__   ,  lowercase__					:  Optional[int]          =							self.model_tester.prepare_config_and_inputs_for_common()

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

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






   def  __UpperCamelCase	( self							:	List[Any]    )   ->		str:
       self.skipTest("Nat's attention operation is handled entirely by NATTEN."    )






   def  __UpperCamelCase	( self							:	List[Any]      ,		lowercase_							:	Union[str, Any]      ,		lowercase_							:	Any      ,		lowercase_							:	int      ,		lowercase_							:	Optional[int]    )   ->		Tuple:
       lowercase__					:  List[Any]          =							model_class(lowercase_    )
       model.to(lowercase_    )
       model.eval()

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

       lowercase__					:  Tuple          =							outputs.hidden_states

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

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

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

       lowercase__					:  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:
           lowercase__					:  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"]
           lowercase__					:  Optional[Any]          =							True

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






   @slow
   def  __UpperCamelCase	( self							:	Dict    )   ->		Optional[Any]:
       for model_name in NAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           lowercase__					:  Optional[Any]          =							NatModel.from_pretrained(lowercase_    )
           self.assertIsNotNone(lowercase_    )







   def  __UpperCamelCase	( self							:	Union[str, Any]    )   ->		Tuple:
       lowercase__   ,  lowercase__					:  Any          =							self.model_tester.prepare_config_and_inputs_for_common()

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






   @cached_property
   def  __UpperCamelCase	( self							:	Any    )   ->		Union[str, Any]:
       return AutoImageProcessor.from_pretrained("shi-labs/nat-mini-in1k-224"    ) if is_vision_available() else None







   @slow
   def  __UpperCamelCase	( self							:	Any    )   ->		Union[str, Any]:
       lowercase__					:  Optional[Any]          =							NatForImageClassification.from_pretrained("shi-labs/nat-mini-in1k-224"    ).to(lowercase_    )
       lowercase__					:  Optional[int]          =							self.default_image_processor

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

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

       # verify the logits
       lowercase__					:  Dict          =							torch.Size((1, 10_00)    )
       self.assertEqual(outputs.logits.shape      ,		lowercase_    )
       lowercase__					:  Dict          =							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     snake_case_			(						unittest.TestCase    ,__A			):
   __A      :						int   =							(NatBackbone,) if is_torch_available() else ()
   __A      :						Optional[Any]   =							NatConfig






   def  __UpperCamelCase	( self							:	Union[str, Any]    )   ->		Optional[int]:
       lowercase__					:  Dict          =							NatModelTester(self    )




import unittest
from pathlib import Path

from transformers import is_vision_available, load_tool
from transformers.testing_utils import get_tests_dir

from .test_tools_common import ToolTesterMixin


if is_vision_available():
      from PIL import Image



class     snake_case_			(						unittest.TestCase    ,__A			):






   def  __UpperCamelCase	( self							:	int    )   ->		Union[str, Any]:
       lowercase__					:  Any          =							load_tool("image-segmentation"    )
       self.tool.setup()
       lowercase__					:  Optional[int]          =							load_tool("image-segmentation"      ,		remote=lowercase_    )






   def  __UpperCamelCase	( self							:	Any    )   ->		Tuple:
       lowercase__					:  Any          =							Image.open(Path(get_tests_dir("fixtures/tests_samples/COCO"    )    ) / "000000039769.png"    )
       lowercase__					:  Tuple          =							self.tool(lowercase_      ,		"cat"    )
       self.assertTrue(isinstance(lowercase_      ,		Image.Image    )    )






   def  __UpperCamelCase	( self							:	Optional[int]    )   ->		Dict:
       lowercase__					:  List[Any]          =							Image.open(Path(get_tests_dir("fixtures/tests_samples/COCO"    )    ) / "000000039769.png"    )
       lowercase__					:  Optional[int]          =							self.remote_tool(lowercase_      ,		"cat"    )
       self.assertTrue(isinstance(lowercase_      ,		Image.Image    )    )






   def  __UpperCamelCase	( self							:	List[Any]    )   ->		Dict:
       lowercase__					:  Optional[Any]          =							Image.open(Path(get_tests_dir("fixtures/tests_samples/COCO"    )    ) / "000000039769.png"    )
       lowercase__					:  int          =							self.tool(image=lowercase_      ,		label="cat"    )
       self.assertTrue(isinstance(lowercase_      ,		Image.Image    )    )







   def  __UpperCamelCase	( self							:	Dict    )   ->		List[Any]:
       lowercase__					:  Tuple          =							Image.open(Path(get_tests_dir("fixtures/tests_samples/COCO"    )    ) / "000000039769.png"    )
       lowercase__					:  Union[str, Any]          =							self.remote_tool(image=lowercase_      ,		label="cat"    )
       self.assertTrue(isinstance(lowercase_      ,		Image.Image    )    )



