



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





							def __init__(				self						:						List[Any] ,			__lowercase						:						Any ,			__lowercase						:						Optional[int]=2 ,			__lowercase						:						List[str]=True ,			__lowercase						:						Optional[Any]=True ,			__lowercase						:						Optional[Any]=10 ,			__lowercase						:						Optional[int]=10 ,			__lowercase						:						Any=8 ,			__lowercase						:						Dict=[2, 4, 4] ,			__lowercase						:						Optional[Any]=3 ,			__lowercase						:						List[Any]=768 ,			__lowercase						:						Any=5 ,			__lowercase						:						int=4 ,			__lowercase						:						str=37 ,			__lowercase						:						str="gelu_fast" ,			__lowercase						:						Any=0.0 ,			__lowercase						:						List[str]=0.0 ,			__lowercase						:						Union[str, Any]=0.02 ,			__lowercase						:						Optional[int]=1E-06 ,			__lowercase						:						Dict=True ,			__lowercase						:						List[Any]=None ,			):

											'''simple docstring'''
											__a					       =						parent
											__a					       =						batch_size
											__a					       =						is_training
											__a					       =						use_labels
											__a					       =						num_labels
											__a					       =						image_size
											__a					       =						num_frames
											__a					       =						tubelet_size
											__a					       =						num_channels
											__a					       =						hidden_size
											__a					       =						num_hidden_layers
											__a					       =						num_attention_heads
											__a					       =						intermediate_size
											__a					       =						hidden_act
											__a					       =						hidden_dropout_prob
											__a					       =						attention_probs_dropout_prob
											__a					       =						initializer_range
											__a					       =						layer_norm_eps
											__a					       =						qkv_bias
											__a					       =						scope

											__a					       =						(
											    (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   UpperCamelCase_   (				self						:						List[str]		):

											'''simple docstring'''
											__a					       =						floats_tensor(
											    [self.batch_size, self.num_frames, 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   UpperCamelCase_   (				self						:						str		):

											'''simple docstring'''
											__a					       =						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 ,			)
											__a					       =						self.num_labels
											return config





							def   UpperCamelCase_   (				self						:						Tuple ,			__lowercase						:						Union[str, Any] ,			__lowercase						:						Any ,			__lowercase						:						Any		):

											'''simple docstring'''
											__a					       =						VivitModel(config=__lowercase		)
											model.to(__lowercase		)
											model.eval()
											__a					       =						model(__lowercase		)
											self.parent.assertEqual(result.last_hidden_state.shape ,			(self.batch_size, self.seq_length, self.hidden_size)		)





							def   UpperCamelCase_   (				self						:						str ,			__lowercase						:						int ,			__lowercase						:						List[str] ,			__lowercase						:						Optional[Any]		):

											'''simple docstring'''
											__a					       =						VivitForVideoClassification(__lowercase		)
											model.to(__lowercase		)
											model.eval()

											__a					       =						model(__lowercase		)

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







							def   UpperCamelCase_   (				self						:						Optional[Any]		):

											'''simple docstring'''
											__a					       =						self.prepare_config_and_inputs()
											__a       ,   __a       ,   __a					       =						config_and_inputs
											__a					       =						{"""pixel_values""": pixel_values}
											return config, inputs_dict






@require_torch
class 				SCREAMING_SNAKE_CASE     (							lowerCamelCase__     , unittest.TestCase   ):
							__lowerCamelCase		:		List[Any]      =(VivitModel, VivitForVideoClassification) if is_torch_available() else ()

							__lowerCamelCase		:		Dict      =False
							__lowerCamelCase		:		Dict      =False
							__lowerCamelCase		:		Optional[Any]      =False
							__lowerCamelCase		:		Dict      =False





							def   UpperCamelCase_   (				self						:						Any		):

											'''simple docstring'''
											__a					       =						VivitModelTester(self		)
											__a					       =						ConfigTester(self ,			config_class=__lowercase ,			has_text_modality=__lowercase ,			hidden_size=37		)





							def   UpperCamelCase_   (				self						:						Optional[Any] ,			__lowercase						:						List[Any] ,			__lowercase						:						Union[str, Any] ,			__lowercase						:						List[Any]=False		):

											'''simple docstring'''
											__a					       =						copy.deepcopy(__lowercase		)

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

											return inputs_dict





							def   UpperCamelCase_   (				self						:						Dict		):

											'''simple docstring'''
											self.config_tester.run_common_tests()





							@unittest.skip(reason="""Vivit does not use inputs_embeds"""		)
							def   UpperCamelCase_   (				self						:						List[str]		):

											'''simple docstring'''
											pass





							def   UpperCamelCase_   (				self						:						List[str]		):

											'''simple docstring'''
											__a       ,   __a					       =						self.model_tester.prepare_config_and_inputs_for_common()

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





							def   UpperCamelCase_   (				self						:						Union[str, Any]		):

											'''simple docstring'''
											__a       ,   __a					       =						self.model_tester.prepare_config_and_inputs_for_common()

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

															__a					       =						["""pixel_values""", """head_mask"""]
															self.assertListEqual(arg_names[:2] ,			__lowercase		)





							def   UpperCamelCase_   (				self						:						List[str]		):

											'''simple docstring'''
											__a					       =						self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*__lowercase		)





							def   UpperCamelCase_   (				self						:						Optional[int]		):

											'''simple docstring'''
											__a					       =						self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_video_classification(*__lowercase		)





							@slow
							def   UpperCamelCase_   (				self						:						str		):

											'''simple docstring'''
											for model_name in VIVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															__a					       =						VivitModel.from_pretrained(__lowercase		)
															self.assertIsNotNone(__lowercase		)





							def   UpperCamelCase_   (				self						:						Tuple		):

											'''simple docstring'''
											__a       ,   __a					       =						self.model_tester.prepare_config_and_inputs_for_common()
											__a					       =						True

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

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

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

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

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

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

															__a					       =						outputs.attentions

															self.assertEqual(len(__lowercase		) ,			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   UpperCamelCase_   (				self						:						Optional[int]		):

											'''simple docstring'''
											def check_hidden_states_output(__lowercase						:						Optional[int] ,			__lowercase						:						str ,			__lowercase						:						Union[str, Any]		):
															__a					       =						model_class(__lowercase		)
															model.to(__lowercase		)
															model.eval()

															with torch.no_grad():
																			__a					       =						model(**self._prepare_for_class(__lowercase ,			__lowercase		)		)

															__a					       =						outputs.hidden_states
															__a					       =						self.model_tester.num_hidden_layers + 1
															self.assertEqual(len(__lowercase		) ,			__lowercase		)

															__a					       =						self.model_tester.seq_length

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

											__a       ,   __a					       =						self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
															__a					       =						True
															check_hidden_states_output(__lowercase ,			__lowercase ,			__lowercase		)

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

															check_hidden_states_output(__lowercase ,			__lowercase ,			__lowercase		)







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

											'''simple docstring'''
											pass



def 					lowerCAmelCase__  (							):
				"""simple docstring"""







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






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





							@cached_property
							def   UpperCamelCase_   (				self						:						Tuple		):

											'''simple docstring'''
											return VivitImageProcessor() if is_vision_available() else None







							@slow
							def   UpperCamelCase_   (				self						:						Dict		):

											'''simple docstring'''
											__a					       =						VivitForVideoClassification.from_pretrained("""google/vivit-b-16x2-kinetics400"""		).to(__lowercase		)

											__a					       =						self.default_image_processor
											__a					       =						prepare_video()
											__a					       =						image_processor(__lowercase ,			return_tensors="""pt"""		).to(__lowercase		)

											# forward pass
											with torch.no_grad():
															__a					       =						model(**__lowercase		)

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

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

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





#
# Sort large text files in a minimum amount of memory
#
import argparse
import os






class 				SCREAMING_SNAKE_CASE     :
							__lowerCamelCase		:		Tuple      ='block_{0}.dat'





							def __init__(				self						:						int ,			__lowercase						:						Dict		):

											'''simple docstring'''
											__a					       =						filename
											__a					       =						[]





							def   UpperCamelCase_   (				self						:						int ,			__lowercase						:						Dict ,			__lowercase						:						Optional[Any]		):

											'''simple docstring'''
											__a					       =						self.BLOCK_FILENAME_FORMAT.format(__lowercase		)
											with open(__lowercase ,			"""w"""		) as file:
															file.write(__lowercase		)
											self.block_filenames.append(__lowercase		)





							def   UpperCamelCase_   (				self						:						List[str]		):

											'''simple docstring'''
											return self.block_filenames





							def   UpperCamelCase_   (				self						:						Any ,			__lowercase						:						str ,			__lowercase						:						Tuple=None		):

											'''simple docstring'''
											__a					       =						0
											with open(self.filename		) as file:
															while True:
																			__a					       =						file.readlines(__lowercase		)

																			if lines == []:
																							break

																			if sort_key is None:
																							lines.sort()
																			else:
																							lines.sort(key=__lowercase		)

																			self.write_block("""""".join(__lowercase		) ,			__lowercase		)
																			i += 1







							def   UpperCamelCase_   (				self						:						List[Any]		):

											'''simple docstring'''
											map(os.remove ,			self.block_filenames		)






class 				SCREAMING_SNAKE_CASE     :





							def   UpperCamelCase_   (				self						:						Dict ,			__lowercase						:						Optional[Any]		):

											'''simple docstring'''
											__a					       =						-1
											__a					       =						None

											for i in range(len(__lowercase		)		):
															if min_str is None or choices[i] < min_str:
																			__a					       =						i

											return min_index






class 				SCREAMING_SNAKE_CASE     :





							def __init__(				self						:						Optional[int] ,			__lowercase						:						Optional[Any]		):

											'''simple docstring'''
											__a					       =						files
											__a					       =						set()
											__a					       =						len(__lowercase		)
											__a					       =						{i: None for i in range(self.num_buffers		)}





							def   UpperCamelCase_   (				self						:						int		):

											'''simple docstring'''
											return {
											    i: self.buffers[i] for i in range(self.num_buffers		) if i not in self.empty
											}





							def   UpperCamelCase_   (				self						:						List[str]		):

											'''simple docstring'''
											for i in range(self.num_buffers		):
															if self.buffers[i] is None and i not in self.empty:
																			__a					       =						self.files[i].readline()

																			if self.buffers[i] == "":
																							self.empty.add(__lowercase		)
																							self.files[i].close()

											if len(self.empty		) == self.num_buffers:
															return False

											return True







							def   UpperCamelCase_   (				self						:						List[str] ,			__lowercase						:						List[Any]		):

											'''simple docstring'''
											__a					       =						self.buffers[index]
											__a					       =						None

											return value






class 				SCREAMING_SNAKE_CASE     :





							def __init__(				self						:						Optional[int] ,			__lowercase						:						Any		):

											'''simple docstring'''
											__a					       =						merge_strategy





							def   UpperCamelCase_   (				self						:						str ,			__lowercase						:						Union[str, Any] ,			__lowercase						:						List[str] ,			__lowercase						:						Dict		):

											'''simple docstring'''
											__a					       =						FilesArray(self.get_file_handles(__lowercase ,			__lowercase		)		)
											with open(__lowercase ,			"""w""" ,			__lowercase		) as outfile:
															while buffers.refresh():
																			__a					       =						self.merge_strategy.select(buffers.get_dict()		)
																			outfile.write(buffers.unshift(__lowercase		)		)







							def   UpperCamelCase_   (				self						:						List[Any] ,			__lowercase						:						str ,			__lowercase						:						Dict		):

											'''simple docstring'''
											__a					       =						{}

											for i in range(len(__lowercase		)		):
															__a					       =						open(filenames[i] ,			"""r""" ,			__lowercase		)  # noqa: UP015

											return files






class 				SCREAMING_SNAKE_CASE     :





							def __init__(				self						:						str ,			__lowercase						:						Tuple		):

											'''simple docstring'''
											__a					       =						block_size





							def   UpperCamelCase_   (				self						:						Tuple ,			__lowercase						:						str ,			__lowercase						:						Union[str, Any]=None		):

											'''simple docstring'''
											__a					       =						self.get_number_blocks(__lowercase ,			self.block_size		)
											__a					       =						FileSplitter(__lowercase		)
											splitter.split(self.block_size ,			__lowercase		)

											__a					       =						FileMerger(NWayMerge()		)
											__a					       =						self.block_size / (num_blocks + 1)
											merger.merge(splitter.get_block_filenames() ,			filename + """.out""" ,			__lowercase		)

											splitter.cleanup()







							def   UpperCamelCase_   (				self						:						str ,			__lowercase						:						Optional[Any] ,			__lowercase						:						Dict		):

											'''simple docstring'''
											return (os.stat(__lowercase		).st_size / block_size) + 1



def 					lowerCAmelCase__  (							_SCREAMING_SNAKE_CASE    :     Tuple      ):
				"""simple docstring"""







				if string[-1].lower() == "k":
								return int(string[:-1]      ) * 1024
				elif string[-1].lower() == "m":
								return int(string[:-1]      ) * 1024 * 1024
				elif string[-1].lower() == "g":
								return int(string[:-1]      ) * 1024 * 1024 * 1024
				else:
								return int(_SCREAMING_SNAKE_CASE      )

def 					lowerCAmelCase__  (							):
				"""simple docstring"""







				__a					       =						argparse.ArgumentParser()
				parser.add_argument(
				    """-m"""			,       """--mem"""			,       help="""amount of memory to use for sorting"""			,       default="""100M"""      )
				parser.add_argument(
				    """filename"""			,       metavar="""<filename>"""			,       nargs=1			,       help="""name of file to sort"""      )
				__a					       =						parser.parse_args()

				__a					       =						ExternalSort(parse_memory(args.mem      )      )
				sorter.sort(args.filename[0]      )


if __name__ == "__main__":
						main()
