





'''simple docstring'''






import inspect
import os
from argparse import ArgumentParser, Namespace
from importlib import import_module

import huggingface_hub
import numpy as np
from packaging import version

from .. import (
    FEATURE_EXTRACTOR_MAPPING,
    IMAGE_PROCESSOR_MAPPING,
    PROCESSOR_MAPPING,
    TOKENIZER_MAPPING,
    AutoConfig,
    AutoFeatureExtractor,
    AutoImageProcessor,
    AutoProcessor,
    AutoTokenizer,
    is_datasets_available,
    is_tf_available,
    is_torch_available,
)
from ..utils import TF2_WEIGHTS_INDEX_NAME, TF2_WEIGHTS_NAME, logging
from . import BaseTransformersCLICommand


if is_tf_available():
	import tensorflow as tf

	tf.config.experimental.enable_tensor_float_aa_execution(False)

if is_torch_available():
	import torch

if is_datasets_available():
	from datasets import load_dataset


__UpperCamelCase     :   int										=  5E-5  # larger error tolerance than in our internal tests, to avoid flaky user-facing errors







def 				_a	(					_lowerCamelCase   )   ->							List[str]:



			"""simple docstring"""


			return PTtoTFCommand(
			    args.model_name						,  args.local_dir						,  args.max_error						,  args.new_weights						,  args.no_pr						,  args.push						,  args.extra_commit_description						,  args.override_model_class						,  )




