



import copy
import inspect
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import VivitConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor


if is_torch_available():
	import torch
	from torch import nn

	from transformers import MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING, VivitForVideoClassification, VivitModel
	from transformers.models.vivit.modeling_vivit import VIVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
	from transformers import VivitImageProcessor





class    __lowerCAmelCase					:




				def __init__(				self		,							_snake_case		,							_snake_case=2		,							_snake_case=True		,							_snake_case=True		,							_snake_case=10		,							_snake_case=10		,							_snake_case=8		,							_snake_case=[2, 4, 4]		,							_snake_case=3		,							_snake_case=768		,							_snake_case=5		,							_snake_case=4		,							_snake_case=37		,							_snake_case="gelu_fast"		,							_snake_case=0.0		,							_snake_case=0.0		,							_snake_case=0.02		,							_snake_case=1e-06		,							_snake_case=True		,							_snake_case=None		,							):



									"""simple docstring"""





									_lowerCAmelCase         =       parent
									_lowerCAmelCase         =       batch_size
									_lowerCAmelCase         =       is_training
									_lowerCAmelCase         =       use_labels
									_lowerCAmelCase         =       num_labels
									_lowerCAmelCase         =       image_size
									_lowerCAmelCase         =       num_frames
									_lowerCAmelCase         =       tubelet_size
									_lowerCAmelCase         =       num_channels
									_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         =       layer_norm_eps
									_lowerCAmelCase         =       qkv_bias
									_lowerCAmelCase         =       scope

									_lowerCAmelCase         =       (
									    (self.image_size // self.tubelet_size[2])
									    * (self.image_size // self.tubelet_size[1])
									    * (self.num_frames // self.tubelet_size[0])
									) + 1  # CLS token




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       floats_tensor(
									    [self.batch_size, self.num_frames, 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 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       VivitConfig(
									    num_frames=self.num_frames		,							image_size=self.image_size		,							tubelet_size=self.tubelet_size		,							num_channels=self.num_channels		,							hidden_size=self.hidden_size		,							num_hidden_layers=self.num_hidden_layers		,							num_attention_heads=self.num_attention_heads		,							intermediate_size=self.intermediate_size		,							hidden_act=self.hidden_act		,							hidden_dropout_prob=self.hidden_dropout_prob		,							attention_probs_dropout_prob=self.attention_probs_dropout_prob		,							initializer_range=self.initializer_range		,							layer_norm_eps=self.layer_norm_eps		,							qkv_bias=self.qkv_bias		,							)
									_lowerCAmelCase         =       self.num_labels
									return config




				def 				snake_case  (				self		,							_snake_case		,							_snake_case		,							_snake_case						):



									"""simple docstring"""





									_lowerCAmelCase         =       VivitModel(config=_snake_case						)
									model.to(_snake_case						)
									model.eval()
									_lowerCAmelCase         =       model(_snake_case						)
									self.parent.assertEqual(result.last_hidden_state.shape		,							(self.batch_size, self.seq_length, self.hidden_size)						)




				def 				snake_case  (				self		,							_snake_case		,							_snake_case		,							_snake_case						):



									"""simple docstring"""





									_lowerCAmelCase         =       VivitForVideoClassification(_snake_case						)
									model.to(_snake_case						)
									model.eval()

									_lowerCAmelCase         =       model(_snake_case						)

									# verify the logits shape
									_lowerCAmelCase         =       torch.Size((self.batch_size, self.num_labels)						)
									self.parent.assertEqual(result.logits.shape		,							_snake_case						)




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       self.prepare_config_and_inputs()
									_lowerCAmelCase      ,      _lowerCAmelCase      ,      _lowerCAmelCase         =       config_and_inputs
									_lowerCAmelCase         =       {"""pixel_values""": pixel_values}
									return config, inputs_dict





@require_torch
class    __lowerCAmelCase					(	lowerCamelCase__			,		unittest.TestCase			):
				__lowerCamelCase					=							(VivitModel, VivitForVideoClassification) if is_torch_available() else ()

				__lowerCamelCase					=							False
				__lowerCamelCase					=							False
				__lowerCamelCase					=							False
				__lowerCamelCase					=							False




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       VivitModelTester(self						)
									_lowerCAmelCase         =       ConfigTester(self		,							config_class=_snake_case		,							has_text_modality=_snake_case		,							hidden_size=37						)




				def 				snake_case  (				self		,							_snake_case		,							_snake_case		,							_snake_case=False						):



									"""simple docstring"""





									_lowerCAmelCase         =       copy.deepcopy(_snake_case						)

									if return_labels:
														if model_class in get_values(_snake_case						):
																			_lowerCAmelCase         =       torch.zeros(
																			    self.model_tester.batch_size		,							dtype=torch.long		,							device=_snake_case						)

									return inputs_dict




				def 				snake_case  (				self						):



									"""simple docstring"""





									self.config_tester.run_common_tests()




				@unittest.skip(reason="""Vivit does not use inputs_embeds"""						)
				def 				snake_case  (				self						):



									"""simple docstring"""





									pass




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase      ,      _lowerCAmelCase         =       self.model_tester.prepare_config_and_inputs_for_common()

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




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase      ,      _lowerCAmelCase         =       self.model_tester.prepare_config_and_inputs_for_common()

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

														_lowerCAmelCase         =       ["""pixel_values""", """head_mask"""]
														self.assertListEqual(arg_names[:2]		,							_snake_case						)




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*_snake_case						)




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_video_classification(*_snake_case						)




				@slow
				def 				snake_case  (				self						):



									"""simple docstring"""





									for model_name in VIVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
														_lowerCAmelCase         =       VivitModel.from_pretrained(_snake_case						)
														self.assertIsNotNone(_snake_case						)




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase      ,      _lowerCAmelCase         =       self.model_tester.prepare_config_and_inputs_for_common()
									_lowerCAmelCase         =       True

									for model_class in self.all_model_classes:
														_lowerCAmelCase         =       self.model_tester.seq_length

														_lowerCAmelCase         =       True
														_lowerCAmelCase         =       False
														_lowerCAmelCase         =       True
														_lowerCAmelCase         =       model_class(_snake_case						)
														model.to(_snake_case						)
														model.eval()
														with torch.no_grad():
																			_lowerCAmelCase         =       model(**self._prepare_for_class(_snake_case		,							_snake_case						)						)
														_lowerCAmelCase         =       outputs.attentions
														self.assertEqual(len(_snake_case						)		,							self.model_tester.num_hidden_layers						)

														# check that output_attentions also work using config
														del inputs_dict["output_attentions"]
														_lowerCAmelCase         =       True
														_lowerCAmelCase         =       model_class(_snake_case						)
														model.to(_snake_case						)
														model.eval()
														with torch.no_grad():
																			_lowerCAmelCase         =       model(**self._prepare_for_class(_snake_case		,							_snake_case						)						)
														_lowerCAmelCase         =       outputs.attentions
														self.assertEqual(len(_snake_case						)		,							self.model_tester.num_hidden_layers						)

														self.assertListEqual(
														    list(attentions[0].shape[-3:]						)		,							[self.model_tester.num_attention_heads, seq_len, seq_len]		,							)
														_lowerCAmelCase         =       len(_snake_case						)

														# Check attention is always last and order is fine
														_lowerCAmelCase         =       True
														_lowerCAmelCase         =       True
														_lowerCAmelCase         =       model_class(_snake_case						)
														model.to(_snake_case						)
														model.eval()
														with torch.no_grad():
																			_lowerCAmelCase         =       model(**self._prepare_for_class(_snake_case		,							_snake_case						)						)

														self.assertEqual(out_len + 1		,							len(_snake_case						)						)

														_lowerCAmelCase         =       outputs.attentions

														self.assertEqual(len(_snake_case						)		,							self.model_tester.num_hidden_layers						)

														self.assertListEqual(
														    list(self_attentions[0].shape[-3:]						)		,							[self.model_tester.num_attention_heads, seq_len, seq_len]		,							)




				def 				snake_case  (				self						):



									"""simple docstring"""





									def check_hidden_states_output(_snake_case		,							_snake_case		,							_snake_case						):
														_lowerCAmelCase         =       model_class(_snake_case						)
														model.to(_snake_case						)
														model.eval()

														with torch.no_grad():
																			_lowerCAmelCase         =       model(**self._prepare_for_class(_snake_case		,							_snake_case						)						)

														_lowerCAmelCase         =       outputs.hidden_states
														_lowerCAmelCase         =       self.model_tester.num_hidden_layers + 1
														self.assertEqual(len(_snake_case						)		,							_snake_case						)

														_lowerCAmelCase         =       self.model_tester.seq_length

														self.assertListEqual(
														    list(hidden_states[0].shape[-2:]						)		,							[seq_length, self.model_tester.hidden_size]		,							)

									_lowerCAmelCase      ,      _lowerCAmelCase         =       self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
														_lowerCAmelCase         =       True
														check_hidden_states_output(_snake_case		,							_snake_case		,							_snake_case						)

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

														check_hidden_states_output(_snake_case		,							_snake_case		,							_snake_case						)




				@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""						)
				def 				snake_case  (				self						):



									"""simple docstring"""





									pass


def 						_UpperCAmelCase							( ):





					"""simple docstring"""

					_lowerCAmelCase         =       hf_hub_download(
					    repo_id="""hf-internal-testing/spaghetti-video"""      ,		filename="""eating_spaghetti_32_frames.npy"""      ,		repo_type="""dataset"""     )
					_lowerCAmelCase         =       np.load(snake_case     )
					return list(snake_case     )





@require_torch
@require_vision
class    __lowerCAmelCase					(	unittest.TestCase			):




				@cached_property
				def 				snake_case  (				self						):



									"""simple docstring"""





									return VivitImageProcessor() if is_vision_available() else None




				@slow
				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       VivitForVideoClassification.from_pretrained("""google/vivit-b-16x2-kinetics400"""						).to(_snake_case						)

									_lowerCAmelCase         =       self.default_image_processor
									_lowerCAmelCase         =       prepare_video()
									_lowerCAmelCase         =       image_processor(_snake_case		,							return_tensors="""pt"""						).to(_snake_case						)

									# forward pass
									with torch.no_grad():
														_lowerCAmelCase         =       model(**_snake_case						)

									# verify the logits
									_lowerCAmelCase         =       torch.Size((1, 400)						)
									self.assertEqual(outputs.logits.shape		,							_snake_case						)

									# taken from original model
									_lowerCAmelCase         =       torch.tensor([-1.0543, 2.0764, -0.2104, 0.4439, -0.9658]						).to(_snake_case						)

									self.assertTrue(torch.allclose(outputs.logits[0, :5]		,							_snake_case		,							atol=1e-4						)						)








import json
import os
from collections.abc import Mapping
from typing import Dict, List, Optional, Union

import numpy as np
import torch

from .imports import is_safetensors_available


def 						_UpperCAmelCase							( snake_case      ,		snake_case      ,		snake_case      ,		snake_case=None     ):





					"""simple docstring"""

					_lowerCAmelCase         =       None
					# Check the string instead of the dtype to be compatible with versions of PyTorch that don't have bfloat16.
					if str(weight.dtype     ) == "torch.bfloat16":
										# Need to reinterpret the underlined data as int16 since NumPy does not handle bfloat16s.
										_lowerCAmelCase         =       weight.view(torch.intaa     )
										_lowerCAmelCase         =       """bfloat16"""
					_lowerCAmelCase         =       weight.cpu().numpy()
					_lowerCAmelCase         =       os.path.join(snake_case      ,		F'{weight_name}.dat'     )
					if index is not None:
										if dtype is None:
															_lowerCAmelCase         =       str(array.dtype     )
										_lowerCAmelCase         =       {"""dtype""": dtype, """shape""": list(array.shape     )}
					if array.ndim == 0:
										_lowerCAmelCase         =       array[None]
					_lowerCAmelCase         =       np.memmap(snake_case      ,		dtype=array.dtype      ,		mode="""w+"""      ,		shape=array.shape     )
					_lowerCAmelCase         =       array[:]
					file_array.flush()
					return index


def 						_UpperCAmelCase							( snake_case      ,		snake_case     ):





					"""simple docstring"""

					_lowerCAmelCase         =       tuple(weight_info["""shape"""]     )
					if shape == ():
										# NumPy memory-mapped arrays can't have 0 dims so it was saved as 1d tensor
										_lowerCAmelCase         =       (1,)

					_lowerCAmelCase         =       weight_info["""dtype"""]
					if dtype == "bfloat16":
										# NumPy does not support bfloat16 so this was saved as a int16
										_lowerCAmelCase         =       """int16"""

					_lowerCAmelCase         =       np.memmap(snake_case      ,		dtype=snake_case      ,		shape=snake_case      ,		mode="""r"""     )

					if len(weight_info["""shape"""]     ) == 0:
										_lowerCAmelCase         =       weight[0]
					_lowerCAmelCase         =       torch.tensor(snake_case     )
					if weight_info["dtype"] == "bfloat16":
										_lowerCAmelCase         =       weight.view(torch.bfloataa     )

					return weight


def 						_UpperCAmelCase							( snake_case      ,		snake_case     ):





					"""simple docstring"""

					if index is None or len(snake_case     ) == 0:
										# Nothing to save
										return

					_lowerCAmelCase         =       os.path.join(snake_case      ,		"""index.json"""     )
					if os.path.isfile(snake_case     ):
										with open(snake_case      ,		"""r"""      ,		encoding="""utf-8"""     ) as f:
															_lowerCAmelCase         =       json.load(snake_case     )
					else:
										_lowerCAmelCase         =       {}
					current_index.update(snake_case     )

					with open(snake_case      ,		"""w"""      ,		encoding="""utf-8"""     ) as f:
										json.dump(snake_case      ,		snake_case      ,		indent=2     )


def 						_UpperCAmelCase							( snake_case      ,		snake_case     ):





					"""simple docstring"""

					os.makedirs(snake_case      ,		exist_ok=snake_case     )
					_lowerCAmelCase         =       {}
					for name, parameter in state_dict.items():
										_lowerCAmelCase         =       offload_weight(snake_case      ,		snake_case      ,		snake_case      ,		index=snake_case     )

					# Update index
					save_offload_index(snake_case      ,		snake_case     )





class    __lowerCAmelCase					(	lowerCamelCase__			):




				def __init__(				self		,							_snake_case		,							_snake_case						):



									"""simple docstring"""





									_lowerCAmelCase         =       dataset
									_lowerCAmelCase         =       prefix




				def __getitem__(				self		,							_snake_case						):



									"""simple docstring"""





									return self.dataset[F'{self.prefix}{key}']




				def __iter__(				self						):



									"""simple docstring"""





									return iter([key for key in self.dataset if key.startswith(self.prefix						)]						)




				def __len__(				self						):



									"""simple docstring"""





									return len(self.dataset						)





class    __lowerCAmelCase					(	lowerCamelCase__			):




				def __init__(				self		,							_snake_case = None		,							_snake_case = None		,							_snake_case = None		,							_snake_case=None		,							):



									"""simple docstring"""





									if state_dict is None and save_folder is None:
														raise ValueError("""Need either a `state_dict` or a `save_folder` containing offloaded weights."""						)

									_lowerCAmelCase         =       {} if state_dict is None else state_dict
									_lowerCAmelCase         =       save_folder
									if index is None and save_folder is not None:
														with open(os.path.join(_snake_case		,							"""index.json"""						)						) as f:
																			_lowerCAmelCase         =       json.load(_snake_case						)
									_lowerCAmelCase         =       {} if index is None else index
									_lowerCAmelCase         =       list(self.state_dict.keys()						)
									self.all_keys.extend([key for key in self.index if key not in self.all_keys]						)
									_lowerCAmelCase         =       device




				def __getitem__(				self		,							_snake_case						):



									"""simple docstring"""





									if key in self.state_dict:
														return self.state_dict[key]
									_lowerCAmelCase         =       self.index[key]
									if weight_info.get("""safetensors_file"""						) is not None:
														if not is_safetensors_available():
																			raise ImportError("""These offloaded weights require the use of safetensors: `pip install safetensors`."""						)

														from safetensors import safe_open

														_lowerCAmelCase         =       """cpu""" if self.device is None else self.device
														with safe_open(weight_info["""safetensors_file"""]		,							framework="""pt"""		,							device=_snake_case						) as f:
																			_lowerCAmelCase         =       f.get_tensor(weight_info.get("""weight_name"""		,							_snake_case						)						)

														if "dtype" in weight_info:
																			return tensor.to(getattr(_snake_case		,							weight_info["""dtype"""]						)						)
														else:
																			return tensor

									_lowerCAmelCase         =       os.path.join(self.save_folder		,							F'{key}.dat'						)
									return load_offloaded_weight(_snake_case		,							_snake_case						)




				def __iter__(				self						):



									"""simple docstring"""





									return iter(self.all_keys						)




				def __len__(				self						):



									"""simple docstring"""





									return len(self.all_keys						)


def 						_UpperCAmelCase							( snake_case      ,		snake_case     ):





					"""simple docstring"""

					_lowerCAmelCase         =       {}
					for module_name in submodule_names:
										# We want to catch module_name parameter (module_name.xxx) or potentially module_name, but not any of the
										# submodules that could being like module_name (transformers.h.1 and transformers.h.10 for instance)
										result.update(
										    {
										        key: param
										        for key, param in state_dict.items()
										        if key == module_name or key.startswith(module_name + """."""     )
										    }     )
					return result



