
"""simple docstring"""


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



						'''simple docstring'''


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


						def   _SCREAMING_SNAKE_CASE (	self       : Tuple )       ->				Tuple:
							A  				=				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

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

							A  				=				self.get_config()

							return config, pixel_values, labels


						def   _SCREAMING_SNAKE_CASE (	self       : List[Any] )       ->				Union[str, Any]:
							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   _SCREAMING_SNAKE_CASE (	self       : str					,A_       : Tuple					,A_       : List[str]					,A_       : Any )       ->				Optional[int]:
							A  				=				NatModel(config=A_ )
							model.to(A_ )
							model.eval()
							A  				=				model(A_ )

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

							# test greyscale images
							A  				=				1
							A  				=				NatForImageClassification(A_ )
							model.to(A_ )
							model.eval()

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


						def   _SCREAMING_SNAKE_CASE (	self       : Optional[int]					,A_       : str					,A_       : Dict					,A_       : Optional[Any] )       ->				Optional[int]:
							A  				=				NatBackbone(config=A_ )
							model.to(A_ )
							model.eval()
							A  				=				model(A_ )

							# 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
							A  				=				None
							A  				=				NatBackbone(config=A_ )
							model.to(A_ )
							model.eval()
							A  				=				model(A_ )

							# 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   _SCREAMING_SNAKE_CASE (	self       : int )       ->				List[str]:
							A  				=				self.prepare_config_and_inputs()
							A				, A				, A  				=				config_and_inputs
							A  				=				{'pixel_values': pixel_values}
							return config, inputs_dict



@require_natten
@require_torch
class 			lowerCAmelCase_			( _lowercase						,   _lowercase						,   unittest.TestCase      ):



						'''simple docstring'''


						_lowerCamelCase:      Optional[Any] =      (
						    (
						        NatModel,
						        NatForImageClassification,
						        NatBackbone,
						    )
						    if is_torch_available()
						    else ()
						)
						_lowerCamelCase:      Dict =      (
						    {'''feature-extraction''': NatModel, '''image-classification''': NatForImageClassification}
						    if is_torch_available()
						    else {}
						)
						_lowerCamelCase:      Optional[int] =      False

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


						def   _SCREAMING_SNAKE_CASE (	self       : List[Any] )       ->				Any:
							A  				=				NatModelTester(self )
							A  				=				ConfigTester(self					,config_class=A_					,embed_dim=37 )


						def   _SCREAMING_SNAKE_CASE (	self       : Any )       ->				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   _SCREAMING_SNAKE_CASE (	self       : Any )       ->				str:
							return


						def   _SCREAMING_SNAKE_CASE (	self       : Optional[Any] )       ->				Union[str, Any]:
							A  				=				self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*A_ )


						def   _SCREAMING_SNAKE_CASE (	self       : List[str] )       ->				Dict:
							A  				=				self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_image_classification(*A_ )


						def   _SCREAMING_SNAKE_CASE (	self       : int )       ->				Optional[Any]:
							A  				=				self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_backbone(*A_ )


						@unittest.skip(reason='Nat does not use inputs_embeds' )
						def   _SCREAMING_SNAKE_CASE (	self       : List[Any] )       ->				int:
							pass


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


						def   _SCREAMING_SNAKE_CASE (	self       : List[str] )       ->				Union[str, Any]:
							A				, A  				=				self.model_tester.prepare_config_and_inputs_for_common()

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


						def   _SCREAMING_SNAKE_CASE (	self       : List[str] )       ->				Tuple:
							A				, A  				=				self.model_tester.prepare_config_and_inputs_for_common()

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

								A  				=				['pixel_values']
								self.assertListEqual(arg_names[:1]					,A_ )


						def   _SCREAMING_SNAKE_CASE (	self       : List[Any] )       ->				int:
							self.skipTest('Nat\'s attention operation is handled entirely by NATTEN.' )


						def   _SCREAMING_SNAKE_CASE (	self       : List[str]					,A_       : Any					,A_       : str					,A_       : Dict					,A_       : Any )       ->				str:
							A  				=				model_class(A_ )
							model.to(A_ )
							model.eval()

							with torch.no_grad():
								A  				=				model(**self._prepare_for_class(A_					,A_ ) )

							A  				=				outputs.hidden_states

							A  				=				getattr(
							    self.model_tester					,'expected_num_hidden_layers'					,len(self.model_tester.depths ) + 1 )
							self.assertEqual(len(A_ )					,A_ )

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

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

								A				, A				, A				, A  				=				reshaped_hidden_states[0].shape
								A  				=				(
								    reshaped_hidden_states[0].view(A_					,A_					,A_					,A_ ).permute(0					,2					,3					,1 )
								)
								self.assertListEqual(
								    list(reshaped_hidden_states.shape[-3:] )					,[height, width, self.model_tester.embed_dim]					,)


						def   _SCREAMING_SNAKE_CASE (	self       : Tuple )       ->				List[str]:
							A				, A  				=				self.model_tester.prepare_config_and_inputs_for_common()

							A  				=				(
							    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:
								A  				=				True
								self.check_hidden_states_output(A_					,A_					,A_					,A_ )

								# check that output_hidden_states also work using config
								del inputs_dict["output_hidden_states"]
								A  				=				True

								self.check_hidden_states_output(A_					,A_					,A_					,A_ )


						@slow
						def   _SCREAMING_SNAKE_CASE (	self       : Optional[Any] )       ->				Tuple:
							for model_name in NAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
								A  				=				NatModel.from_pretrained(A_ )
								self.assertIsNotNone(A_ )




						def   _SCREAMING_SNAKE_CASE (	self       : List[str] )       ->				List[Any]:
							A				, A  				=				self.model_tester.prepare_config_and_inputs_for_common()

							A  				=				_config_zero_init(A_ )
							for model_class in self.all_model_classes:
								A  				=				model_class(config=A_ )
								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      ):



						'''simple docstring'''


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




						@slow
						def   _SCREAMING_SNAKE_CASE (	self       : Optional[Any] )       ->				Union[str, Any]:
							A  				=				NatForImageClassification.from_pretrained('shi-labs/nat-mini-in1k-224' ).to(A_ )
							A  				=				self.default_image_processor

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

							# forward pass
							with torch.no_grad():
								A  				=				model(**A_ )

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



