
"""simple docstring"""


import argparse
import pathlib
from pathlib import Path
from tempfile import TemporaryDirectory

import esm as esm_module
import torch
from esm.esmfold.va.misc import batch_encode_sequences as esmfold_encode_sequences
from esm.esmfold.va.pretrained import esmfold_va

from transformers.models.esm.configuration_esm import EsmConfig, EsmFoldConfig
from transformers.models.esm.modeling_esm import (
    EsmForMaskedLM,
    EsmForSequenceClassification,
    EsmIntermediate,
    EsmLayer,
    EsmOutput,
    EsmSelfAttention,
    EsmSelfOutput,
)
from transformers.models.esm.modeling_esmfold import EsmForProteinFolding
from transformers.models.esm.tokenization_esm import EsmTokenizer
from transformers.utils import logging


logging.set_verbosity_info()
_lowercase											=			logging.get_logger(__name__)

_lowercase											=			[
    (
        '''protein1''',
        '''MNGTEGPNFYVPFSNATGVVRSPFEYPQYYLAEPWQFSMLAAYMFLLIVLGFPINFLTLYVTVQHKKLRTPLNYILLNLAVADLFMVLGGFTSTLYTSLHGYFVFGPTGCNLEGFFATLGGEIALWSLVVLAIERYVVVCKPMSNFRFGENHAIMGVAFTWVMALACAAPPLAGWSRYIPEGLQCSCGIDYYTLKPEVNNESFVIYMFVVHFTIPMIIIFFCYGQLVFTVKEAAAQQQESATTQKAEKEVTRMVIIMVIAFLICWVPYASVAFYIFTHQGSNFGPIFMTIPAFFAKSAAIYNPVIYIMMNKQFRNCMLTTICCGKNPLGDDEASATVSKTETSQVAPA''',
    ),
    ('''protein2''', '''MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLA'''),
    ('''protein3''', '''MKTVRQERLKSI<mask>RILERSKEPVSGAQLAEELS<mask>SRQVIVQDIAYLRSLGYN<mask>VATPRGYVLAGG'''),
    ('''protein4''', '''MKTVRQERLKSI<mask>RILERSKEPVSGAQLAEELS<mask>SRQVIVQDIAYLRSLGYN<mask>VATPRGYVLA'''),
]

_lowercase											=			{
    '''esm1b_t33_650M_UR50S''': esm_module.pretrained.esm1b_t33_650M_UR50S,
    '''esm1v_t33_650M_UR90S_1''': esm_module.pretrained.esm1v_t33_650M_UR90S_1,
    '''esm1v_t33_650M_UR90S_2''': esm_module.pretrained.esm1v_t33_650M_UR90S_2,
    '''esm1v_t33_650M_UR90S_3''': esm_module.pretrained.esm1v_t33_650M_UR90S_3,
    '''esm1v_t33_650M_UR90S_4''': esm_module.pretrained.esm1v_t33_650M_UR90S_4,
    '''esm1v_t33_650M_UR90S_5''': esm_module.pretrained.esm1v_t33_650M_UR90S_5,
    '''esm2_t48_15B_UR50D''': esm_module.pretrained.esm2_t48_15B_UR50D,
    '''esm2_t36_3B_UR50D''': esm_module.pretrained.esm2_t36_3B_UR50D,
    '''esm2_t33_650M_UR50D''': esm_module.pretrained.esm2_t33_650M_UR50D,
    '''esm2_t30_150M_UR50D''': esm_module.pretrained.esm2_t30_150M_UR50D,
    '''esm2_t12_35M_UR50D''': esm_module.pretrained.esm2_t12_35M_UR50D,
    '''esm2_t6_8M_UR50D''': esm_module.pretrained.esm2_t6_8M_UR50D,
    '''esmfold_v1''': esmfold_va,
}

_lowercase											=			list('''ARNDCQEGHILKMFPSTWYV''')

_lowercase											=			restypes + ['''X''']
_lowercase											=			restypes_with_x + ['''<pad>''', '''<mask>''', '''<cls>''', '''<sep>''', '''<eos>''']






def  _snake_case			(						):
	with TemporaryDirectory() as tempdir:
		A  				=				'\n'.join(snake_case__       )
		A  				=				Path(snake_case__       ) / 'vocab.txt'
		vocab_file.write_text(snake_case__       )
		A  				=				EsmTokenizer(vocab_file=str(snake_case__       )       )
	A  				=				0  # Overlaps with 'A' but that seems to be what they want
	return hf_tokenizer






