
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       ,	UpperCAmelCase       ,	UpperCAmelCase=13       ,	UpperCAmelCase=64       ,	UpperCAmelCase=4       ,	UpperCAmelCase=3       ,	UpperCAmelCase=16       ,	UpperCAmelCase=[1, 2, 1]       ,	UpperCAmelCase=[2, 4, 8]       ,	UpperCAmelCase=3       ,	UpperCAmelCase=2.0       ,	UpperCAmelCase=True       ,	UpperCAmelCase=0.0       ,	UpperCAmelCase=0.0       ,	UpperCAmelCase=0.1       ,	UpperCAmelCase="gelu"       ,	UpperCAmelCase=True       ,	UpperCAmelCase=0.0_2       ,	UpperCAmelCase=1e-5       ,	UpperCAmelCase=True       ,	UpperCAmelCase=None       ,	UpperCAmelCase=True       ,	UpperCAmelCase=10       ,	UpperCAmelCase=["stage1", "stage2"]       ,	UpperCAmelCase=[1, 2]       ,	):
							lowerCamelCase_ 		=     parent
							lowerCamelCase_ 		=     batch_size
							lowerCamelCase_ 		=     image_size
							lowerCamelCase_ 		=     patch_size
							lowerCamelCase_ 		=     num_channels
							lowerCamelCase_ 		=     embed_dim
							lowerCamelCase_ 		=     depths
							lowerCamelCase_ 		=     num_heads
							lowerCamelCase_ 		=     kernel_size
							lowerCamelCase_ 		=     mlp_ratio
							lowerCamelCase_ 		=     qkv_bias
							lowerCamelCase_ 		=     hidden_dropout_prob
							lowerCamelCase_ 		=     attention_probs_dropout_prob
							lowerCamelCase_ 		=     drop_path_rate
							lowerCamelCase_ 		=     hidden_act
							lowerCamelCase_ 		=     patch_norm
							lowerCamelCase_ 		=     layer_norm_eps
							lowerCamelCase_ 		=     initializer_range
							lowerCamelCase_ 		=     is_training
							lowerCamelCase_ 		=     scope
							lowerCamelCase_ 		=     use_labels
							lowerCamelCase_ 		=     num_labels
							lowerCamelCase_ 		=     out_features
							lowerCamelCase_ 		=     out_indices




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

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

							lowerCamelCase_ 		=     self.get_config()

							return config, pixel_values, labels




				def 					UpperCAmelCase__	(							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 					UpperCAmelCase__	(							self       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase      ):
							lowerCamelCase_ 		=     NatModel(config=UpperCAmelCase      )
							model.to(UpperCAmelCase      )
							model.eval()
							lowerCamelCase_ 		=     model(UpperCAmelCase      )

							lowerCamelCase_ 		=     lowerCamelCase_ 		=     (config.image_size // config.patch_size) // (2 ** (len(config.depths      ) - 1))
							lowerCamelCase_ 		=     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       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase      ):
							lowerCamelCase_ 		=     NatForImageClassification(UpperCAmelCase      )
							model.to(UpperCAmelCase      )
							model.eval()
							lowerCamelCase_ 		=     model(UpperCAmelCase       ,	labels=UpperCAmelCase      )
							self.parent.assertEqual(result.logits.shape       ,	(self.batch_size, self.num_labels)      )

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

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




				def 					UpperCAmelCase__	(							self       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase      ):
							lowerCamelCase_ 		=     NatBackbone(config=UpperCAmelCase      )
							model.to(UpperCAmelCase      )
							model.eval()
							lowerCamelCase_ 		=     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], 16, 16]      )

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

							# verify backbone works with out_features=None
							lowerCamelCase_ 		=     None
							lowerCamelCase_ 		=     NatBackbone(config=UpperCAmelCase      )
							model.to(UpperCAmelCase      )
							model.eval()
							lowerCamelCase_ 		=     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 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     self.prepare_config_and_inputs()
							lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_ 		=     config_and_inputs
							lowerCamelCase_ 		=     {'''pixel_values''': pixel_values}
							return config, inputs_dict