@require_torch
@require_natten
class 			lowerCAmelCase_			( unittest.TestCase						,   _lowercase      ):



						'''simple docstring'''


						_lowerCamelCase:      List[str] =      (NatBackbone,) if is_torch_available() else ()
						_lowerCamelCase:      Tuple =      NatConfig


						def   _SCREAMING_SNAKE_CASE (	self       : Tuple )       ->				List[str]:
							A  				=				NatModelTester(self )

"""simple docstring"""


import argparse
import json
import os.path
from collections import OrderedDict

import numpy as np
import requests
import torch
from flax.training.checkpoints import restore_checkpoint
from huggingface_hub import hf_hub_download

from transformers import VivitConfig, VivitForVideoClassification, VivitImageProcessor
from transformers.image_utils import PILImageResampling






def  _snake_case			(						snake_case__				:       Tuple       ):
	A  				=				'https://storage.googleapis.com/scenic-bucket/vivit/kinetics_400/vivit_base_16x2_unfactorized/checkpoint'

	with open(snake_case__		, 'wb'       ) as f:
		with requests.get(snake_case__		, stream=snake_case__       ) as req:
			for chunk in req.iter_content(chunk_size=2048       ):
				f.write(snake_case__       )






def  _snake_case			(						):
	A  				=				VivitConfig()

	A  				=				400
	A  				=				'huggingface/label-files'
	A  				=				'kinetics400-id2label.json'

	A  				=				json.load(open(hf_hub_download(snake_case__		, snake_case__		, repo_type='dataset'       )		, 'r'       )       )
	A  				=				{int(snake_case__       ): v for k, v in idalabel.items()}
	A  				=				idalabel
	A  				=				{v: k for k, v in idalabel.items()}
	return config






def  _snake_case			(						):
	A  				=				hf_hub_download(
	    repo_id='hf-internal-testing/spaghetti-video'		, filename='eating_spaghetti_32_frames.npy'		, repo_type='dataset'       )
	A  				=				np.load(snake_case__       )
	return list(snake_case__       )






def  _snake_case			(						snake_case__				:       np.ndarray       ):
	if np.ndim(snake_case__       ) == 2:
		return transform_attention_bias(snake_case__       )

	elif np.ndim(snake_case__       ) == 3:
		return transform_attention_kernel(snake_case__       )

	else:
		raise Exception(F'Invalid number of dimesions: {np.ndim(snake_case__       )}'       )






def  _snake_case			(						snake_case__				:       np.ndarray       ):
	return current.flatten()






def  _snake_case			(						snake_case__				:       np.ndarray       ):
	return np.reshape(snake_case__		, (current.shape[0], current.shape[1] * current.shape[2])       ).T






def  _snake_case			(						snake_case__				:       np.ndarray       ):
	return np.reshape(snake_case__		, (current.shape[0] * current.shape[1], current.shape[2])       ).T