def  _snake_case			(						snake_case__				:       List[str]		, snake_case__				:       Optional[Any]       ):
	A  				=				our_module.load_state_dict(original_module.state_dict()       )
	if status.missing_keys:
		raise ValueError(F'Missing keys: {status.missing_keys}'       )
	if status.unexpected_keys:
		raise ValueError(F'Unexpected keys: {status.unexpected_keys}'       )






def  _snake_case			(						snake_case__				:       str		, snake_case__				:       str		, snake_case__				:       bool		, snake_case__				:       str		, snake_case__				:       str       ):
	if model.startswith('esmfold'       ):
		A  				=				MODEL_MAPPING[model]()
	else:
		A				, A  				=				MODEL_MAPPING[model]()
	esm.eval()  # disable dropout

	if model.startswith('esmfold'       ):
		A  				=				esm.esm.embed_dim
		A  				=				esm.esm.num_layers
		A  				=				esm.esm.attention_heads
		A  				=				4 * embed_dim
		A  				=				esm.esm.token_dropout
		A  				=				False  # This code path does not exist in ESM-2
		A  				=				'rotary'
		A  				=				True
		A  				=				EsmFoldConfig()
		for key, val in esm.cfg.items():
			if hasattr(snake_case__		, snake_case__       ) and key != "trunk":
				setattr(snake_case__		, snake_case__		, snake_case__       )
		for key, val in esm.cfg.trunk.items():
			if hasattr(esmfold_config.trunk		, snake_case__       ) and key != "structure_module":
				setattr(esmfold_config.trunk		, snake_case__		, snake_case__       )
		for key, val in esm.cfg.trunk.structure_module.items():
			if hasattr(esmfold_config.trunk.structure_module		, snake_case__       ):
				setattr(esmfold_config.trunk.structure_module		, snake_case__		, snake_case__       )
	elif hasattr(snake_case__		, 'args'       ):
		# Indicates an ESM-1b or ESM-1v model
		A  				=				esm.args.embed_dim
		A  				=				esm.args.layers
		A  				=				esm.args.attention_heads
		A  				=				esm.args.ffn_embed_dim
		A  				=				esm.args.token_dropout
		A  				=				True if esm.emb_layer_norm_before else False
		A  				=				'absolute'
		A  				=				False
		A  				=				None
	else:
		# Indicates an ESM-2 model
		A  				=				esm.embed_dim
		A  				=				esm.num_layers
		A  				=				esm.attention_heads
		A  				=				4 * embed_dim  # This is hardcoded in ESM-2
		A  				=				esm.token_dropout
		A  				=				False  # This code path does not exist in ESM-2
		A  				=				'rotary'
		A  				=				False
		A  				=				None

	if is_folding_model:
		A  				=				esm.esm.alphabet
	A  				=				tuple(alphabet.all_toks       )
	A  				=				alphabet.mask_idx
	A  				=				alphabet.padding_idx

	if is_folding_model:
		A  				=				esm.esm
	else:
		A  				=				esm

	A  				=				EsmConfig(
	    vocab_size=original_esm_model.embed_tokens.num_embeddings		, mask_token_id=snake_case__		, hidden_size=snake_case__		, num_hidden_layers=snake_case__		, num_attention_heads=snake_case__		, intermediate_size=snake_case__		, max_position_embeddings=1026		, layer_norm_eps=1e-5		, attention_probs_dropout_prob=0.0		, hidden_dropout_prob=0.0		, pad_token_id=snake_case__		, emb_layer_norm_before=snake_case__		, token_dropout=snake_case__		, position_embedding_type=snake_case__		, is_folding_model=snake_case__		, esmfold_config=snake_case__		, vocab_list=snake_case__		, )
	if classification_head:
		A  				=				esm.classification_heads['mnli'].out_proj.weight.shape[0]
	print('Our ESM config:'		, snake_case__       )

	if model.startswith('esmfold'       ):
		A  				=				EsmForProteinFolding
	elif classification_head:
		A  				=				EsmForSequenceClassification
	else:
		A  				=				EsmForMaskedLM
	A  				=				model_class(snake_case__       )
	model.eval()

	# Now let's copy all the weights.
	# Embeddings
	A  				=				original_esm_model.embed_tokens.weight
	if position_embedding_type == "absolute":
		A  				=				original_esm_model.embed_positions.weight

	if config.emb_layer_norm_before:
		A  				=				original_esm_model.emb_layer_norm_before.weight
		A  				=				original_esm_model.emb_layer_norm_before.bias

	A  				=				original_esm_model.emb_layer_norm_after.weight
	A  				=				original_esm_model.emb_layer_norm_after.bias

	for i in range(config.num_hidden_layers       ):
		# Encoder: start of layer
		A  				=				model.esm.encoder.layer[i]
		# esm_layer: TransformerSentenceEncoderLayer = original_esm_model.layers[i]
		A  				=				original_esm_model.layers[i]

		# self attention
		A  				=				layer.attention.self
		assert (
		    esm_layer.self_attn.k_proj.weight.data.shape
		    == esm_layer.self_attn.q_proj.weight.data.shape
		    == esm_layer.self_attn.v_proj.weight.data.shape
		    == torch.Size((config.hidden_size, config.hidden_size)       )
		)

		A  				=				esm_layer.self_attn.q_proj.weight
		A  				=				esm_layer.self_attn.q_proj.bias
		A  				=				esm_layer.self_attn.k_proj.weight
		A  				=				esm_layer.self_attn.k_proj.bias
		A  				=				esm_layer.self_attn.v_proj.weight
		A  				=				esm_layer.self_attn.v_proj.bias

		if getattr(esm_layer.self_attn		, 'rot_emb'		, snake_case__       ) is not None:
			# Matt: Although inv_freq is not a trainable weight, it is computed at model init and cached.
			# During the training of ESM-2 the model was converted to float16 precision, which also converts
			# the inv_freq tensor, and the loss of precision remains even if the model is loaded later as float32.
			# If we recompute inv_freq without this loss of precision then we will get subtly different rotary
			# embeddings, which are enough to cause significant discrepancies in model outputs. To avoid this,
			# we make sure the new model copies the data from the old inv_freq.
			A  				=				esm_layer.self_attn.rot_emb.inv_freq

		# LayerNorm changes for pre-activation
		A  				=				esm_layer.self_attn_layer_norm.weight
		A  				=				esm_layer.self_attn_layer_norm.bias
		A  				=				esm_layer.final_layer_norm.weight
		A  				=				esm_layer.final_layer_norm.bias

		# self-attention output
		A  				=				layer.attention.output
		assert self_output.dense.weight.shape == esm_layer.self_attn.out_proj.weight.shape
		A  				=				esm_layer.self_attn.out_proj.weight
		A  				=				esm_layer.self_attn.out_proj.bias

		# intermediate
		A  				=				layer.intermediate
		assert intermediate.dense.weight.shape == esm_layer.fca.weight.shape
		A  				=				esm_layer.fca.weight
		A  				=				esm_layer.fca.bias

		# output
		A  				=				layer.output
		assert bert_output.dense.weight.shape == esm_layer.fca.weight.shape
		A  				=				esm_layer.fca.weight
		A  				=				esm_layer.fca.bias
		# end of layer

	if is_folding_model:
		A  				=				esm.esm_s_combine.data
		A  				=				esm.afa_to_esm.data
		transfer_and_check_weights(esm.embedding		, model.embedding       )
		transfer_and_check_weights(esm.esm_s_mlp		, model.esm_s_mlp       )
		transfer_and_check_weights(esm.trunk		, model.trunk       )
		transfer_and_check_weights(esm.distogram_head		, model.distogram_head       )
		transfer_and_check_weights(esm.ptm_head		, model.ptm_head       )
		transfer_and_check_weights(esm.lm_head		, model.lm_head       )
		transfer_and_check_weights(esm.lddt_head		, model.lddt_head       )

	elif classification_head:
		A  				=				esm.esm.classification_heads['mnli'].dense.weight
		A  				=				esm.classification_heads['mnli'].dense.bias
		A  				=				esm.classification_heads['mnli'].out_proj.weight
		A  				=				esm.classification_heads['mnli'].out_proj.bias
	else:
		# LM Head
		A  				=				esm.lm_head.dense.weight
		A  				=				esm.lm_head.dense.bias
		A  				=				esm.lm_head.layer_norm.weight
		A  				=				esm.lm_head.layer_norm.bias
		A  				=				esm.lm_head.weight
		A  				=				esm.lm_head.bias

	# Contact prediction head
	transfer_and_check_weights(esm.contact_head		, model.esm.contact_head       )

	# Prepare data (first 2 sequences from ESMStructuralSplitDataset superfamily / 4)
	if is_folding_model:
		# Folding models aren't trained on masked inputs and don't like mask tokens.
		A  				=				SAMPLE_DATA[:2]
	else:
		A  				=				SAMPLE_DATA

	if is_folding_model:
		A  				=				get_esmfold_tokenizer()
		A  				=				hf_tokenizer(
		    [row[1] for row in sample_data]		, return_tensors='pt'		, padding=snake_case__		, add_special_tokens=snake_case__       )
		A				, A				, A				, A				, A  				=				esmfold_encode_sequences([row[1] for row in sample_data]       )
		A  				=				torch.all(hf_tokens['input_ids'] == esmfold_aas       ) and torch.all(
		    hf_tokens['attention_mask'] == esmfold_mask       )
	else:
		# Let's check that we get the same results.
		A  				=				alphabet.get_batch_converter()
		A				, A				, A  				=				batch_converter(snake_case__       )
		# Prepare tokenizer and make sure it matches
		with TemporaryDirectory() as tempdir:
			A  				=				'\n'.join(alphabet.all_toks       )
			A  				=				Path(snake_case__       ) / 'vocab.txt'
			vocab_file.write_text(snake_case__       )
			A  				=				EsmTokenizer(vocab_file=str(snake_case__       )       )

		A  				=				hf_tokenizer([row[1] for row in sample_data]		, return_tensors='pt'		, padding=snake_case__       )
		A  				=				torch.all(hf_tokens['input_ids'] == batch_tokens       )

	print('Do both models tokenizers output the same tokens?'		, '🔥' if success else '💩'       )
	if not success:
		raise Exception('Tokenization does not match!'       )

	with torch.no_grad():
		if is_folding_model:
			# Let's test the model in parts
			# ESMFold always converts the ESM stem to float16, which requires float16 ops
			# that don't exist on CPU. Therefore, to test it we need to run it on GPU. However,
			# ESMFold is what we in the community call a "big boy" and so we desperately avoid putting both the
			# original and the converted model on the GPU at the same time.
			A  				=				esm.cuda().infer([row[1] for row in sample_data]       )
			A  				=				model.cuda()(
			    input_ids=hf_tokens['input_ids'].cuda()		, attention_mask=hf_tokens['attention_mask'].cuda()       )
		else:
			A  				=				model(**snake_case__		, output_hidden_states=snake_case__       )
			A  				=				our_output['logits']
			if classification_head:
				A  				=				esm.model.classification_heads['mnli'](esm.extract_features(snake_case__       )       )
			else:
				A  				=				esm(hf_tokens['input_ids']		, repr_layers=list(range(999       )       )       )
				A  				=				their_output['logits']

		if is_folding_model:
			A  				=				torch.max(torch.abs(our_output['positions'] - their_output['positions']       )       ).item()
			A  				=				torch.allclose(our_output['positions']		, their_output['positions']		, atol=1e-5       )
		else:
			A  				=				torch.max(torch.abs(our_output - their_output       )       ).item()
			A  				=				torch.allclose(snake_case__		, snake_case__		, atol=1e-5       )

		print(F'max_absolute_diff = {max_absolute_diff}'       )  # ~ 1e-5
		print('Do both models output the same tensors?'		, '🔥' if success else '💩'       )

		if not success:
			raise Exception('Something went wRoNg'       )

		if not is_folding_model:
			# Let's check contact prediction too
			A  				=				model.predict_contacts(hf_tokens['input_ids']		, hf_tokens['attention_mask']       )
			A  				=				esm.predict_contacts(hf_tokens['input_ids']       )
			A  				=				torch.max(torch.abs(our_output - their_output       )       ).item()
			A  				=				torch.allclose(snake_case__		, snake_case__		, atol=1e-5       )

			print('Contact prediction testing:'       )
			print(F'max_absolute_diff = {max_absolute_diff}'       )  # ~ 1e-5
			print('Do both models output the same tensors?'		, '🔥' if success else '💩'       )

			if not success:
				raise Exception('Something went wRoNg'       )

		pathlib.Path(snake_case__       ).mkdir(parents=snake_case__		, exist_ok=snake_case__       )
		print(F'Saving model to {pytorch_dump_folder_path}'       )
		model.save_pretrained(snake_case__       )

		del esm  # Free up some memory before continuing

	print(F'Saving tokenizer to {pytorch_dump_folder_path}'       )
	hf_tokenizer.save_pretrained(snake_case__       )

	if push_to_repo:
		model.push_to_hub(repo_id=snake_case__		, use_auth_token=snake_case__       )
		hf_tokenizer.push_to_hub(repo_id=snake_case__		, use_auth_token=snake_case__       )