@require_natten
@require_torch
class    __lowerCamelCase     (   lowerCAmelCase	,  lowerCAmelCase	,  unittest.TestCase ):
				a__:			Union[str, Any]									=						(
				    (
				        NatModel,
				        NatForImageClassification,
				        NatBackbone,
				    )
				    if is_torch_available()
				    else ()
				)
				a__:			List[str]									=						(
				    {'feature-extraction': NatModel, 'image-classification': NatForImageClassification}
				    if is_torch_available()
				    else {}
				)
				a__:			Tuple									=						False

				a__:			Optional[int]									=						False
				a__:			int									=						False
				a__:			str									=						False
				a__:			str									=						False




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     NatModelTester(self      )
							lowerCamelCase_ 		=     ConfigTester(self       ,	config_class=UpperCAmelCase       ,	embed_dim=37      )




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




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*UpperCAmelCase      )




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_image_classification(*UpperCAmelCase      )




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     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 					UpperCAmelCase__	(							self      ):
							pass




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




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.model_tester.prepare_config_and_inputs_for_common()

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




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.model_tester.prepare_config_and_inputs_for_common()

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

										lowerCamelCase_ 		=     ['''pixel_values''']
										self.assertListEqual(arg_names[:1]       ,	UpperCAmelCase      )




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




				def 					UpperCAmelCase__	(							self       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase      ):
							lowerCamelCase_ 		=     model_class(UpperCAmelCase      )
							model.to(UpperCAmelCase      )
							model.eval()

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

							lowerCamelCase_ 		=     outputs.hidden_states

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

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

										lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_ 		=     reshaped_hidden_states[0].shape
										lowerCamelCase_ 		=     (
										    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 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.model_tester.prepare_config_and_inputs_for_common()

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

										self.check_hidden_states_output(UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase      )




				@slow
				def 					UpperCAmelCase__	(							self      ):
							for model_name in NAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										lowerCamelCase_ 		=     NatModel.from_pretrained(UpperCAmelCase      )
										self.assertIsNotNone(UpperCAmelCase      )





				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.model_tester.prepare_config_and_inputs_for_common()

							lowerCamelCase_ 		=     _config_zero_init(UpperCAmelCase      )
							for model_class in self.all_model_classes:
										lowerCamelCase_ 		=     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    __lowerCamelCase     (   unittest.TestCase ):




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





				@slow
				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     NatForImageClassification.from_pretrained('''shi-labs/nat-mini-in1k-224'''      ).to(UpperCAmelCase      )
							lowerCamelCase_ 		=     self.default_image_processor

							lowerCamelCase_ 		=     Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''      )
							lowerCamelCase_ 		=     image_processor(images=UpperCAmelCase       ,	return_tensors='''pt'''      ).to(UpperCAmelCase      )

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

							# verify the logits
							lowerCamelCase_ 		=     torch.Size((1, 1000)      )
							self.assertEqual(outputs.logits.shape       ,	UpperCAmelCase      )
							lowerCamelCase_ 		=     torch.tensor([0.3_8_0_5, -0.8_6_7_6, -0.3_9_1_2]      ).to(UpperCAmelCase      )
							self.assertTrue(torch.allclose(outputs.logits[0, :3]       ,	UpperCAmelCase       ,	atol=1e-4      )      )




@require_torch
@require_natten
class    __lowerCamelCase     (   unittest.TestCase	,  lowerCAmelCase ):
				a__:			List[str]									=						(NatBackbone,) if is_torch_available() else ()
				a__:			Optional[Any]									=						NatConfig




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     NatModelTester(self      )





"""simple docstring"""


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


A_						       =				logging.get_logger(__name__)

A_						       =				{
    """uw-madison/yoso-4096""": """https://huggingface.co/uw-madison/yoso-4096/resolve/main/config.json""",
    # See all YOSO models at https://huggingface.co/models?filter=yoso
}




class    __lowerCamelCase     (   lowerCAmelCase ):
				a__:			Dict									=						'yoso'




				def __init__(							self       ,	UpperCAmelCase=5_0265       ,	UpperCAmelCase=768       ,	UpperCAmelCase=12       ,	UpperCAmelCase=12       ,	UpperCAmelCase=3072       ,	UpperCAmelCase="gelu"       ,	UpperCAmelCase=0.1       ,	UpperCAmelCase=0.1       ,	UpperCAmelCase=4096       ,	UpperCAmelCase=1       ,	UpperCAmelCase=0.0_2       ,	UpperCAmelCase=1e-1_2       ,	UpperCAmelCase="absolute"       ,	UpperCAmelCase=True       ,	UpperCAmelCase=9       ,	UpperCAmelCase=64       ,	UpperCAmelCase=None       ,	UpperCAmelCase=True       ,	UpperCAmelCase=True       ,	UpperCAmelCase=1       ,	UpperCAmelCase=0       ,	UpperCAmelCase=2       ,	**UpperCAmelCase       ,	):
							super().__init__(pad_token_id=UpperCAmelCase       ,	bos_token_id=UpperCAmelCase       ,	eos_token_id=UpperCAmelCase       ,	**UpperCAmelCase      )

							lowerCamelCase_ 		=     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_ 		=     type_vocab_size
							lowerCamelCase_ 		=     layer_norm_eps
							lowerCamelCase_ 		=     position_embedding_type
							lowerCamelCase_ 		=     use_expectation
							lowerCamelCase_ 		=     hash_code_len
							lowerCamelCase_ 		=     num_hash
							lowerCamelCase_ 		=     conv_window
							lowerCamelCase_ 		=     use_fast_hash
							lowerCamelCase_ 		=     lsh_backward



