

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






     def __init__(							self						,		_lowercase						,		_lowercase=13						,		_lowercase=64						,		_lowercase=4						,		_lowercase=3						,		_lowercase=16						,		_lowercase=[1, 2, 1]						,		_lowercase=[2, 4, 8]						,		_lowercase=3						,		_lowercase=2.0						,		_lowercase=True						,		_lowercase=0.0						,		_lowercase=0.0						,		_lowercase=0.1						,		_lowercase="gelu"						,		_lowercase=True						,		_lowercase=0.02						,		_lowercase=1e-5						,		_lowercase=True						,		_lowercase=None						,		_lowercase=True						,		_lowercase=10						,		_lowercase=["stage1", "stage2"]						,		_lowercase=[1, 2]						,		)->			Union[str, Any]:
         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     )->			Any:
         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     )->			Tuple:
         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						,		_lowercase						,		_lowercase						,		_lowercase     )->			Tuple:
         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						,		_lowercase						,		_lowercase						,		_lowercase     )->			Optional[int]:
         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						,		_lowercase						,		_lowercase						,		_lowercase     )->			Dict:
         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     )->			Optional[Any]:
         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 		__magic_name__ (				snake_case       ,	snake_case       ,	unittest.TestCase							):
     UpperCamelCase_     :List[str]									=	(
         (
             NatModel,
             NatForImageClassification,
             NatBackbone,
         )
         if is_torch_available()
         else ()
     )
     UpperCamelCase_     :str									=	(
         {"""feature-extraction""": NatModel, """image-classification""": NatForImageClassification}
         if is_torch_available()
         else {}
     )
     UpperCamelCase_     :Dict									=	False

     UpperCamelCase_     :Union[str, Any]									=	False
     UpperCamelCase_     :int									=	False
     UpperCamelCase_     :Tuple									=	False
     UpperCamelCase_     :Optional[Any]									=	False






     def      UpperCAmelCase_							(							self     )->			List[Any]:
         UpperCamelCase_          =		NatModelTester(self     )
         UpperCamelCase_          =		ConfigTester(self						,		config_class=_lowercase						,		embed_dim=37     )






     def      UpperCAmelCase_							(							self     )->			Optional[int]:
         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     )->			Tuple:
         return






     def      UpperCAmelCase_							(							self     )->			int:
         UpperCamelCase_          =		self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*_lowercase     )






     def      UpperCAmelCase_							(							self     )->			Optional[int]:
         UpperCamelCase_          =		self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*_lowercase     )






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






     @unittest.skip(reason="Nat does not use feedforward chunking"     )
     def      UpperCAmelCase_							(							self     )->			List[Any]:
         pass






     def      UpperCAmelCase_							(							self     )->			int:
         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     )->			Tuple:
         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     )->			Tuple:
         self.skipTest("Nat's attention operation is handled entirely by NATTEN."     )






     def      UpperCAmelCase_							(							self						,		_lowercase						,		_lowercase						,		_lowercase						,		_lowercase     )->			Dict:
         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     )->			List[str]:
         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     )->			Dict:
         for model_name in NAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             UpperCamelCase_          =		NatModel.from_pretrained(_lowercase     )
             self.assertIsNotNone(_lowercase     )






     def      UpperCAmelCase_							(							self     )->			Dict:
         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 		__magic_name__ (				unittest.TestCase							):






     @cached_property
     def      UpperCAmelCase_							(							self     )->			str:
         return AutoImageProcessor.from_pretrained("shi-labs/nat-mini-in1k-224"     ) if is_vision_available() else None






     @slow
     def      UpperCAmelCase_							(							self     )->			List[Any]:
         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, 1_000)     )
         self.assertEqual(outputs.logits.shape						,		_lowercase     )
         UpperCamelCase_          =		torch.tensor([0.3_805, -0.8_676, -0.3_912]     ).to(_lowercase     )
         self.assertTrue(torch.allclose(outputs.logits[0, :3]						,		_lowercase						,		atol=1e-4     )     )






@require_torch
@require_natten
class 		__magic_name__ (				unittest.TestCase       ,	snake_case							):
     UpperCamelCase_     :int									=	(NatBackbone,) if is_torch_available() else ()
     UpperCamelCase_     :Optional[int]									=	NatConfig






     def      UpperCAmelCase_							(							self     )->			Tuple:
         UpperCamelCase_          =		NatModelTester(self     )






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


_lowerCamelCase					:     List[str]					  =  logging.get_logger(__name__)

