


'''simple docstring'''




import argparse
import json

import gdown
import numpy as np
import torch
from huggingface_hub import hf_hub_download

from transformers import (
    VideoMAEConfig,
    VideoMAEForPreTraining,
    VideoMAEForVideoClassification,
    VideoMAEImageProcessor,
)

def 			__UpperCamelCase    (    _A    :				Dict				)    ->      Any:




	"""simple docstring"""





	lowerCAmelCase					:    List[str]      							= VideoMAEConfig()

	set_architecture_configs(lowerCAmelCase__					,   lowerCAmelCase__				)

	if "finetuned" not in model_name:
		lowerCAmelCase					:    Optional[Any]      							= False

	if "finetuned" in model_name:
		lowerCAmelCase					:    Optional[Any]      							= 'huggingface/label-files'
		if "kinetics" in model_name:
			lowerCAmelCase					:    Optional[Any]      							= 4_00
			lowerCAmelCase					:    List[Any]      							= 'kinetics400-id2label.json'
		elif "ssv2" in model_name:
			lowerCAmelCase					:    Union[str, Any]      							= 1_74
			lowerCAmelCase					:    Dict      							= 'something-something-v2-id2label.json'
		else:
			raise ValueError('Model name should either contain \'kinetics\' or \'ssv2\' in case it\'s fine-tuned.'				)
		lowerCAmelCase					:    Optional[Any]      							= json.load(open(hf_hub_download(lowerCAmelCase__					,   lowerCAmelCase__					,   repo_type='dataset'				)					,   'r'				)				)
		lowerCAmelCase					:    str      							= {int(lowerCAmelCase__				): v for k, v in idalabel.items()}
		lowerCAmelCase					:    Union[str, Any]      							= idalabel
		lowerCAmelCase					:    Union[str, Any]      							= {v: k for k, v in idalabel.items()}

	return config

def 			__UpperCamelCase    (    _A    :				Tuple					,   _A    :				Any				)    ->      List[str]:




	"""simple docstring"""





	if "small" in model_name:
		lowerCAmelCase					:    Dict      							= 3_84
		lowerCAmelCase					:    List[str]      							= 15_36
		lowerCAmelCase					:    Tuple      							= 12
		lowerCAmelCase					:    int      							= 16
		lowerCAmelCase					:    Optional[int]      							= 12
		lowerCAmelCase					:    Optional[Any]      							= 3
		lowerCAmelCase					:    List[Any]      							= 1_92
		lowerCAmelCase					:    Any      							= 7_68
	elif "large" in model_name:
		lowerCAmelCase					:    List[str]      							= 10_24
		lowerCAmelCase					:    List[str]      							= 40_96
		lowerCAmelCase					:    List[Any]      							= 24
		lowerCAmelCase					:    Any      							= 16
		lowerCAmelCase					:    Optional[Any]      							= 12
		lowerCAmelCase					:    List[Any]      							= 8
		lowerCAmelCase					:    Union[str, Any]      							= 5_12
		lowerCAmelCase					:    Any      							= 20_48
	elif "huge" in model_name:
		lowerCAmelCase					:    Optional[int]      							= 12_80
		lowerCAmelCase					:    Tuple      							= 51_20
		lowerCAmelCase					:    List[Any]      							= 32
		lowerCAmelCase					:    Dict      							= 16
		lowerCAmelCase					:    Union[str, Any]      							= 12
		lowerCAmelCase					:    Any      							= 8
		lowerCAmelCase					:    Optional[Any]      							= 6_40
		lowerCAmelCase					:    str      							= 25_60
	elif "base" not in model_name:
		raise ValueError('Model name should include either \"small\", \"base\", \"large\", or \"huge\"'				)