if __name__ == "__main__":
		_lowercase											=			argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    '''--pytorch_dump_folder_path''', type=str, required=True, help='''Path to the output PyTorch model.'''
		)
		parser.add_argument(
		    '''--classification_head''', action='''store_true''', help='''Whether to convert a final classification head.'''
		)
		parser.add_argument('''--model''', default=None, type=str, required=True, help='''Name of model to convert.''')
		parser.add_argument('''--push_to_repo''', type=str, help='''Repo to upload to (including username!).''')
		parser.add_argument('''--auth_token''', type=str, help='''HuggingFace auth token.''')
		_lowercase											=			parser.parse_args()
		convert_esm_checkpoint_to_pytorch(
		    args.model, args.pytorch_dump_folder_path, args.classification_head, args.push_to_repo, args.auth_token
		)

"""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  _snake_case			(						snake_case__				:       Tuple       ):
	A  				=				'https://storage.googleapis.com/scenic-bucket/vivit/kinetics_400/vivit_base_16x2_unfactorized/checkpoint'

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






def  _snake_case			(						):
	A  				=				VivitConfig()

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

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






def  _snake_case			(						):
	A  				=				hf_hub_download(
	    repo_id='hf-internal-testing/spaghetti-video'		, filename='eating_spaghetti_32_frames.npy'		, repo_type='dataset'       )
	A  				=				np.load(snake_case__       )
	return list(snake_case__       )






def  _snake_case			(						snake_case__				:       np.ndarray       ):
	if np.ndim(snake_case__       ) == 2:
		return transform_attention_bias(snake_case__       )

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

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






def  _snake_case			(						snake_case__				:       np.ndarray       ):
	return current.flatten()






def  _snake_case			(						snake_case__				:       np.ndarray       ):
	return np.reshape(snake_case__		, (current.shape[0], current.shape[1] * current.shape[2])       ).T






def  _snake_case			(						snake_case__				:       np.ndarray       ):
	return np.reshape(snake_case__		, (current.shape[0] * current.shape[1], current.shape[2])       ).T






def  _snake_case			(						snake_case__				:       str		, snake_case__				:       List[str]       ):
	A  				=				state_dict['optimizer']['target']['Transformer'][F'encoderblock_{i}']

	A  				=				F'encoder.layer.{i}.'
	A  				=				{
	    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  _snake_case			(						snake_case__				:       Dict       ):
	return sum([1 if 'encoderblock_' in k else 0 for k in state_dict['optimizer']['target']['Transformer'].keys()]       )






def  _snake_case			(						snake_case__				:       Union[str, Any]		, snake_case__				:       Optional[int]=False       ):
	A  				=				get_n_layers(snake_case__       )

	A  				=				OrderedDict()

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

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

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

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

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

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






def  _snake_case			(						):
	A  				=				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  _snake_case			(						snake_case__				:       str       ):
	A  				=				'checkpoint'

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

	A  				=				restore_checkpoint(snake_case__		, snake_case__       )
	A  				=				transform_state(snake_case__		, classification_head=snake_case__       )

	A  				=				get_vivit_config()

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

	A  				=				VivitForVideoClassification(snake_case__       )
	model.load_state_dict(snake_case__       )
	model.eval()

	A  				=				get_processor()

	A  				=				prepare_video()
	A  				=				extractor(snake_case__		, return_tensors='pt'       )

	A  				=				model(**snake_case__       )

	A  				=				torch.Size([1, 400]       )
	A  				=				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]		, snake_case__		, atol=1e-4       ), outputs.logits[0, :5]

	model.save_pretrained(snake_case__       )
	extractor.save_pretrained(snake_case__       )


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

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

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