





'''simple docstring'''






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

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


					"""simple docstring"""

					__snake_case				:			Tuple       			=    parent
					__snake_case				:			Optional[Any]       			=    batch_size
					__snake_case				:			int       			=    is_training
					__snake_case				:			Any       			=    use_labels
					__snake_case				:			Optional[Any]       			=    num_labels
					__snake_case				:			List[Any]       			=    image_size
					__snake_case				:			Optional[int]       			=    num_frames
					__snake_case				:			List[Any]       			=    tubelet_size
					__snake_case				:			str       			=    num_channels
					__snake_case				:			Tuple       			=    hidden_size
					__snake_case				:			int       			=    num_hidden_layers
					__snake_case				:			Optional[int]       			=    num_attention_heads
					__snake_case				:			Union[str, Any]       			=    intermediate_size
					__snake_case				:			Tuple       			=    hidden_act
					__snake_case				:			List[Any]       			=    hidden_dropout_prob
					__snake_case				:			Union[str, Any]       			=    attention_probs_dropout_prob
					__snake_case				:			Tuple       			=    initializer_range
					__snake_case				:			Any       			=    layer_norm_eps
					__snake_case				:			Optional[Any]       			=    qkv_bias
					__snake_case				:			str       			=    scope

					__snake_case				:			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   lowercase__    (      self				:      str )				->						Optional[Any]:


					"""simple docstring"""

					__snake_case				:			List[Any]       			=    floats_tensor(
					    [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size] )

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

					__snake_case				:			Dict       			=    self.get_config()

					return config, pixel_values, labels

		def   lowercase__    (      self				:      Union[str, Any] )				->						Union[str, Any]:


					"""simple docstring"""

					__snake_case				:			int       			=    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     ,			)
					__snake_case				:			str       			=    self.num_labels
					return config

		def   lowercase__    (      self				:      List[str]     ,			__magic_name__				:      int     ,			__magic_name__				:      Optional[int]     ,			__magic_name__				:      Any )				->						Tuple:


					"""simple docstring"""

					__snake_case				:			Dict       			=    VivitModel(config=__magic_name__ )
					model.to(__magic_name__ )
					model.eval()
					__snake_case				:			Union[str, Any]       			=    model(__magic_name__ )
					self.parent.assertEqual(result.last_hidden_state.shape     ,			(self.batch_size, self.seq_length, self.hidden_size) )

		def   lowercase__    (      self				:      Tuple     ,			__magic_name__				:      List[Any]     ,			__magic_name__				:      List[str]     ,			__magic_name__				:      Optional[int] )				->						Union[str, Any]:


					"""simple docstring"""

					__snake_case				:			List[Any]       			=    VivitForVideoClassification(__magic_name__ )
					model.to(__magic_name__ )
					model.eval()

					__snake_case				:			List[str]       			=    model(__magic_name__ )

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

		def   lowercase__    (      self				:      Tuple )				->						Optional[int]:


					"""simple docstring"""

					__snake_case				:			Union[str, Any]       			=    self.prepare_config_and_inputs()
					__snake_case			,    __snake_case			,    __snake_case				:			str       			=    config_and_inputs
					__snake_case				:			str       			=    {"""pixel_values""": pixel_values}
					return config, inputs_dict