def  _snake_case			(						snake_case__				:       str		, snake_case__				:       List[str]       ):
	A  				=				state_dict['optimizer']['target']['Transformer'][F'encoderblock_{i}']

	A  				=				F'encoder.layer.{i}.'
	A  				=				{
	    prefix + 'intermediate.dense.bias': state['MlpBlock_0']['Dense_0']['bias'],
	    prefix + 'intermediate.dense.weight': np.transpose(state['MlpBlock_0']['Dense_0']['kernel']       ),
	    prefix + 'output.dense.bias': state['MlpBlock_0']['Dense_1']['bias'],
	    prefix + 'output.dense.weight': np.transpose(state['MlpBlock_0']['Dense_1']['kernel']       ),
	    prefix + 'layernorm_before.bias': state['LayerNorm_0']['bias'],
	    prefix + 'layernorm_before.weight': state['LayerNorm_0']['scale'],
	    prefix + 'layernorm_after.bias': state['LayerNorm_1']['bias'],
	    prefix + 'layernorm_after.weight': state['LayerNorm_1']['scale'],
	    prefix
	    + 'attention.attention.query.bias': transform_attention(
	        state['MultiHeadDotProductAttention_0']['query']['bias']       ),
	    prefix
	    + 'attention.attention.query.weight': transform_attention(
	        state['MultiHeadDotProductAttention_0']['query']['kernel']       ),
	    prefix
	    + 'attention.attention.key.bias': transform_attention(state['MultiHeadDotProductAttention_0']['key']['bias']       ),
	    prefix
	    + 'attention.attention.key.weight': transform_attention(
	        state['MultiHeadDotProductAttention_0']['key']['kernel']       ),
	    prefix
	    + 'attention.attention.value.bias': transform_attention(
	        state['MultiHeadDotProductAttention_0']['value']['bias']       ),
	    prefix
	    + 'attention.attention.value.weight': transform_attention(
	        state['MultiHeadDotProductAttention_0']['value']['kernel']       ),
	    prefix + 'attention.output.dense.bias': state['MultiHeadDotProductAttention_0']['out']['bias'],
	    prefix
	    + 'attention.output.dense.weight': transform_attention_output_weight(
	        state['MultiHeadDotProductAttention_0']['out']['kernel']       ),
	}

	return new_state






def  _snake_case			(						snake_case__				:       Dict       ):
	return sum([1 if 'encoderblock_' in k else 0 for k in state_dict['optimizer']['target']['Transformer'].keys()]       )






def  _snake_case			(						snake_case__				:       Union[str, Any]		, snake_case__				:       Optional[int]=False       ):
	A  				=				get_n_layers(snake_case__       )

	A  				=				OrderedDict()

	A  				=				state_dict['optimizer']['target']['Transformer']['encoder_norm']['bias']
	A  				=				state_dict['optimizer']['target']['Transformer']['encoder_norm']['scale']

	A  				=				np.transpose(
	    state_dict['optimizer']['target']['embedding']['kernel']		, (4, 3, 0, 1, 2)       )
	A  				=				state_dict['optimizer']['target']['embedding']['bias']

	A  				=				state_dict['optimizer']['target']['cls']
	A  				=				state_dict['optimizer']['target']['Transformer']['posembed_input'][
	    'pos_embedding'
	]

	for i in range(snake_case__       ):
		new_state.update(transform_state_encoder_block(snake_case__		, snake_case__       )       )

	if classification_head:
		A  				=				{'vivit.' + k: v for k, v in new_state.items()}
		A  				=				np.transpose(state_dict['optimizer']['target']['output_projection']['kernel']       )
		A  				=				np.transpose(state_dict['optimizer']['target']['output_projection']['bias']       )

	return {k: torch.tensor(snake_case__       ) for k, v in new_state.items()}






def  _snake_case			(						):
	A  				=				VivitImageProcessor()

	assert extractor.do_resize is True
	assert extractor.size == {"shortest_edge": 256}
	assert extractor.do_center_crop is True
	assert extractor.crop_size == {"width": 224, "height": 224}
	assert extractor.resample == PILImageResampling.BILINEAR

	# here: https://github.com/deepmind/dmvr/blob/master/dmvr/modalities.py
	# one can seen that add_image has default values for normalization_mean and normalization_std set to 0 and 1
	# which effectively means no normalization (and ViViT does not overwrite those when calling this func)
	assert extractor.do_normalize is False
	assert extractor.do_rescale is True
	assert extractor.rescale_factor == 1 / 255

	# zero-centering = True in original implementation
	assert extractor.do_zero_centering is True

	return extractor






def  _snake_case			(						snake_case__				:       str       ):
	A  				=				'checkpoint'

	if not os.path.exists(snake_case__       ):
		download_checkpoint(snake_case__       )

	A  				=				restore_checkpoint(snake_case__		, snake_case__       )
	A  				=				transform_state(snake_case__		, classification_head=snake_case__       )

	A  				=				get_vivit_config()

	assert config.image_size == 224
	assert config.num_frames == 32

	A  				=				VivitForVideoClassification(snake_case__       )
	model.load_state_dict(snake_case__       )
	model.eval()

	A  				=				get_processor()

	A  				=				prepare_video()
	A  				=				extractor(snake_case__		, return_tensors='pt'       )

	A  				=				model(**snake_case__       )

	A  				=				torch.Size([1, 400]       )
	A  				=				torch.tensor([-1.0543, 2.0764, -0.2104, 0.4439, -0.9658]       )

	assert outputs.logits.shape == expected_shape
	assert torch.allclose(outputs.logits[0, :5]		, snake_case__		, atol=1e-4       ), outputs.logits[0, :5]

	model.save_pretrained(snake_case__       )
	extractor.save_pretrained(snake_case__       )


if __name__ == "__main__":
		_lowercase											=			argparse.ArgumentParser()

		parser.add_argument('''--output_model_name''', '''-o''', type=str, help='''Output path for the converted HuggingFace model''')

		_lowercase											=			parser.parse_args()
		convert(args.output_model_name)