def 			__UpperCamelCase    (    _A    :				Optional[Any]				)    ->      Optional[Any]:




	"""simple docstring"""





	if "encoder." in name:
		lowerCAmelCase					:    Tuple      							= name.replace('encoder.'					,   ''				)
	if "cls_token" in name:
		lowerCAmelCase					:    Union[str, Any]      							= name.replace('cls_token'					,   'videomae.embeddings.cls_token'				)
	if "decoder_pos_embed" in name:
		lowerCAmelCase					:    str      							= name.replace('decoder_pos_embed'					,   'decoder.decoder_pos_embed'				)
	if "pos_embed" in name and "decoder" not in name:
		lowerCAmelCase					:    Tuple      							= name.replace('pos_embed'					,   'videomae.embeddings.position_embeddings'				)
	if "patch_embed.proj" in name:
		lowerCAmelCase					:    str      							= name.replace('patch_embed.proj'					,   'videomae.embeddings.patch_embeddings.projection'				)
	if "patch_embed.norm" in name:
		lowerCAmelCase					:    Tuple      							= name.replace('patch_embed.norm'					,   'videomae.embeddings.norm'				)
	if "decoder.blocks" in name:
		lowerCAmelCase					:    int      							= name.replace('decoder.blocks'					,   'decoder.decoder_layers'				)
	if "blocks" in name:
		lowerCAmelCase					:    List[Any]      							= name.replace('blocks'					,   'videomae.encoder.layer'				)
	if "attn.proj" in name:
		lowerCAmelCase					:    Any      							= name.replace('attn.proj'					,   'attention.output.dense'				)
	if "attn" in name and "bias" not in name:
		lowerCAmelCase					:    List[str]      							= name.replace('attn'					,   'attention.self'				)
	if "attn" in name:
		lowerCAmelCase					:    List[str]      							= name.replace('attn'					,   'attention.attention'				)
	if "norm1" in name:
		lowerCAmelCase					:    Optional[int]      							= name.replace('norm1'					,   'layernorm_before'				)
	if "norm2" in name:
		lowerCAmelCase					:    str      							= name.replace('norm2'					,   'layernorm_after'				)
	if "mlp.fc1" in name:
		lowerCAmelCase					:    Optional[int]      							= name.replace('mlp.fc1'					,   'intermediate.dense'				)
	if "mlp.fc2" in name:
		lowerCAmelCase					:    Union[str, Any]      							= name.replace('mlp.fc2'					,   'output.dense'				)
	if "decoder_embed" in name:
		lowerCAmelCase					:    int      							= name.replace('decoder_embed'					,   'decoder.decoder_embed'				)
	if "decoder_norm" in name:
		lowerCAmelCase					:    Optional[int]      							= name.replace('decoder_norm'					,   'decoder.decoder_norm'				)
	if "decoder_pred" in name:
		lowerCAmelCase					:    Tuple      							= name.replace('decoder_pred'					,   'decoder.decoder_pred'				)
	if "norm.weight" in name and "decoder" not in name and "fc" not in name:
		lowerCAmelCase					:    Optional[int]      							= name.replace('norm.weight'					,   'videomae.layernorm.weight'				)
	if "norm.bias" in name and "decoder" not in name and "fc" not in name:
		lowerCAmelCase					:    Optional[int]      							= name.replace('norm.bias'					,   'videomae.layernorm.bias'				)
	if "head" in name and "decoder" not in name:
		lowerCAmelCase					:    str      							= name.replace('head'					,   'classifier'				)

	return name

