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





		def __init__(							self    :	Union[str, Any]						,  lowerCamelCase__    :	str						,  lowerCamelCase__    :	Optional[int]=2						,  lowerCamelCase__    :	Union[str, Any]=True						,  lowerCamelCase__    :	List[Any]=True						,  lowerCamelCase__    :	List[Any]=1_0						,  lowerCamelCase__    :	Any=1_0						,  lowerCamelCase__    :	Dict=8						,  lowerCamelCase__    :	Optional[Any]=[2, 4, 4]						,  lowerCamelCase__    :	List[str]=3						,  lowerCamelCase__    :	Dict=7_6_8						,  lowerCamelCase__    :	Tuple=5						,  lowerCamelCase__    :	Dict=4						,  lowerCamelCase__    :	Dict=3_7						,  lowerCamelCase__    :	List[str]="gelu_fast"						,  lowerCamelCase__    :	List[str]=0.0						,  lowerCamelCase__    :	Union[str, Any]=0.0						,  lowerCamelCase__    :	int=0.0_2						,  lowerCamelCase__    :	List[str]=1E-06						,  lowerCamelCase__    :	Optional[int]=True						,  lowerCamelCase__    :	List[str]=None						,  ):
							lowerCAmelCase		:					str  	=     parent
							lowerCAmelCase		:					int  	=     batch_size
							lowerCAmelCase		:					int  	=     is_training
							lowerCAmelCase		:					str  	=     use_labels
							lowerCAmelCase		:					Dict  	=     num_labels
							lowerCAmelCase		:					Optional[Any]  	=     image_size
							lowerCAmelCase		:					Any  	=     num_frames
							lowerCAmelCase		:					Tuple  	=     tubelet_size
							lowerCAmelCase		:					Dict  	=     num_channels
							lowerCAmelCase		:					Optional[int]  	=     hidden_size
							lowerCAmelCase		:					int  	=     num_hidden_layers
							lowerCAmelCase		:					Dict  	=     num_attention_heads
							lowerCAmelCase		:					Optional[int]  	=     intermediate_size
							lowerCAmelCase		:					str  	=     hidden_act
							lowerCAmelCase		:					Optional[int]  	=     hidden_dropout_prob
							lowerCAmelCase		:					Tuple  	=     attention_probs_dropout_prob
							lowerCAmelCase		:					Any  	=     initializer_range
							lowerCAmelCase		:					str  	=     layer_norm_eps
							lowerCAmelCase		:					Dict  	=     qkv_bias
							lowerCAmelCase		:					Optional[Any]  	=     scope

							lowerCAmelCase		:					Union[str, Any]  	=     (
							    (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       _A							(							self    :	int				):
							lowerCAmelCase		:					int  	=     floats_tensor(
							    [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size]				)

							lowerCAmelCase		:					Tuple  	=     None
							if self.use_labels:
												lowerCAmelCase		:					Any  	=     ids_tensor([self.batch_size]						,  self.num_labels				)

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

							return config, pixel_values, labels





		def       _A							(							self    :	Union[str, Any]				):
							lowerCAmelCase		:					Optional[Any]  	=     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		:					Any  	=     self.num_labels
							return config





		def       _A							(							self    :	Tuple						,  lowerCamelCase__    :	Optional[Any]						,  lowerCamelCase__    :	int						,  lowerCamelCase__    :	Dict				):
							lowerCAmelCase		:					List[str]  	=     VivitModel(config=lowerCamelCase__				)
							model.to(lowerCamelCase__				)
							model.eval()
							lowerCAmelCase		:					Union[str, Any]  	=     model(lowerCamelCase__				)
							self.parent.assertEqual(result.last_hidden_state.shape						,  (self.batch_size, self.seq_length, self.hidden_size)				)





		def       _A							(							self    :	Tuple						,  lowerCamelCase__    :	Dict						,  lowerCamelCase__    :	Any						,  lowerCamelCase__    :	Optional[Any]				):
							lowerCAmelCase		:					Tuple  	=     VivitForVideoClassification(lowerCamelCase__				)
							model.to(lowerCamelCase__				)
							model.eval()

							lowerCAmelCase		:					List[Any]  	=     model(lowerCamelCase__				)

							# verify the logits shape
							lowerCAmelCase		:					Union[str, Any]  	=     torch.Size((self.batch_size, self.num_labels)				)
							self.parent.assertEqual(result.logits.shape						,  lowerCamelCase__				)





		def       _A							(							self    :	Optional[Any]				):
							lowerCAmelCase		:					str  	=     self.prepare_config_and_inputs()
							lowerCAmelCase	,							lowerCAmelCase	,							lowerCAmelCase		:					Dict  	=     config_and_inputs
							lowerCAmelCase		:					Tuple  	=     {'''pixel_values''': pixel_values}
							return config, inputs_dict



@require_torch
class 		__magic_name__   (						snake_case,				unittest.TestCase    ):
		_lowerCAmelCase		=	(VivitModel, VivitForVideoClassification) if is_torch_available() else ()

		_lowerCAmelCase		=	False
		_lowerCAmelCase		=	False
		_lowerCAmelCase		=	False
		_lowerCAmelCase		=	False





		def       _A							(							self    :	Optional[int]				):
							lowerCAmelCase		:					Union[str, Any]  	=     VivitModelTester(self				)
							lowerCAmelCase		:					Any  	=     ConfigTester(self						,  config_class=lowerCamelCase__						,  has_text_modality=lowerCamelCase__						,  hidden_size=3_7				)





		def       _A							(							self    :	Optional[Any]						,  lowerCamelCase__    :	Dict						,  lowerCamelCase__    :	Optional[int]						,  lowerCamelCase__    :	List[Any]=False				):
							lowerCAmelCase		:					Union[str, Any]  	=     copy.deepcopy(lowerCamelCase__				)

							if return_labels:
												if model_class in get_values(lowerCamelCase__				):
																	lowerCAmelCase		:					Optional[Any]  	=     torch.zeros(
																	    self.model_tester.batch_size						,  dtype=torch.long						,  device=lowerCamelCase__				)

							return inputs_dict





		def       _A							(							self    :	Tuple				):
							self.config_tester.run_common_tests()





		@unittest.skip(reason='''Vivit does not use inputs_embeds'''				)
		def       _A							(							self    :	int				):
							pass





		def       _A							(							self    :	Union[str, Any]				):
							lowerCAmelCase	,							lowerCAmelCase		:					Optional[int]  	=     self.model_tester.prepare_config_and_inputs_for_common()

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





		def       _A							(							self    :	Dict				):
							lowerCAmelCase	,							lowerCAmelCase		:					Optional[Any]  	=     self.model_tester.prepare_config_and_inputs_for_common()

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

												lowerCAmelCase		:					Tuple  	=     ['''pixel_values''', '''head_mask''']
												self.assertListEqual(arg_names[:2]						,  lowerCamelCase__				)





		def       _A							(							self    :	List[Any]				):
							lowerCAmelCase		:					Tuple  	=     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*lowerCamelCase__				)





		def       _A							(							self    :	List[Any]				):
							lowerCAmelCase		:					Tuple  	=     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_video_classification(*lowerCamelCase__				)





		@slow
		def       _A							(							self    :	Optional[Any]				):
							for model_name in VIVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												lowerCAmelCase		:					str  	=     VivitModel.from_pretrained(lowerCamelCase__				)
												self.assertIsNotNone(lowerCamelCase__				)





		def       _A							(							self    :	List[Any]				):
							lowerCAmelCase	,							lowerCAmelCase		:					Tuple  	=     self.model_tester.prepare_config_and_inputs_for_common()
							lowerCAmelCase		:					str  	=     True

							for model_class in self.all_model_classes:
												lowerCAmelCase		:					int  	=     self.model_tester.seq_length

												lowerCAmelCase		:					int  	=     True
												lowerCAmelCase		:					Any  	=     False
												lowerCAmelCase		:					Tuple  	=     True
												lowerCAmelCase		:					List[str]  	=     model_class(lowerCamelCase__				)
												model.to(lowerCamelCase__				)
												model.eval()
												with torch.no_grad():
																	lowerCAmelCase		:					Any  	=     model(**self._prepare_for_class(lowerCamelCase__						,  lowerCamelCase__				)				)
												lowerCAmelCase		:					Any  	=     outputs.attentions
												self.assertEqual(len(lowerCamelCase__				)						,  self.model_tester.num_hidden_layers				)

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

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

												# Check attention is always last and order is fine
												lowerCAmelCase		:					Optional[Any]  	=     True
												lowerCAmelCase		:					Any  	=     True
												lowerCAmelCase		:					List[str]  	=     model_class(lowerCamelCase__				)
												model.to(lowerCamelCase__				)
												model.eval()
												with torch.no_grad():
																	lowerCAmelCase		:					int  	=     model(**self._prepare_for_class(lowerCamelCase__						,  lowerCamelCase__				)				)

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

												lowerCAmelCase		:					str  	=     outputs.attentions

												self.assertEqual(len(lowerCamelCase__				)						,  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       _A							(							self    :	Optional[Any]				):
							def check_hidden_states_output(lowerCamelCase__    :	List[Any]						,  lowerCamelCase__    :	List[str]						,  lowerCamelCase__    :	int				):
												lowerCAmelCase		:					Tuple  	=     model_class(lowerCamelCase__				)
												model.to(lowerCamelCase__				)
												model.eval()

												with torch.no_grad():
																	lowerCAmelCase		:					str  	=     model(**self._prepare_for_class(lowerCamelCase__						,  lowerCamelCase__				)				)

												lowerCAmelCase		:					Union[str, Any]  	=     outputs.hidden_states
												lowerCAmelCase		:					List[Any]  	=     self.model_tester.num_hidden_layers + 1
												self.assertEqual(len(lowerCamelCase__				)						,  lowerCamelCase__				)

												lowerCAmelCase		:					List[Any]  	=     self.model_tester.seq_length

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

							lowerCAmelCase	,							lowerCAmelCase		:					List[str]  	=     self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
												lowerCAmelCase		:					Union[str, Any]  	=     True
												check_hidden_states_output(lowerCamelCase__						,  lowerCamelCase__						,  lowerCamelCase__				)

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

												check_hidden_states_output(lowerCamelCase__						,  lowerCamelCase__						,  lowerCamelCase__				)





		@unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''				)
		def       _A							(							self    :	str				):
							pass



def     UpperCAmelCase__							(       ):


					'''simple docstring'''

					lowerCAmelCase		:					Dict  	=     hf_hub_download(
					    repo_id='''hf-internal-testing/spaghetti-video'''							,  filename='''eating_spaghetti_32_frames.npy'''							,  repo_type='''dataset'''  )
					lowerCAmelCase		:					List[str]  	=     np.load(__magic_name__  )
					return list(__magic_name__  )


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





		@cached_property
		def       _A							(							self    :	List[str]				):
							return VivitImageProcessor() if is_vision_available() else None





		@slow
		def       _A							(							self    :	List[str]				):
							lowerCAmelCase		:					Tuple  	=     VivitForVideoClassification.from_pretrained('''google/vivit-b-16x2-kinetics400'''				).to(lowerCamelCase__				)

							lowerCAmelCase		:					Any  	=     self.default_image_processor
							lowerCAmelCase		:					Any  	=     prepare_video()
							lowerCAmelCase		:					Optional[int]  	=     image_processor(lowerCamelCase__						,  return_tensors='''pt'''				).to(lowerCamelCase__				)

							# forward pass
							with torch.no_grad():
												lowerCAmelCase		:					str  	=     model(**lowerCamelCase__				)

							# verify the logits
							lowerCAmelCase		:					Union[str, Any]  	=     torch.Size((1, 4_0_0)				)
							self.assertEqual(outputs.logits.shape						,  lowerCamelCase__				)

							# taken from original model
							lowerCAmelCase		:					Optional[Any]  	=     torch.tensor([-1.0_5_4_3, 2.0_7_6_4, -0.2_1_0_4, 0.4_4_3_9, -0.9_6_5_8]				).to(lowerCamelCase__				)

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


import json
import os
from dataclasses import dataclass

import numpy as np
import pyarrow as pa

import datasets
from utils import get_duration


__SCREAMING_SNAKE_CASE			:   List[Any]         =     10_00_00_00_00_00
__SCREAMING_SNAKE_CASE			:   Optional[Any]         =     1_00_00

__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE			:   Any         =     os.path.split(__file__)
__SCREAMING_SNAKE_CASE			:   List[str]         =     os.path.join(RESULTS_BASEPATH, 'results', RESULTS_FILENAME.replace('.py', '.json'))
def     UpperCAmelCase__							(       __magic_name__						:     int							,  __magic_name__						:     int  ):


					'''simple docstring'''

					lowerCAmelCase		:					Optional[int]  	=     pa.Table.from_pydict({'''col''': [0] * chunk_size}  )
					lowerCAmelCase		:					int  	=     pa.concat_tables([table] * (num_examples // chunk_size)  )
					return datasets.Dataset(__magic_name__							,  fingerprint='''table_100B'''  )


@dataclass
class 		__magic_name__   :
		_lowerCAmelCase		=	42
		_lowerCAmelCase		=	42
		_lowerCAmelCase		=	42
		_lowerCAmelCase		=	42





		def       _A							(							self    :	Optional[int]				):
							lowerCAmelCase		:					Union[str, Any]  	=     np.random.default_rng(self.seed				)
							lowerCAmelCase		:					Tuple  	=     rng.integers(low=self.low						,  high=self.high						,  size=self.size				).tolist()





		def __iter__(							self    :	Dict				):
							return iter(self._sampled_values				)





		def __len__(							self    :	Tuple				):
							return self.size



@get_duration
def     UpperCAmelCase__							(       __magic_name__						:     datasets.Dataset  ):


					'''simple docstring'''

					lowerCAmelCase		:					Union[str, Any]  	=     dataset[0]
@get_duration
def     UpperCAmelCase__							(       __magic_name__						:     datasets.Dataset  ):


					'''simple docstring'''

					lowerCAmelCase		:					Any  	=     dataset[-1]
@get_duration
def     UpperCAmelCase__							(       __magic_name__						:     datasets.Dataset  ):


					'''simple docstring'''

					lowerCAmelCase		:					int  	=     dataset[range(len(__magic_name__  ) // 2							,  len(__magic_name__  ) // 2 + 10_24  )]
@get_duration
def     UpperCAmelCase__							(       __magic_name__						:     datasets.Dataset  ):


					'''simple docstring'''

					lowerCAmelCase		:					List[Any]  	=     dataset[RandIter(0							,  len(__magic_name__  )							,  10_24							,  seed=42  )]
def     UpperCAmelCase__							(       ):


					'''simple docstring'''

					lowerCAmelCase		:					Dict  	=     {'''num examples''': SPEED_TEST_N_EXAMPLES}
					lowerCAmelCase		:					Tuple  	=     (get_first_row, get_last_row, get_batch_of_aaaa_rows, get_batch_of_aaaa_random_rows)
					print('''generating dataset'''  )
					lowerCAmelCase		:					Dict  	=     generate_100B_dataset(num_examples=__magic_name__							,  chunk_size=__magic_name__  )
					print('''Functions'''  )
					for func in functions:
										print(func.__name__  )
										lowerCAmelCase		:					List[Any]  	=     func(__magic_name__  )

					with open(__magic_name__							,  '''wb'''  ) as f:
										f.write(json.dumps(__magic_name__  ).encode('''utf-8'''  )  )


if __name__ == "__main__":  # useful to run the profiler
		benchmark_table_100B()