@require_torch
class    _A   (				__lowercase	,				unittest.TestCase		):
		lowercase__:      Dict							    =      (VivitModel, VivitForVideoClassification) if is_torch_available() else ()

		lowercase__:      int							    =      False
		lowercase__:      Union[str, Any]							    =      False
		lowercase__:      str							    =      False
		lowercase__:      Optional[Any]							    =      False

		def   lowercase__    (      self				:      Optional[Any] )				->						str:


					"""simple docstring"""

					__snake_case				:			Optional[int]       			=    VivitModelTester(self )
					__snake_case				:			int       			=    ConfigTester(self     ,			config_class=__magic_name__     ,			has_text_modality=__magic_name__     ,			hidden_size=37 )

		def   lowercase__    (      self				:      Optional[int]     ,			__magic_name__				:      Union[str, Any]     ,			__magic_name__				:      Any     ,			__magic_name__				:      Union[str, Any]=False )				->						Optional[int]:


					"""simple docstring"""

					__snake_case				:			Optional[Any]       			=    copy.deepcopy(__magic_name__ )

					if return_labels:
								if model_class in get_values(__magic_name__ ):
											__snake_case				:			Dict       			=    torch.zeros(
											    self.model_tester.batch_size     ,			dtype=torch.long     ,			device=__magic_name__ )

					return inputs_dict

		def   lowercase__    (      self				:      Tuple )				->						List[Any]:


					"""simple docstring"""

					self.config_tester.run_common_tests()

		@unittest.skip(reason="""Vivit does not use inputs_embeds""" )
		def   lowercase__    (      self				:      Any )				->						int:


					"""simple docstring"""

					pass

		def   lowercase__    (      self				:      Optional[int] )				->						Optional[Any]:


					"""simple docstring"""

					__snake_case			,    __snake_case				:			Optional[Any]       			=    self.model_tester.prepare_config_and_inputs_for_common()

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

		def   lowercase__    (      self				:      int )				->						Dict:


					"""simple docstring"""

					__snake_case			,    __snake_case				:			Any       			=    self.model_tester.prepare_config_and_inputs_for_common()

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

								__snake_case				:			int       			=    ["""pixel_values""", """head_mask"""]
								self.assertListEqual(arg_names[:2]     ,			__magic_name__ )

		def   lowercase__    (      self				:      str )				->						Union[str, Any]:


					"""simple docstring"""

					__snake_case				:			str       			=    self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*__magic_name__ )

		def   lowercase__    (      self				:      List[str] )				->						List[str]:


					"""simple docstring"""

					__snake_case				:			Dict       			=    self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_video_classification(*__magic_name__ )

		@slow
		def   lowercase__    (      self				:      str )				->						Optional[Any]:


					"""simple docstring"""

					for model_name in VIVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
								__snake_case				:			int       			=    VivitModel.from_pretrained(__magic_name__ )
								self.assertIsNotNone(__magic_name__ )

		def   lowercase__    (      self				:      Tuple )				->						str:


					"""simple docstring"""

					__snake_case			,    __snake_case				:			Tuple       			=    self.model_tester.prepare_config_and_inputs_for_common()
					__snake_case				:			Any       			=    True

					for model_class in self.all_model_classes:
								__snake_case				:			List[Any]       			=    self.model_tester.seq_length

								__snake_case				:			Any       			=    True
								__snake_case				:			Union[str, Any]       			=    False
								__snake_case				:			Optional[Any]       			=    True
								__snake_case				:			Tuple       			=    model_class(__magic_name__ )
								model.to(__magic_name__ )
								model.eval()
								with torch.no_grad():
											__snake_case				:			Dict       			=    model(**self._prepare_for_class(__magic_name__     ,			__magic_name__ ) )
								__snake_case				:			List[Any]       			=    outputs.attentions
								self.assertEqual(len(__magic_name__ )     ,			self.model_tester.num_hidden_layers )

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

								self.assertListEqual(
								    list(attentions[0].shape[-3:] )     ,			[self.model_tester.num_attention_heads, seq_len, seq_len]     ,			)
								__snake_case				:			List[Any]       			=    len(__magic_name__ )

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

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

								__snake_case				:			List[Any]       			=    outputs.attentions

								self.assertEqual(len(__magic_name__ )     ,			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   lowercase__    (      self				:      Any )				->						List[str]:


					"""simple docstring"""

					def check_hidden_states_output(__magic_name__				:      Union[str, Any]     ,			__magic_name__				:      Optional[Any]     ,			__magic_name__				:      int ):
								__snake_case				:			str       			=    model_class(__magic_name__ )
								model.to(__magic_name__ )
								model.eval()

								with torch.no_grad():
											__snake_case				:			Union[str, Any]       			=    model(**self._prepare_for_class(__magic_name__     ,			__magic_name__ ) )

								__snake_case				:			Any       			=    outputs.hidden_states
								__snake_case				:			str       			=    self.model_tester.num_hidden_layers + 1
								self.assertEqual(len(__magic_name__ )     ,			__magic_name__ )

								__snake_case				:			Dict       			=    self.model_tester.seq_length

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

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

					for model_class in self.all_model_classes:
								__snake_case				:			Optional[int]       			=    True
								check_hidden_states_output(__magic_name__     ,			__magic_name__     ,			__magic_name__ )

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

								check_hidden_states_output(__magic_name__     ,			__magic_name__     ,			__magic_name__ )

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


					"""simple docstring"""

					pass







def 				_a	(					)   ->							str:



			"""simple docstring"""


			__snake_case				:			int       			=    hf_hub_download(
			    repo_id="""hf-internal-testing/spaghetti-video"""						,  filename="""eating_spaghetti_32_frames.npy"""						,  repo_type="""dataset"""   )
			__snake_case				:			str       			=    np.load(_lowerCamelCase   )
			return list(_lowerCamelCase   )




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

		@cached_property
		def   lowercase__    (      self				:      Any )				->						Dict:


					"""simple docstring"""

					return VivitImageProcessor() if is_vision_available() else None

		@slow
		def   lowercase__    (      self				:      Optional[Any] )				->						Optional[Any]:


					"""simple docstring"""

					__snake_case				:			Tuple       			=    VivitForVideoClassification.from_pretrained("""google/vivit-b-16x2-kinetics400""" ).to(__magic_name__ )

					__snake_case				:			int       			=    self.default_image_processor
					__snake_case				:			List[Any]       			=    prepare_video()
					__snake_case				:			List[Any]       			=    image_processor(__magic_name__     ,			return_tensors="""pt""" ).to(__magic_name__ )

					# forward pass
					with torch.no_grad():
								__snake_case				:			Tuple       			=    model(**__magic_name__ )

					# verify the logits
					__snake_case				:			Dict       			=    torch.Size((1, 4_00) )
					self.assertEqual(outputs.logits.shape     ,			__magic_name__ )

					# taken from original model
					__snake_case				:			Optional[Any]       			=    torch.tensor([-1.0543, 2.0764, -0.2104, 0.4439, -0.9658] ).to(__magic_name__ )

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









'''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 				_a	(					_lowerCamelCase   )   ->							Optional[Any]:



			"""simple docstring"""


			__snake_case				:			str       			=    """https://storage.googleapis.com/scenic-bucket/vivit/kinetics_400/vivit_base_16x2_unfactorized/checkpoint"""

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







def 				_a	(					)   ->							VivitConfig:



			"""simple docstring"""


			__snake_case				:			str       			=    VivitConfig()

			__snake_case				:			int       			=    400
			__snake_case				:			str       			=    """huggingface/label-files"""
			__snake_case				:			Any       			=    """kinetics400-id2label.json"""

			__snake_case				:			str       			=    json.load(open(hf_hub_download(_lowerCamelCase						,  _lowerCamelCase						,  repo_type="""dataset"""   )						,  """r"""   )   )
			__snake_case				:			Optional[int]       			=    {int(_lowerCamelCase   ): v for k, v in idalabel.items()}
			__snake_case				:			Optional[Any]       			=    idalabel
			__snake_case				:			Dict       			=    {v: k for k, v in idalabel.items()}
			return config







def 				_a	(					)   ->							Any:



			"""simple docstring"""


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







def 				_a	(					_lowerCamelCase   )   ->							Union[str, Any]:



			"""simple docstring"""


			if np.ndim(_lowerCamelCase   ) == 2:
						return transform_attention_bias(_lowerCamelCase   )

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

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







def 				_a	(					_lowerCamelCase   )   ->							Any:



			"""simple docstring"""


			return current.flatten()







def 				_a	(					_lowerCamelCase   )   ->							Optional[int]:



			"""simple docstring"""


			return np.reshape(_lowerCamelCase						,  (current.shape[0], current.shape[1] * current.shape[2])   ).T







def 				_a	(					_lowerCamelCase   )   ->							int:



			"""simple docstring"""


			return np.reshape(_lowerCamelCase						,  (current.shape[0] * current.shape[1], current.shape[2])   ).T







def 				_a	(					_lowerCamelCase						,  _lowerCamelCase   )   ->							Dict:



			"""simple docstring"""


			__snake_case				:			Optional[Any]       			=    state_dict["""optimizer"""]["""target"""]["""Transformer"""][F'''encoderblock_{i}''']

			__snake_case				:			str       			=    F'''encoder.layer.{i}.'''
			__snake_case				:			Optional[Any]       			=    {
			    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 				_a	(					_lowerCamelCase   )   ->							List[Any]:



			"""simple docstring"""


			return sum([1 if """encoderblock_""" in k else 0 for k in state_dict["""optimizer"""]["""target"""]["""Transformer"""].keys()]   )







def 				_a	(					_lowerCamelCase						,  _lowerCamelCase=False   )   ->							str:



			"""simple docstring"""


			__snake_case				:			Optional[int]       			=    get_n_layers(_lowerCamelCase   )

			__snake_case				:			Optional[Any]       			=    OrderedDict()

			__snake_case				:			Tuple       			=    state_dict["""optimizer"""]["""target"""]["""Transformer"""]["""encoder_norm"""]["""bias"""]
			__snake_case				:			str       			=    state_dict["""optimizer"""]["""target"""]["""Transformer"""]["""encoder_norm"""]["""scale"""]

			__snake_case				:			Optional[Any]       			=    np.transpose(
			    state_dict["""optimizer"""]["""target"""]["""embedding"""]["""kernel"""]						,  (4, 3, 0, 1, 2)   )
			__snake_case				:			Tuple       			=    state_dict["""optimizer"""]["""target"""]["""embedding"""]["""bias"""]

			__snake_case				:			Optional[int]       			=    state_dict["""optimizer"""]["""target"""]["""cls"""]
			__snake_case				:			Dict       			=    state_dict["""optimizer"""]["""target"""]["""Transformer"""]["""posembed_input"""][
			    """pos_embedding"""
			]

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

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

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







def 				_a	(					)   ->							VivitImageProcessor:



			"""simple docstring"""


			__snake_case				:			Optional[Any]       			=    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 				_a	(					_lowerCamelCase   )   ->							Optional[int]:



			"""simple docstring"""


			__snake_case				:			Optional[int]       			=    """checkpoint"""

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

			__snake_case				:			int       			=    restore_checkpoint(_lowerCamelCase						,  _lowerCamelCase   )
			__snake_case				:			Union[str, Any]       			=    transform_state(_lowerCamelCase						,  classification_head=_lowerCamelCase   )

			__snake_case				:			Optional[Any]       			=    get_vivit_config()

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

			__snake_case				:			List[Any]       			=    VivitForVideoClassification(_lowerCamelCase   )
			model.load_state_dict(_lowerCamelCase   )
			model.eval()

			__snake_case				:			int       			=    get_processor()

			__snake_case				:			int       			=    prepare_video()
			__snake_case				:			List[str]       			=    extractor(_lowerCamelCase						,  return_tensors="""pt"""   )

			__snake_case				:			Optional[int]       			=    model(**_lowerCamelCase   )

			__snake_case				:			List[Any]       			=    torch.Size([1, 400]   )
			__snake_case				:			Tuple       			=    torch.tensor([-1.05_43, 2.07_64, -0.21_04, 0.44_39, -0.96_58]   )

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

			model.save_pretrained(_lowerCamelCase   )
			extractor.save_pretrained(_lowerCamelCase   )


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

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

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