def 			__UpperCamelCase    (    _A    :				int					,   _A    :				Optional[Any]				)    ->      Any:




	"""simple docstring"""





	for key in orig_state_dict.copy().keys():
		lowerCAmelCase					:    List[str]      							= orig_state_dict.pop(lowerCAmelCase__				)

		if key.startswith('encoder.'				):
			lowerCAmelCase					:    Optional[Any]      							= key.replace('encoder.'					,   ''				)

		if "qkv" in key:
			lowerCAmelCase					:    str      							= key.split('.'				)
			if key.startswith('decoder.blocks'				):
				lowerCAmelCase					:    List[Any]      							= config.decoder_hidden_size
				lowerCAmelCase					:    Dict      							= int(key_split[2]				)
				lowerCAmelCase					:    List[Any]      							= 'decoder.decoder_layers.'
				if "weight" in key:
					lowerCAmelCase					:    List[str]      							= val[:dim, :]
					lowerCAmelCase					:    Optional[Any]      							= val[dim : dim * 2, :]
					lowerCAmelCase					:    Dict      							= val[-dim:, :]
			else:
				lowerCAmelCase					:    List[str]      							= config.hidden_size
				lowerCAmelCase					:    List[str]      							= int(key_split[1]				)
				lowerCAmelCase					:    int      							= 'videomae.encoder.layer.'
				if "weight" in key:
					lowerCAmelCase					:    Union[str, Any]      							= val[:dim, :]
					lowerCAmelCase					:    int      							= val[dim : dim * 2, :]
					lowerCAmelCase					:    Optional[Any]      							= val[-dim:, :]
		else:
			lowerCAmelCase					:    Optional[int]      							= val

	return orig_state_dict

def 			__UpperCamelCase    (    )    ->      Any:




	"""simple docstring"""





	lowerCAmelCase					:    Union[str, Any]      							= hf_hub_download(
	    repo_id='hf-internal-testing/spaghetti-video'					,   filename='eating_spaghetti.npy'					,   repo_type='dataset'				)
	lowerCAmelCase					:    Tuple      							= np.load(lowerCAmelCase__				)
	return list(lowerCAmelCase__				)