class    _A   (				a_		):

		@staticmethod
		def   lowercase__    (      __magic_name__				:      int )				->						Any:


					"""simple docstring"""

					__snake_case				:			Any       			=    parser.add_parser(
					    """pt-to-tf"""     ,			help=(
					        """CLI tool to run convert a transformers model from a PyTorch checkpoint to a TensorFlow checkpoint."""
					        """ Can also be used to validate existing weights without opening PRs, with --no-pr."""
					    )     ,			)
					train_parser.add_argument(
					    """--model-name"""     ,			type=_snake_case     ,			required=_snake_case     ,			help="""The model name, including owner/organization, as seen on the hub."""     ,			)
					train_parser.add_argument(
					    """--local-dir"""     ,			type=_snake_case     ,			default=""""""     ,			help="""Optional local directory of the model repository. Defaults to /tmp/{model_name}"""     ,			)
					train_parser.add_argument(
					    """--max-error"""     ,			type=_snake_case     ,			default=_snake_case     ,			help=(
					        f'''Maximum error tolerance. Defaults to {MAX_ERROR}. This flag should be avoided, use at your own risk.'''
					    )     ,			)
					train_parser.add_argument(
					    """--new-weights"""     ,			action="""store_true"""     ,			help="""Optional flag to create new TensorFlow weights, even if they already exist."""     ,			)
					train_parser.add_argument(
					    """--no-pr"""     ,			action="""store_true"""     ,			help="""Optional flag to NOT open a PR with converted weights.""" )
					train_parser.add_argument(
					    """--push"""     ,			action="""store_true"""     ,			help="""Optional flag to push the weights directly to `main` (requires permissions)"""     ,			)
					train_parser.add_argument(
					    """--extra-commit-description"""     ,			type=_snake_case     ,			default=""""""     ,			help="""Optional additional commit description to use when opening a PR (e.g. to tag the owner)."""     ,			)
					train_parser.add_argument(
					    """--override-model-class"""     ,			type=_snake_case     ,			default=_snake_case     ,			help="""If you think you know better than the auto-detector, you can specify the model class here. """
					    """Can be either an AutoModel class or a specific model class like BertForSequenceClassification."""     ,			)
					train_parser.set_defaults(func=_snake_case )

		@staticmethod
		def   lowercase__    (      __magic_name__				:      Optional[int]     ,			__magic_name__				:      int )				->						List[str]:


					"""simple docstring"""

					__snake_case				:			Any       			=    set(pt_outputs.keys() )
					__snake_case				:			Dict       			=    set(tf_outputs.keys() )
					if pt_out_attrs != tf_out_attrs:
								raise ValueError(
								    f'''The model outputs have different attributes, aborting. (Pytorch: {pt_out_attrs}, TensorFlow:'''
								    f''' {tf_out_attrs})''' )

					# 2. For each output attribute, computes the difference
					def _find_pt_tf_differences(__magic_name__				:      str     ,			__magic_name__				:      Any     ,			__magic_name__				:      str     ,			__magic_name__				:      Dict="" ):
								# If the current attribute is a tensor, it is a leaf and we make the comparison. Otherwise, we will dig in
								# recursivelly, keeping the name of the attribute.
								if isinstance(_snake_case     ,			torch.Tensor ):
											__snake_case				:			Optional[int]       			=    np.max(np.abs(pt_out.numpy() - tf_out.numpy() ) )
											__snake_case				:			str       			=    tensor_difference
								else:
											__snake_case				:			List[Any]       			=    attr_name
											for i, pt_item in enumerate(_snake_case ):
														# If it is a named attribute, we keep the name. Otherwise, just its index.
														if isinstance(_snake_case     ,			_snake_case ):
																	__snake_case				:			Dict       			=    root_name + pt_item
																	__snake_case				:			str       			=    tf_out[pt_item]
																	__snake_case				:			List[str]       			=    pt_out[pt_item]
														else:
																	__snake_case				:			List[str]       			=    root_name + f'''[{i}]'''
																	__snake_case				:			List[str]       			=    tf_out[i]
														__snake_case				:			List[Any]       			=    _find_pt_tf_differences(_snake_case     ,			_snake_case     ,			_snake_case     ,			_snake_case )

								return differences

					return _find_pt_tf_differences(_snake_case     ,			_snake_case     ,			{} )

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


					"""simple docstring"""

					__snake_case				:			int       			=    logging.get_logger("""transformers-cli/pt_to_tf""" )
					__snake_case				:			List[Any]       			=    model_name
					__snake_case				:			Any       			=    local_dir if local_dir else os.path.join("""/tmp"""     ,			_snake_case )
					__snake_case				:			Optional[int]       			=    max_error
					__snake_case				:			Optional[int]       			=    new_weights
					__snake_case				:			Dict       			=    no_pr
					__snake_case				:			Tuple       			=    push
					__snake_case				:			Tuple       			=    extra_commit_description
					__snake_case				:			List[Any]       			=    override_model_class

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


					"""simple docstring"""

					def _get_audio_input():
								__snake_case				:			Dict       			=    load_dataset("""hf-internal-testing/librispeech_asr_dummy"""     ,			"""clean"""     ,			split="""validation""" )
								__snake_case				:			Union[str, Any]       			=    ds.sort("""id""" ).select(range(2 ) )[:2]['''audio''']
								__snake_case				:			Dict       			=    [x['''array'''] for x in speech_samples]
								return raw_samples

					__snake_case				:			Optional[int]       			=    type(pt_model.config )
					if model_config_class in PROCESSOR_MAPPING:
								__snake_case				:			Optional[Any]       			=    AutoProcessor.from_pretrained(self._local_dir )
								if model_config_class in TOKENIZER_MAPPING and processor.tokenizer.pad_token is None:
											__snake_case				:			List[Any]       			=    processor.tokenizer.eos_token
					elif model_config_class in IMAGE_PROCESSOR_MAPPING:
								__snake_case				:			Optional[Any]       			=    AutoImageProcessor.from_pretrained(self._local_dir )
					elif model_config_class in FEATURE_EXTRACTOR_MAPPING:
								__snake_case				:			Any       			=    AutoFeatureExtractor.from_pretrained(self._local_dir )
					elif model_config_class in TOKENIZER_MAPPING:
								__snake_case				:			Any       			=    AutoTokenizer.from_pretrained(self._local_dir )
								if processor.pad_token is None:
											__snake_case				:			Optional[Any]       			=    processor.eos_token
					else:
								raise ValueError(f'''Unknown data processing type (model config type: {model_config_class})''' )

					__snake_case				:			Optional[int]       			=    set(inspect.signature(pt_model.forward ).parameters.keys() )
					__snake_case				:			Optional[int]       			=    {}
					if "input_ids" in model_forward_signature:
								processor_inputs.update(
								    {
								        """text""": ["""Hi there!""", """I am a batch with more than one row and different input lengths."""],
								        """padding""": True,
								        """truncation""": True,
								    } )
					if "pixel_values" in model_forward_signature:
								__snake_case				:			Optional[int]       			=    load_dataset("""cifar10"""     ,			"""plain_text"""     ,			split="""test""" )[:2]['''img''']
								processor_inputs.update({"""images""": sample_images} )
					if "input_features" in model_forward_signature:
								__snake_case				:			Optional[Any]       			=    inspect.signature(processor.feature_extractor ).parameters
								# Pad to the largest input length by default but take feature extractor default
								# padding value if it exists e.g. "max_length" and is not False or None
								if "padding" in feature_extractor_signature:
											__snake_case				:			Any       			=    feature_extractor_signature['''padding'''].default
											if default_strategy is not False and default_strategy is not None:
														__snake_case				:			int       			=    default_strategy
											else:
														__snake_case				:			Any       			=    True
								else:
											__snake_case				:			int       			=    True
								processor_inputs.update({"""audio""": _get_audio_input(), """padding""": padding_strategy} )
					if "input_values" in model_forward_signature:  # Wav2Vec2 audio input
								processor_inputs.update({"""audio""": _get_audio_input(), """padding""": True} )
					__snake_case				:			Optional[int]       			=    processor(**_snake_case     ,			return_tensors="""pt""" )
					__snake_case				:			List[str]       			=    processor(**_snake_case     ,			return_tensors="""tf""" )

					# Extra input requirements, in addition to the input modality
					if (
					    config.is_encoder_decoder
					    or (hasattr(_snake_case     ,			"""encoder""" ) and hasattr(_snake_case     ,			"""decoder""" ))
					    or "decoder_input_ids" in tf_dummy_inputs
					):
								__snake_case				:			Dict       			=    np.asarray([[1], [1]]     ,			dtype=_snake_case ) * (pt_model.config.decoder_start_token_id or 0)
								pt_input.update({"""decoder_input_ids""": torch.tensor(_snake_case )} )
								tf_input.update({"""decoder_input_ids""": tf.convert_to_tensor(_snake_case )} )

					return pt_input, tf_input

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


					"""simple docstring"""

					if version.parse(huggingface_hub.__version__ ) < version.parse("""0.9.0""" ):
								raise ImportError(
								    """The huggingface_hub version must be >= 0.9.0 to use this command. Please update your huggingface_hub"""
								    """ installation.""" )
					else:
								from huggingface_hub import Repository, create_commit
								from huggingface_hub._commit_api import CommitOperationAdd

					# Fetch remote data
					__snake_case				:			Dict       			=    Repository(local_dir=self._local_dir     ,			clone_from=self._model_name )

					# Load config and get the appropriate architecture -- the latter is needed to convert the head's weights
					__snake_case				:			int       			=    AutoConfig.from_pretrained(self._local_dir )
					__snake_case				:			int       			=    config.architectures
					if self._override_model_class is not None:
								if self._override_model_class.startswith("""TF""" ):
											__snake_case				:			str       			=    [self._override_model_class[2:]]
								else:
											__snake_case				:			int       			=    [self._override_model_class]
								try:
											__snake_case				:			Any       			=    getattr(import_module("""transformers""" )     ,			architectures[0] )
								except AttributeError:
											raise ValueError(f'''Model class {self._override_model_class} not found in transformers.''' )
								try:
											__snake_case				:			Union[str, Any]       			=    getattr(import_module("""transformers""" )     ,			"""TF""" + architectures[0] )
								except AttributeError:
											raise ValueError(f'''TF model class TF{self._override_model_class} not found in transformers.''' )
					elif architectures is None:  # No architecture defined -- use auto classes
								__snake_case				:			List[str]       			=    getattr(import_module("""transformers""" )     ,			"""AutoModel""" )
								__snake_case				:			Tuple       			=    getattr(import_module("""transformers""" )     ,			"""TFAutoModel""" )
								self._logger.warning("""No detected architecture, using AutoModel/TFAutoModel""" )
					else:  # Architecture defined -- use it
								if len(_snake_case ) > 1:
											raise ValueError(f'''More than one architecture was found, aborting. (architectures = {architectures})''' )
								self._logger.warning(f'''Detected architecture: {architectures[0]}''' )
								__snake_case				:			Tuple       			=    getattr(import_module("""transformers""" )     ,			architectures[0] )
								try:
											__snake_case				:			List[str]       			=    getattr(import_module("""transformers""" )     ,			"""TF""" + architectures[0] )
								except AttributeError:
											raise AttributeError(f'''The TensorFlow equivalent of {architectures[0]} doesn\'t exist in transformers.''' )

        # Check the TF dummy inputs to see what keys we need in the forward pass
					__snake_case				:			Tuple       			=    tf_class.from_config(_snake_case )
					__snake_case				:			List[Any]       			=    tf_from_pt_model.dummy_inputs

					del tf_from_pt_model  # Try to keep only one model in memory at a time

					# Load the model and get some basic inputs
					__snake_case				:			Any       			=    pt_class.from_pretrained(self._local_dir )
					pt_model.eval()

					__snake_case				:			List[Any]       			=    self.get_inputs(_snake_case     ,			_snake_case     ,			_snake_case )

					with torch.no_grad():
								__snake_case				:			List[Any]       			=    pt_model(**_snake_case     ,			output_hidden_states=_snake_case )
					del pt_model  # will no longer be used, and may have a large memory footprint

					__snake_case				:			Union[str, Any]       			=    tf_class.from_pretrained(self._local_dir     ,			from_pt=_snake_case )
					__snake_case				:			Tuple       			=    tf_from_pt_model(**_snake_case     ,			output_hidden_states=_snake_case     ,			training=_snake_case )

					# Confirms that cross loading PT weights into TF worked.
					__snake_case				:			Any       			=    self.find_pt_tf_differences(_snake_case     ,			_snake_case )
					__snake_case				:			List[str]       			=    {k: v for k, v in crossload_differences.items() if '''hidden''' not in k}
					__snake_case				:			Optional[int]       			=    {k: v for k, v in crossload_differences.items() if '''hidden''' in k}
					if len(_snake_case ) == 0 and architectures is not None:
								raise ValueError(
								    f'''Something went wrong -- the config file has architectures ({architectures}), but no model head'''
								    """ output was found. All outputs start with \'hidden\'""" )
					__snake_case				:			Union[str, Any]       			=    max(output_differences.values() ) if output_differences else 0.0
					__snake_case				:			int       			=    max(hidden_differences.values() )
					if max_crossload_output_diff > self._max_error or max_crossload_hidden_diff > self._max_error:
								raise ValueError(
								    """The cross-loaded TensorFlow model has different outputs, something went wrong!\n"""
								    + f'''\nList of maximum output differences above the threshold ({self._max_error}):\n'''
								    + """\n""".join([f'''{k}: {v:.3e}''' for k, v in output_differences.items() if v > self._max_error] )
								    + f'''\n\nList of maximum hidden layer differences above the threshold ({self._max_error}):\n'''
								    + """\n""".join([f'''{k}: {v:.3e}''' for k, v in hidden_differences.items() if v > self._max_error] ) )

					# Save the weights in a TF format (if needed) and confirms that the results are still good
					__snake_case				:			Dict       			=    os.path.join(self._local_dir     ,			_snake_case )
					__snake_case				:			Union[str, Any]       			=    os.path.join(self._local_dir     ,			_snake_case )
					if (not os.path.exists(_snake_case ) and not os.path.exists(_snake_case )) or self._new_weights:
								tf_from_pt_model.save_pretrained(self._local_dir )
					del tf_from_pt_model  # will no longer be used, and may have a large memory footprint

					__snake_case				:			Optional[int]       			=    tf_class.from_pretrained(self._local_dir )
					__snake_case				:			Any       			=    tf_model(**_snake_case     ,			output_hidden_states=_snake_case )

					__snake_case				:			int       			=    self.find_pt_tf_differences(_snake_case     ,			_snake_case )
					__snake_case				:			str       			=    {k: v for k, v in conversion_differences.items() if '''hidden''' not in k}
					__snake_case				:			List[str]       			=    {k: v for k, v in conversion_differences.items() if '''hidden''' in k}
					if len(_snake_case ) == 0 and architectures is not None:
								raise ValueError(
								    f'''Something went wrong -- the config file has architectures ({architectures}), but no model head'''
								    """ output was found. All outputs start with \'hidden\'""" )
					__snake_case				:			Any       			=    max(output_differences.values() ) if output_differences else 0.0
					__snake_case				:			Union[str, Any]       			=    max(hidden_differences.values() )
					if max_conversion_output_diff > self._max_error or max_conversion_hidden_diff > self._max_error:
								raise ValueError(
								    """The converted TensorFlow model has different outputs, something went wrong!\n"""
								    + f'''\nList of maximum output differences above the threshold ({self._max_error}):\n'''
								    + """\n""".join([f'''{k}: {v:.3e}''' for k, v in output_differences.items() if v > self._max_error] )
								    + f'''\n\nList of maximum hidden layer differences above the threshold ({self._max_error}):\n'''
								    + """\n""".join([f'''{k}: {v:.3e}''' for k, v in hidden_differences.items() if v > self._max_error] ) )

					__snake_case				:			Dict       			=    '''Update TF weights''' if self._new_weights else '''Add TF weights'''
					if self._push:
								repo.git_add(auto_lfs_track=_snake_case )
								repo.git_commit(_snake_case )
								repo.git_push(blocking=_snake_case )  # this prints a progress bar with the upload
								self._logger.warning(f'''TF weights pushed into {self._model_name}''' )
					elif not self._no_pr:
								self._logger.warning("""Uploading the weights into a new PR...""" )
								__snake_case				:			int       			=    (
								    '''Model converted by the [`transformers`\' `pt_to_tf`'''
								    ''' CLI](https://github.com/huggingface/transformers/blob/main/src/transformers/commands/pt_to_tf.py). '''
								    '''All converted model outputs and hidden layers were validated against its PyTorch counterpart.\n\n'''
								    f'''Maximum crossload output difference={max_crossload_output_diff:.3e}; '''
								    f'''Maximum crossload hidden layer difference={max_crossload_hidden_diff:.3e};\n'''
								    f'''Maximum conversion output difference={max_conversion_output_diff:.3e}; '''
								    f'''Maximum conversion hidden layer difference={max_conversion_hidden_diff:.3e};\n'''
								)
								if self._max_error > MAX_ERROR:
											commit_descrition += (
											    f'''\n\nCAUTION: The maximum admissible error was manually increased to {self._max_error}!'''
											)
								if self._extra_commit_description:
											commit_descrition += "\n\n" + self._extra_commit_description

								# sharded model -> adds all related files (index and .h5 shards)
								if os.path.exists(_snake_case ):
											__snake_case				:			int       			=    [
											    CommitOperationAdd(path_in_repo=_snake_case     ,			path_or_fileobj=_snake_case )
											]
											for shard_path in tf.io.gfile.glob(self._local_dir + """/tf_model-*.h5""" ):
														operations += [
														    CommitOperationAdd(path_in_repo=os.path.basename(_snake_case )     ,			path_or_fileobj=_snake_case )
														]
								else:
											__snake_case				:			str       			=    [CommitOperationAdd(path_in_repo=_snake_case     ,			path_or_fileobj=_snake_case )]

								__snake_case				:			List[Any]       			=    create_commit(
								    repo_id=self._model_name     ,			operations=_snake_case     ,			commit_message=_snake_case     ,			commit_description=_snake_case     ,			repo_type="""model"""     ,			create_pr=_snake_case     ,			).pr_url
								self._logger.warning(f'''PR open in {hub_pr_url}''' )









'''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)


