


"""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


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






def 					a__ (       _UpperCamelCase ):





						"""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 							lowerCamelCase       (			a_ ):


							"""simple docstring"""






							@staticmethod
							def 					lowerCAmelCase     (						lowerCamelCase_ :     Optional[Any]				):
													UpperCAmelCase__		=    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 					lowerCAmelCase     (						lowerCamelCase_ :     Any		, lowerCamelCase_ :     Optional[Any]				):
													UpperCAmelCase__		=    set(pt_outputs.keys()				)
													UpperCAmelCase__		=    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(lowerCamelCase_ :     Optional[int]		, lowerCamelCase_ :     Optional[int]		, lowerCamelCase_ :     Optional[int]		, lowerCamelCase_ :     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				):
																									UpperCAmelCase__		=    np.max(np.abs(pt_out.numpy() - tf_out.numpy()				)				)
																									UpperCAmelCase__		=    tensor_difference
																			else:
																									UpperCAmelCase__		=    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				):
																																					UpperCAmelCase__		=    root_name + pt_item
																																					UpperCAmelCase__		=    tf_out[pt_item]
																																					UpperCAmelCase__		=    pt_out[pt_item]
																															else:
																																					UpperCAmelCase__		=    root_name + f"""[{i}]"""
																																					UpperCAmelCase__		=    tf_out[i]
																															UpperCAmelCase__		=    _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 :     Dict		, lowerCamelCase_ :     Tuple		, lowerCamelCase_ :     List[str]		, lowerCamelCase_ :     Optional[int]		, lowerCamelCase_ :     List[str]		, lowerCamelCase_ :     Union[str, Any]		, lowerCamelCase_ :     Union[str, Any]		, lowerCamelCase_ :     Dict		, lowerCamelCase_ :     List[str]		, *lowerCamelCase_ :     str		, ):
													UpperCAmelCase__		=    logging.get_logger('''transformers-cli/pt_to_tf'''				)
													UpperCAmelCase__		=    model_name
													UpperCAmelCase__		=    local_dir if local_dir else os.path.join('''/tmp'''		, _snake_case				)
													UpperCAmelCase__		=    max_error
													UpperCAmelCase__		=    new_weights
													UpperCAmelCase__		=    no_pr
													UpperCAmelCase__		=    push
													UpperCAmelCase__		=    extra_commit_description
													UpperCAmelCase__		=    override_model_class






							def 					lowerCAmelCase     (						self :     str		, lowerCamelCase_ :     Tuple		, lowerCamelCase_ :     Optional[int]		, lowerCamelCase_ :     Optional[int]				):

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

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

													UpperCAmelCase__		=    set(inspect.signature(pt_model.forward				).parameters.keys()				)
													UpperCAmelCase__		=    {}
													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:
																			UpperCAmelCase__		=    load_dataset('''cifar10'''		, '''plain_text'''		, split='''test'''				)[:2]['''img''']
																			processor_inputs.update({'''images''': sample_images}				)
													if "input_features" in model_forward_signature:
																			UpperCAmelCase__		=    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:
																									UpperCAmelCase__		=    feature_extractor_signature['''padding'''].default
																									if default_strategy is not False and default_strategy is not None:
																															UpperCAmelCase__		=    default_strategy
																									else:
																															UpperCAmelCase__		=    True
																			else:
																									UpperCAmelCase__		=    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}				)
													UpperCAmelCase__		=    processor(**_snake_case		, return_tensors='''pt'''				)
													UpperCAmelCase__		=    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
													):
																			UpperCAmelCase__		=    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 					lowerCAmelCase     (						self :     Optional[int]				):
													# hub version 0.9.0 introduced the possibility of programmatically opening PRs with normal write tokens.
													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
													UpperCAmelCase__		=    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
													UpperCAmelCase__		=    AutoConfig.from_pretrained(self._local_dir				)
													UpperCAmelCase__		=    config.architectures
													if self._override_model_class is not None:
																			if self._override_model_class.startswith('''TF'''				):
																									UpperCAmelCase__		=    [self._override_model_class[2:]]
																			else:
																									UpperCAmelCase__		=    [self._override_model_class]
																			try:
																									UpperCAmelCase__		=    getattr(import_module('''transformers'''				)		, architectures[0]				)
																			except AttributeError:
																									raise ValueError(f"""Model class {self._override_model_class} not found in transformers."""				)
																			try:
																									UpperCAmelCase__		=    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
																			UpperCAmelCase__		=    getattr(import_module('''transformers'''				)		, '''AutoModel'''				)
																			UpperCAmelCase__		=    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]}"""				)
																			UpperCAmelCase__		=    getattr(import_module('''transformers'''				)		, architectures[0]				)
																			try:
																									UpperCAmelCase__		=    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
													UpperCAmelCase__		=    tf_class.from_config(_snake_case				)
													UpperCAmelCase__		=    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
													UpperCAmelCase__		=    pt_class.from_pretrained(self._local_dir				)
													pt_model.eval()

													UpperCAmelCase__		=    self.get_inputs(_snake_case		, _snake_case		, _snake_case				)

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

													UpperCAmelCase__		=    tf_class.from_pretrained(self._local_dir		, from_pt=_snake_case				)
													UpperCAmelCase__		=    tf_from_pt_model(**_snake_case		, output_hidden_states=_snake_case		, training=_snake_case				)

													# Confirms that cross loading PT weights into TF worked.
													UpperCAmelCase__		=    self.find_pt_tf_differences(_snake_case		, _snake_case				)
													UpperCAmelCase__		=    {k: v for k, v in crossload_differences.items() if '''hidden''' not in k}
													UpperCAmelCase__		=    {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\''''				)
													UpperCAmelCase__		=    max(output_differences.values()				) if output_differences else 0.0
													UpperCAmelCase__		=    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
													UpperCAmelCase__		=    os.path.join(self._local_dir		, _snake_case				)
													UpperCAmelCase__		=    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

													UpperCAmelCase__		=    tf_class.from_pretrained(self._local_dir				)
													UpperCAmelCase__		=    tf_model(**_snake_case		, output_hidden_states=_snake_case				)

													UpperCAmelCase__		=    self.find_pt_tf_differences(_snake_case		, _snake_case				)
													UpperCAmelCase__		=    {k: v for k, v in conversion_differences.items() if '''hidden''' not in k}
													UpperCAmelCase__		=    {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\''''				)
													UpperCAmelCase__		=    max(output_differences.values()				) if output_differences else 0.0
													UpperCAmelCase__		=    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]				)				)

													UpperCAmelCase__		=    '''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...'''				)
																			UpperCAmelCase__		=    (
																			    '''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				):
																									UpperCAmelCase__		=    [
																									    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:
																									UpperCAmelCase__		=    [CommitOperationAdd(path_in_repo=_snake_case		, path_or_fileobj=_snake_case				)]

																			UpperCAmelCase__		=    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__ (       _UpperCamelCase ):





						"""simple docstring"""

						UpperCAmelCase__		=    '''https://storage.googleapis.com/scenic-bucket/vivit/kinetics_400/vivit_base_16x2_unfactorized/checkpoint'''

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






def 					a__ (       ):





						"""simple docstring"""

						UpperCAmelCase__		=    VivitConfig()

						UpperCAmelCase__		=    400
						UpperCAmelCase__		=    '''huggingface/label-files'''
						UpperCAmelCase__		=    '''kinetics400-id2label.json'''

						UpperCAmelCase__		=    json.load(open(hf_hub_download(_UpperCamelCase	,			_UpperCamelCase	,			repo_type='''dataset''' )	,			'''r''' ) )
						UpperCAmelCase__		=    {int(_UpperCamelCase ): v for k, v in idalabel.items()}
						UpperCAmelCase__		=    idalabel
						UpperCAmelCase__		=    {v: k for k, v in idalabel.items()}
						return config






def 					a__ (       ):





						"""simple docstring"""

						UpperCAmelCase__		=    hf_hub_download(
						    repo_id='''hf-internal-testing/spaghetti-video'''	,			filename='''eating_spaghetti_32_frames.npy'''	,			repo_type='''dataset''' )
						UpperCAmelCase__		=    np.load(_UpperCamelCase )
						return list(_UpperCamelCase )






def 					a__ (       _UpperCamelCase ):





						"""simple docstring"""

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

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

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






def 					a__ (       _UpperCamelCase ):





						"""simple docstring"""

						return current.flatten()






def 					a__ (       _UpperCamelCase ):





						"""simple docstring"""

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






def 					a__ (       _UpperCamelCase ):





						"""simple docstring"""

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






def 					a__ (       _UpperCamelCase	,			_UpperCamelCase ):





						"""simple docstring"""

						UpperCAmelCase__		=    state_dict['''optimizer''']['''target''']['''Transformer'''][F"""encoderblock_{i}"""]

						UpperCAmelCase__		=    F"""encoder.layer.{i}."""
						UpperCAmelCase__		=    {
						    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__ (       _UpperCamelCase ):





						"""simple docstring"""

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






def 					a__ (       _UpperCamelCase	,			_UpperCamelCase=False ):





						"""simple docstring"""

						UpperCAmelCase__		=    get_n_layers(_UpperCamelCase )

						UpperCAmelCase__		=    OrderedDict()

						UpperCAmelCase__		=    state_dict['''optimizer''']['''target''']['''Transformer''']['''encoder_norm''']['''bias''']
						UpperCAmelCase__		=    state_dict['''optimizer''']['''target''']['''Transformer''']['''encoder_norm''']['''scale''']

						UpperCAmelCase__		=    np.transpose(
						    state_dict['''optimizer''']['''target''']['''embedding''']['''kernel''']	,			(4, 3, 0, 1, 2) )
						UpperCAmelCase__		=    state_dict['''optimizer''']['''target''']['''embedding''']['''bias''']

						UpperCAmelCase__		=    state_dict['''optimizer''']['''target''']['''cls''']
						UpperCAmelCase__		=    state_dict['''optimizer''']['''target''']['''Transformer''']['''posembed_input'''][
						    '''pos_embedding'''
						]

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

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

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






def 					a__ (       ):





						"""simple docstring"""

						UpperCAmelCase__		=    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__ (       _UpperCamelCase ):





						"""simple docstring"""

						UpperCAmelCase__		=    '''checkpoint'''

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

						UpperCAmelCase__		=    restore_checkpoint(_UpperCamelCase	,			_UpperCamelCase )
						UpperCAmelCase__		=    transform_state(_UpperCamelCase	,			classification_head=_UpperCamelCase )

						UpperCAmelCase__		=    get_vivit_config()

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

						UpperCAmelCase__		=    VivitForVideoClassification(_UpperCamelCase )
						model.load_state_dict(_UpperCamelCase )
						model.eval()

						UpperCAmelCase__		=    get_processor()

						UpperCAmelCase__		=    prepare_video()
						UpperCAmelCase__		=    extractor(_UpperCamelCase	,			return_tensors='''pt''' )

						UpperCAmelCase__		=    model(**_UpperCamelCase )

						UpperCAmelCase__		=    torch.Size([1, 400] )
						UpperCAmelCase__		=    torch.tensor([-1.0543, 2.0764, -0.2104, 0.4439, -0.9658] )

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

						model.save_pretrained(_UpperCamelCase )
						extractor.save_pretrained(_UpperCamelCase )


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

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

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