def 			__UpperCamelCase    (    _A    :				Tuple					,   _A    :				Any					,   _A    :				Dict					,   _A    :				Tuple				)    ->      Any:




	"""simple docstring"""





	lowerCAmelCase					:    Any      							= get_videomae_config(lowerCAmelCase__				)

	if "finetuned" in model_name:
		lowerCAmelCase					:    Tuple      							= VideoMAEForVideoClassification(lowerCAmelCase__				)
	else:
		lowerCAmelCase					:    Any      							= VideoMAEForPreTraining(lowerCAmelCase__				)

	# download original checkpoint, hosted on Google Drive
	lowerCAmelCase					:    Optional[Any]      							= 'pytorch_model.bin'
	gdown.cached_download(lowerCAmelCase__					,   lowerCAmelCase__					,   quiet=lowerCAmelCase__				)
	lowerCAmelCase					:    str      							= torch.load(lowerCAmelCase__					,   map_location='cpu'				)
	if "model" in files:
		lowerCAmelCase					:    Tuple      							= files['model']
	else:
		lowerCAmelCase					:    Union[str, Any]      							= files['module']
	lowerCAmelCase					:    Tuple      							= convert_state_dict(lowerCAmelCase__					,   lowerCAmelCase__				)

	model.load_state_dict(lowerCAmelCase__				)
	model.eval()

	# verify model on basic input
	lowerCAmelCase					:    List[Any]      							= VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]					,   image_std=[0.5, 0.5, 0.5]				)
	lowerCAmelCase					:    Union[str, Any]      							= prepare_video()
	lowerCAmelCase					:    int      							= image_processor(lowerCAmelCase__					,   return_tensors='pt'				)

	if "finetuned" not in model_name:
		lowerCAmelCase					:    int      							= hf_hub_download(repo_id='hf-internal-testing/bool-masked-pos'					,   filename='bool_masked_pos.pt'				)
		lowerCAmelCase					:    Union[str, Any]      							= torch.load(lowerCAmelCase__				)

	lowerCAmelCase					:    Optional[int]      							= model(**lowerCAmelCase__				)
	lowerCAmelCase					:    int      							= outputs.logits

	lowerCAmelCase					:    Tuple      							= [
	    'videomae-small-finetuned-kinetics',
	    'videomae-small-finetuned-ssv2',
	    # Kinetics-400 checkpoints (short = pretrained only for 800 epochs instead of 1600)
	    'videomae-base-short',
	    'videomae-base-short-finetuned-kinetics',
	    'videomae-base',
	    'videomae-base-finetuned-kinetics',
	    'videomae-large',
	    'videomae-large-finetuned-kinetics',
	    'videomae-huge-finetuned-kinetics',
	    # Something-Something-v2 checkpoints (short = pretrained only for 800 epochs instead of 2400)
	    'videomae-base-short-ssv2',
	    'videomae-base-short-finetuned-ssv2',
	    'videomae-base-ssv2',
	    'videomae-base-finetuned-ssv2',
	]

	# NOTE: logits were tested with image_mean and image_std equal to [0.5, 0.5, 0.5] and [0.5, 0.5, 0.5]
	if model_name == "videomae-small-finetuned-kinetics":
		lowerCAmelCase					:    Dict      							= torch.Size([1, 4_00]				)
		lowerCAmelCase					:    Tuple      							= torch.tensor([-0.92_91, -0.40_61, -0.93_07]				)
	elif model_name == "videomae-small-finetuned-ssv2":
		lowerCAmelCase					:    Optional[int]      							= torch.Size([1, 1_74]				)
		lowerCAmelCase					:    Union[str, Any]      							= torch.tensor([0.26_71, -0.46_89, -0.82_35]				)
	elif model_name == "videomae-base":
		lowerCAmelCase					:    int      							= torch.Size([1, 14_08, 15_36]				)
		lowerCAmelCase					:    Any      							= torch.tensor([[0.77_39, 0.79_68, 0.70_89], [0.67_01, 0.74_87, 0.62_09], [0.42_87, 0.51_58, 0.47_73]]				)
	elif model_name == "videomae-base-short":
		lowerCAmelCase					:    Dict      							= torch.Size([1, 14_08, 15_36]				)
		lowerCAmelCase					:    Union[str, Any]      							= torch.tensor([[0.79_94, 0.96_12, 0.85_08], [0.74_01, 0.89_58, 0.83_02], [0.58_62, 0.74_68, 0.73_25]]				)
		# we verified the loss both for normalized and unnormalized targets for this one
		lowerCAmelCase					:    Optional[Any]      							= torch.tensor([0.51_42]				) if config.norm_pix_loss else torch.tensor([0.64_69]				)
	elif model_name == "videomae-large":
		lowerCAmelCase					:    str      							= torch.Size([1, 14_08, 15_36]				)
		lowerCAmelCase					:    List[str]      							= torch.tensor([[0.71_49, 0.79_97, 0.69_66], [0.67_68, 0.78_69, 0.69_48], [0.51_39, 0.62_21, 0.56_05]]				)
	elif model_name == "videomae-large-finetuned-kinetics":
		lowerCAmelCase					:    Optional[int]      							= torch.Size([1, 4_00]				)
		lowerCAmelCase					:    Any      							= torch.tensor([0.07_71, 0.00_11, -0.36_25]				)
	elif model_name == "videomae-huge-finetuned-kinetics":
		lowerCAmelCase					:    int      							= torch.Size([1, 4_00]				)
		lowerCAmelCase					:    str      							= torch.tensor([0.24_33, 0.16_32, -0.48_94]				)
	elif model_name == "videomae-base-short-finetuned-kinetics":
		lowerCAmelCase					:    Any      							= torch.Size([1, 4_00]				)
		lowerCAmelCase					:    str      							= torch.tensor([0.65_88, 0.09_90, -0.24_93]				)
	elif model_name == "videomae-base-finetuned-kinetics":
		lowerCAmelCase					:    List[Any]      							= torch.Size([1, 4_00]				)
		lowerCAmelCase					:    Union[str, Any]      							= torch.tensor([0.36_69, -0.06_88, -0.24_21]				)
	elif model_name == "videomae-base-short-ssv2":
		lowerCAmelCase					:    Dict      							= torch.Size([1, 14_08, 15_36]				)
		lowerCAmelCase					:    str      							= torch.tensor([[0.47_12, 0.52_96, 0.57_86], [0.22_78, 0.27_29, 0.40_26], [0.03_52, 0.07_30, 0.25_06]]				)
	elif model_name == "videomae-base-short-finetuned-ssv2":
		lowerCAmelCase					:    Tuple      							= torch.Size([1, 1_74]				)
		lowerCAmelCase					:    int      							= torch.tensor([-0.05_37, -0.15_39, -0.32_66]				)
	elif model_name == "videomae-base-ssv2":
		lowerCAmelCase					:    str      							= torch.Size([1, 14_08, 15_36]				)
		lowerCAmelCase					:    Union[str, Any]      							= torch.tensor([[0.81_31, 0.87_27, 0.85_46], [0.73_66, 0.93_77, 0.88_70], [0.59_35, 0.88_74, 0.85_64]]				)
	elif model_name == "videomae-base-finetuned-ssv2":
		lowerCAmelCase					:    Any      							= torch.Size([1, 1_74]				)
		lowerCAmelCase					:    str      							= torch.tensor([0.19_61, -0.83_37, -0.63_89]				)
	else:
		raise ValueError(F"Model name not supported. Should be one of {model_names}"				)

	# verify logits
	assert logits.shape == expected_shape
	if "finetuned" in model_name:
		assert torch.allclose(logits[0, :3]					,   lowerCAmelCase__					,   atol=1e-4				)
	else:
		print('Logits:'					,   logits[0, :3, :3]				)
		assert torch.allclose(logits[0, :3, :3]					,   lowerCAmelCase__					,   atol=1e-4				)
	print('Logits ok!'				)

	# verify loss, if applicable
	if model_name == "videomae-base-short":
		lowerCAmelCase					:    Tuple      							= outputs.loss
		assert torch.allclose(lowerCAmelCase__					,   lowerCAmelCase__					,   atol=1e-4				)
		print('Loss ok!'				)

	if pytorch_dump_folder_path is not None:
		print(F"Saving model and image processor to {pytorch_dump_folder_path}"				)
		image_processor.save_pretrained(lowerCAmelCase__				)
		model.save_pretrained(lowerCAmelCase__				)

	if push_to_hub:
		print('Pushing to the hub...'				)
		model.push_to_hub(lowerCAmelCase__					,   organization='nielsr'				)


