






import collections
import inspect
import unittest

from transformers import DinatConfig
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 DinatBackbone, DinatForImageClassification, DinatModel
							from transformers.models.dinat.modeling_dinat import DINAT_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
							from PIL import Image

							from transformers import AutoImageProcessor




class 						UpperCAmelCase_  :


						'''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=[[3], [1, 2], [1]]						, _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]						, )       -> Tuple:
								snake_case_ :							Optional[int]         =				parent
								snake_case_ :							Optional[Any]         =				batch_size
								snake_case_ :							Any         =				image_size
								snake_case_ :							Tuple         =				patch_size
								snake_case_ :							Any         =				num_channels
								snake_case_ :							List[Any]         =				embed_dim
								snake_case_ :							Optional[int]         =				depths
								snake_case_ :							Any         =				num_heads
								snake_case_ :							Optional[Any]         =				kernel_size
								snake_case_ :							Tuple         =				dilations
								snake_case_ :							Dict         =				mlp_ratio
								snake_case_ :							Optional[int]         =				qkv_bias
								snake_case_ :							Dict         =				hidden_dropout_prob
								snake_case_ :							Tuple         =				attention_probs_dropout_prob
								snake_case_ :							List[str]         =				drop_path_rate
								snake_case_ :							int         =				hidden_act
								snake_case_ :							int         =				patch_norm
								snake_case_ :							Tuple         =				layer_norm_eps
								snake_case_ :							List[Any]         =				initializer_range
								snake_case_ :							Tuple         =				is_training
								snake_case_ :							Optional[int]         =				scope
								snake_case_ :							List[Any]         =				use_labels
								snake_case_ :							Tuple         =				num_labels
								snake_case_ :							Any         =				out_features
								snake_case_ :							Any         =				out_indices






						def 					_lowerCAmelCase				(		self )       -> int:
								snake_case_ :							List[str]         =				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

								snake_case_ :							Dict         =				None
								if self.use_labels:
										snake_case_ :							Union[str, Any]         =				ids_tensor([self.batch_size]						, self.num_labels )

								snake_case_ :							Optional[int]         =				self.get_config()

								return config, pixel_values, labels






						def 					_lowerCAmelCase				(		self )       -> Any:
								return DinatConfig(
								    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						, dilations=self.dilations						, 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 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )       -> Any:
								snake_case_ :							str         =				DinatModel(config=_SCREAMING_SNAKE_CASE )
								model.to(_SCREAMING_SNAKE_CASE )
								model.eval()
								snake_case_ :							List[Any]         =				model(_SCREAMING_SNAKE_CASE )

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

								# test greyscale images
								snake_case_ :							Optional[int]         =				1
								snake_case_ :							str         =				DinatForImageClassification(_SCREAMING_SNAKE_CASE )
								model.to(_SCREAMING_SNAKE_CASE )
								model.eval()

								snake_case_ :							Any         =				floats_tensor([self.batch_size, 1, self.image_size, self.image_size] )
								snake_case_ :							List[str]         =				model(_SCREAMING_SNAKE_CASE )
								self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.num_labels) )






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )       -> Any:
								snake_case_ :							Optional[int]         =				DinatBackbone(config=_SCREAMING_SNAKE_CASE )
								model.to(_SCREAMING_SNAKE_CASE )
								model.eval()
								snake_case_ :							List[Any]         =				model(_SCREAMING_SNAKE_CASE )

								# 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
								snake_case_ :							Optional[Any]         =				None
								snake_case_ :							Tuple         =				DinatBackbone(config=_SCREAMING_SNAKE_CASE )
								model.to(_SCREAMING_SNAKE_CASE )
								model.eval()
								snake_case_ :							int         =				model(_SCREAMING_SNAKE_CASE )

								# 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 					_lowerCAmelCase				(		self )       -> Optional[Any]:
								snake_case_ :							Optional[Any]         =				self.prepare_config_and_inputs()
								snake_case_       ,				snake_case_       ,				snake_case_ :							Tuple         =				config_and_inputs
								snake_case_ :							List[str]         =				{"pixel_values": pixel_values}
								return config, inputs_dict




@require_natten
@require_torch
class 						UpperCAmelCase_  (		SCREAMING_SNAKE_CASE__    ,     SCREAMING_SNAKE_CASE__    ,     unittest.TestCase    ):


						'''simple docstring'''




						A   :      List[Any]						 =    (
						    (
						        DinatModel,
						        DinatForImageClassification,
						        DinatBackbone,
						    )
						    if is_torch_available()
						    else ()
						)
						A   :      List[Any]						 =    (
						    {'feature-extraction': DinatModel, 'image-classification': DinatForImageClassification}
						    if is_torch_available()
						    else {}
						)
						A   :      List[Any]						 =    False

						A   :      int						 =    False
						A   :      Any						 =    False
						A   :      List[Any]						 =    False
						A   :      Dict						 =    False






						def 					_lowerCAmelCase				(		self )       -> str:
								snake_case_ :							Optional[int]         =				DinatModelTester(self )
								snake_case_ :							Optional[int]         =				ConfigTester(self						, config_class=_SCREAMING_SNAKE_CASE						, embed_dim=37 )






						def 					_lowerCAmelCase				(		self )       -> Dict:
								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 					_lowerCAmelCase				(		self )       -> Tuple:
								return






						def 					_lowerCAmelCase				(		self )       -> List[Any]:
								snake_case_ :							Union[str, Any]         =				self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*_SCREAMING_SNAKE_CASE )






						def 					_lowerCAmelCase				(		self )       -> Optional[Any]:
								snake_case_ :							Optional[Any]         =				self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_image_classification(*_SCREAMING_SNAKE_CASE )






						def 					_lowerCAmelCase				(		self )       -> List[Any]:
								snake_case_ :							int         =				self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_backbone(*_SCREAMING_SNAKE_CASE )






						@unittest.skip(reason="Dinat does not use inputs_embeds" )
						def 					_lowerCAmelCase				(		self )       -> Any:
								pass






						@unittest.skip(reason="Dinat does not use feedforward chunking" )
						def 					_lowerCAmelCase				(		self )       -> List[str]:
								pass






						def 					_lowerCAmelCase				(		self )       -> Any:
								snake_case_       ,				snake_case_ :							Union[str, Any]         =				self.model_tester.prepare_config_and_inputs_for_common()

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






						def 					_lowerCAmelCase				(		self )       -> Optional[int]:
								snake_case_       ,				snake_case_ :							str         =				self.model_tester.prepare_config_and_inputs_for_common()

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

										snake_case_ :							List[Any]         =				["pixel_values"]
										self.assertListEqual(arg_names[:1]						, _SCREAMING_SNAKE_CASE )






						def 					_lowerCAmelCase				(		self )       -> Any:
								self.skipTest("Dinat's attention operation is handled entirely by NATTEN." )






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )       -> Dict:
								snake_case_ :							Union[str, Any]         =				model_class(_SCREAMING_SNAKE_CASE )
								model.to(_SCREAMING_SNAKE_CASE )
								model.eval()

								with torch.no_grad():
										snake_case_ :							Any         =				model(**self._prepare_for_class(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE ) )

								snake_case_ :							Optional[Any]         =				outputs.hidden_states

								snake_case_ :							List[str]         =				getattr(
								    self.model_tester						, "expected_num_hidden_layers"						, len(self.model_tester.depths ) + 1 )
								self.assertEqual(len(_SCREAMING_SNAKE_CASE )						, _SCREAMING_SNAKE_CASE )

								# Dinat has a different seq_length
								snake_case_ :							str         =				(
								    config.patch_size
								    if isinstance(config.patch_size						, collections.abc.Iterable )
								    else (config.patch_size, config.patch_size)
								)

								snake_case_ :							int         =				image_size[0] // patch_size[0]
								snake_case_ :							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__ != "DinatBackbone":
										snake_case_ :							Union[str, Any]         =				outputs.reshaped_hidden_states
										self.assertEqual(len(_SCREAMING_SNAKE_CASE )						, _SCREAMING_SNAKE_CASE )

										snake_case_       ,				snake_case_       ,				snake_case_       ,				snake_case_ :							Union[str, Any]         =				reshaped_hidden_states[0].shape
										snake_case_ :							List[Any]         =				(
										    reshaped_hidden_states[0].view(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE ).permute(0						, 2						, 3						, 1 )
										)
										self.assertListEqual(
										    list(reshaped_hidden_states.shape[-3:] )						, [height, width, self.model_tester.embed_dim]						, )






						def 					_lowerCAmelCase				(		self )       -> Optional[Any]:
								snake_case_       ,				snake_case_ :							Tuple         =				self.model_tester.prepare_config_and_inputs_for_common()

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

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

										self.check_hidden_states_output(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )






						@slow
						def 					_lowerCAmelCase				(		self )       -> Union[str, Any]:
								for model_name in DINAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										snake_case_ :							Optional[int]         =				DinatModel.from_pretrained(_SCREAMING_SNAKE_CASE )
										self.assertIsNotNone(_SCREAMING_SNAKE_CASE )






						def 					_lowerCAmelCase				(		self )       -> List[Any]:
								snake_case_       ,				snake_case_ :							Any         =				self.model_tester.prepare_config_and_inputs_for_common()

								snake_case_ :							List[str]         =				_config_zero_init(_SCREAMING_SNAKE_CASE )
								for model_class in self.all_model_classes:
										snake_case_ :							Optional[int]         =				model_class(config=_SCREAMING_SNAKE_CASE )
										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 					_lowerCAmelCase				(		self )       -> str:
								return AutoImageProcessor.from_pretrained("shi-labs/dinat-mini-in1k-224" ) if is_vision_available() else None






						@slow
						def 					_lowerCAmelCase				(		self )       -> Dict:
								snake_case_ :							List[str]         =				DinatForImageClassification.from_pretrained("shi-labs/dinat-mini-in1k-224" ).to(_SCREAMING_SNAKE_CASE )
								snake_case_ :							Optional[int]         =				self.default_image_processor

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

								# forward pass
								with torch.no_grad():
										snake_case_ :							Optional[int]         =				model(**_SCREAMING_SNAKE_CASE )

								# verify the logits
								snake_case_ :							Tuple         =				torch.Size((1, 1000) )
								self.assertEqual(outputs.logits.shape						, _SCREAMING_SNAKE_CASE )
								snake_case_ :							Any         =				torch.tensor([-0.1545, -0.7667, 0.4642] ).to(_SCREAMING_SNAKE_CASE )
								self.assertTrue(torch.allclose(outputs.logits[0, :3]						, _SCREAMING_SNAKE_CASE						, atol=1e-4 ) )




@require_torch
@require_natten
class 						UpperCAmelCase_  (		unittest.TestCase    ,     SCREAMING_SNAKE_CASE__    ):


						'''simple docstring'''




						A   :      str						 =    (DinatBackbone,) if is_torch_available() else ()
						A   :      Optional[int]						 =    DinatConfig






						def 					_lowerCAmelCase				(		self )       -> Tuple:
								snake_case_ :							int         =				DinatModelTester(self )








def    lowerCAmelCase__		(			_a				:	str     ,					_a				:	int = 13			):
		snake_case_ :							List[Any]         =				""
		for c in s:
				if "A" <= c <= "Z":
						out += chr(ord("A"			) + (ord(_a			) - ord("A"			) + n) % 26			)
				elif "a" <= c <= "z":
						out += chr(ord("a"			) + (ord(_a			) - ord("a"			) + n) % 26			)
				else:
						out += c
		return out

def    lowerCAmelCase__		(			):
		snake_case_ :							int         =				input("Enter message: "			)

		snake_case_ :							List[Any]         =				dencrypt(_a     ,					13			)
		print("Encryption:"     ,					_a			)

		snake_case_ :							int         =				dencrypt(_a     ,					13			)
		print("Decryption: "     ,					_a			)


if __name__ == "__main__":
							import doctest

							doctest.testmod()
							main()