_lowerCamelCase					:     List[str]					  =  {
    "ZinengTang/tvlt-base": "https://huggingface.co/ZinengTang/tvlt-base/blob/main/config.json",
}
class 				__snake_case     (_a						):
				lowerCAmelCase__			=   "tvlt"



				def __init__(					self					:						Dict						,		_UpperCAmelCase					:						Any=224						,		_UpperCAmelCase					:						Any=2048						,		_UpperCAmelCase					:						Optional[int]=128						,		_UpperCAmelCase					:						Tuple=[16, 16]						,		_UpperCAmelCase					:						Optional[int]=[16, 16]						,		_UpperCAmelCase					:						Tuple=3						,		_UpperCAmelCase					:						str=1						,		_UpperCAmelCase					:						Dict=8						,		_UpperCAmelCase					:						Any=768						,		_UpperCAmelCase					:						Dict=12						,		_UpperCAmelCase					:						Any=12						,		_UpperCAmelCase					:						Optional[Any]=3072						,		_UpperCAmelCase					:						Tuple="gelu"						,		_UpperCAmelCase					:						List[str]=0.0						,		_UpperCAmelCase					:						Any=0.0						,		_UpperCAmelCase					:						Optional[int]=0.02						,		_UpperCAmelCase					:						int=1E-6						,		_UpperCAmelCase					:						str=True						,		_UpperCAmelCase					:						Dict=False						,		_UpperCAmelCase					:						str=16						,		_UpperCAmelCase					:						List[Any]=512						,		_UpperCAmelCase					:						Optional[Any]=8						,		_UpperCAmelCase					:						Optional[Any]=2048						,		_UpperCAmelCase					:						Optional[int]=0.75						,		_UpperCAmelCase					:						List[Any]=0.15						,		_UpperCAmelCase					:						int="frame-level"						,		_UpperCAmelCase					:						Union[str, Any]=True						,		_UpperCAmelCase					:						Optional[int]=True						,		_UpperCAmelCase					:						Any="classification"						,		**_UpperCAmelCase					:						Tuple						,		)     ->			List[str]:





								'''simple docstring'''




								super().__init__(**_UpperCAmelCase					)

								if audio_mask_type not in ("frame-level", "patch_level"):
												raise ValueError(
												    """audio_mask_type must be one of two acceptable strategies - {'frame_level', 'patch-level') """
												    f"got {audio_mask_type}"					)

								_lowerCAmelCase			:			Tuple        =					image_size
								_lowerCAmelCase			:			List[str]        =					spectrogram_length
								_lowerCAmelCase			:			str        =					frequency_length
								_lowerCAmelCase			:			Dict        =					image_patch_size
								_lowerCAmelCase			:			Tuple        =					audio_patch_size
								_lowerCAmelCase			:			str        =					num_image_channels
								_lowerCAmelCase			:			Dict        =					num_audio_channels
								_lowerCAmelCase			:			str        =					num_frames

								_lowerCAmelCase			:			Any        =					hidden_size
								_lowerCAmelCase			:			List[str]        =					num_hidden_layers
								_lowerCAmelCase			:			Union[str, Any]        =					num_attention_heads
								_lowerCAmelCase			:			List[Any]        =					intermediate_size
								_lowerCAmelCase			:			Any        =					hidden_act
								_lowerCAmelCase			:			Any        =					hidden_dropout_prob
								_lowerCAmelCase			:			Union[str, Any]        =					attention_probs_dropout_prob
								_lowerCAmelCase			:			str        =					initializer_range
								_lowerCAmelCase			:			Tuple        =					layer_norm_eps
								_lowerCAmelCase			:			str        =					qkv_bias
								_lowerCAmelCase			:			Tuple        =					use_mean_pooling

								_lowerCAmelCase			:			Union[str, Any]        =					decoder_num_attention_heads
								_lowerCAmelCase			:			int        =					decoder_hidden_size
								_lowerCAmelCase			:			List[str]        =					decoder_num_hidden_layers
								_lowerCAmelCase			:			Optional[int]        =					decoder_intermediate_size
								_lowerCAmelCase			:			Any        =					pixel_mask_ratio
								_lowerCAmelCase			:			int        =					audio_mask_ratio
								_lowerCAmelCase			:			Optional[int]        =					audio_mask_type

								_lowerCAmelCase			:			Dict        =					task_matching
								_lowerCAmelCase			:			Any        =					task_mae
								_lowerCAmelCase			:			int        =					loss_type