if __name__ == "__main__":
				_lowerCAmelCase  :		Union[str, Any]		     =	argparse.ArgumentParser()
				# Required parameters
				parser.add_argument(
				    '--checkpoint_url',
				    default='https://drive.google.com/u/1/uc?id=1tEhLyskjb755TJ65ptsrafUG2llSwQE1&amp;export=download&amp;confirm=t&amp;uuid=aa3276eb-fb7e-482a-adec-dc7171df14c4',
				    type=str,
				    help=(
				        'URL of the original PyTorch checkpoint (on Google Drive) you\'d like to convert. Should be a direct'
				        ' download link.'
				    ),
				)
				parser.add_argument(
				    '--pytorch_dump_folder_path',
				    default='/Users/nielsrogge/Documents/VideoMAE/Test',
				    type=str,
				    help='Path to the output PyTorch model directory.',
				)
				parser.add_argument('--model_name', default='videomae-base', type=str, help='Name of the model.')
				parser.add_argument(
				    '--push_to_hub', action='store_true', help='Whether or not to push the converted model to the 🤗 hub.'
				)

				_lowerCAmelCase  :		Tuple		     =	parser.parse_args()
				convert_videomae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)









'''simple docstring'''




from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


_lowerCAmelCase  :		Tuple		     =	logging.get_logger(__name__)

_lowerCAmelCase  :		List[str]		     =	{
    'facebook/xmod-base': 'https://huggingface.co/facebook/xmod-base/resolve/main/config.json',
    'facebook/xmod-large-prenorm': 'https://huggingface.co/facebook/xmod-large-prenorm/resolve/main/config.json',
    'facebook/xmod-base-13-125k': 'https://huggingface.co/facebook/xmod-base-13-125k/resolve/main/config.json',
    'facebook/xmod-base-30-125k': 'https://huggingface.co/facebook/xmod-base-30-125k/resolve/main/config.json',
    'facebook/xmod-base-30-195k': 'https://huggingface.co/facebook/xmod-base-30-195k/resolve/main/config.json',
    'facebook/xmod-base-60-125k': 'https://huggingface.co/facebook/xmod-base-60-125k/resolve/main/config.json',
    'facebook/xmod-base-60-265k': 'https://huggingface.co/facebook/xmod-base-60-265k/resolve/main/config.json',
    'facebook/xmod-base-75-125k': 'https://huggingface.co/facebook/xmod-base-75-125k/resolve/main/config.json',
    'facebook/xmod-base-75-269k': 'https://huggingface.co/facebook/xmod-base-75-269k/resolve/main/config.json',
}



class  lowerCAmelCase				(				a							):
	_lowerCamelCase   :					int			      = """xmod"""



	def __init__(					self ,					snake_case__=3_0522 ,					snake_case__=768 ,					snake_case__=12 ,					snake_case__=12 ,					snake_case__=3072 ,					snake_case__="gelu" ,					snake_case__=0.1 ,					snake_case__=0.1 ,					snake_case__=512 ,					snake_case__=2 ,					snake_case__=0.0_2 ,					snake_case__=1e-1_2 ,					snake_case__=1 ,					snake_case__=0 ,					snake_case__=2 ,					snake_case__="absolute" ,					snake_case__=True ,					snake_case__=None ,					snake_case__=False ,					snake_case__=2 ,					snake_case__=False ,					snake_case__=True ,					snake_case__=True ,					snake_case__=("en_XX",) ,					snake_case__=None ,					**snake_case__ ,					):
		super().__init__(pad_token_id=snake_case__ ,					bos_token_id=snake_case__ ,					eos_token_id=snake_case__ ,					**snake_case__    )

		lowerCAmelCase					:    Dict      							= vocab_size
		lowerCAmelCase					:    Optional[Any]      							= hidden_size
		lowerCAmelCase					:    int      							= num_hidden_layers
		lowerCAmelCase					:    List[Any]      							= num_attention_heads
		lowerCAmelCase					:    List[Any]      							= hidden_act
		lowerCAmelCase					:    Optional[int]      							= intermediate_size
		lowerCAmelCase					:    Optional[int]      							= hidden_dropout_prob
		lowerCAmelCase					:    Optional[Any]      							= attention_probs_dropout_prob
		lowerCAmelCase					:    str      							= max_position_embeddings
		lowerCAmelCase					:    int      							= type_vocab_size
		lowerCAmelCase					:    List[Any]      							= initializer_range
		lowerCAmelCase					:    Any      							= layer_norm_eps
		lowerCAmelCase					:    Dict      							= position_embedding_type
		lowerCAmelCase					:    Optional[Any]      							= use_cache
		lowerCAmelCase					:    Union[str, Any]      							= classifier_dropout
		lowerCAmelCase					:    int      							= pre_norm
		lowerCAmelCase					:    Optional[Any]      							= adapter_reduction_factor
		lowerCAmelCase					:    Any      							= adapter_layer_norm
		lowerCAmelCase					:    Dict      							= adapter_reuse_layer_norm
		lowerCAmelCase					:    Any      							= ln_before_adapter
		lowerCAmelCase					:    Optional[Any]      							= list(snake_case__    )
		lowerCAmelCase					:    List[Any]      							= default_language



class  lowerCAmelCase				(				a							):



	@property
	def 				lowercase     (					self    ):
		if self.task == "multiple-choice":
			lowerCAmelCase					:    List[Any]      							= {0: 'batch', 1: 'choice', 2: 'sequence'}
		else:
			lowerCAmelCase					:    Optional[int]      							= {0: 'batch', 1: 'sequence'}
		return OrderedDict(
		    [
		        ('input_ids', dynamic_axis),
		        ('attention_mask', dynamic_axis),
		    ]    )
