
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   __lowerCAmelCase    (				UpperCAmelCase__					:   Union[str, Any]							)	->     Any:
							lowerCamelCase_					=       VideoMAEConfig()

							set_architecture_configs(UpperCAmelCase__      ,					UpperCAmelCase__							)

							if "finetuned" not in model_name:
														lowerCamelCase_					=       False

							if "finetuned" in model_name:
														lowerCamelCase_					=       """huggingface/label-files"""
														if "kinetics" in model_name:
																					lowerCamelCase_					=       4_0_0
																					lowerCamelCase_					=       """kinetics400-id2label.json"""
														elif "ssv2" in model_name:
																					lowerCamelCase_					=       1_7_4
																					lowerCamelCase_					=       """something-something-v2-id2label.json"""
														else:
																					raise ValueError("""Model name should either contain 'kinetics' or 'ssv2' in case it's fine-tuned."""							)
														lowerCamelCase_					=       json.load(open(hf_hub_download(UpperCAmelCase__      ,					UpperCAmelCase__      ,					repo_type="""dataset"""							)      ,					"""r"""							)							)
														lowerCamelCase_					=       {int(UpperCAmelCase__							): v for k, v in idalabel.items()}
														lowerCamelCase_					=       idalabel
														lowerCamelCase_					=       {v: k for k, v in idalabel.items()}

							return config



def   __lowerCAmelCase    (				UpperCAmelCase__					:   int      ,					UpperCAmelCase__					:   Any							)	->     int:
							if "small" in model_name:
														lowerCamelCase_					=       3_8_4
														lowerCamelCase_					=       1_5_3_6
														lowerCamelCase_					=       1_2
														lowerCamelCase_					=       1_6
														lowerCamelCase_					=       1_2
														lowerCamelCase_					=       3
														lowerCamelCase_					=       1_9_2
														lowerCamelCase_					=       7_6_8
							elif "large" in model_name:
														lowerCamelCase_					=       1_0_2_4
														lowerCamelCase_					=       4_0_9_6
														lowerCamelCase_					=       2_4
														lowerCamelCase_					=       1_6
														lowerCamelCase_					=       1_2
														lowerCamelCase_					=       8
														lowerCamelCase_					=       5_1_2
														lowerCamelCase_					=       2_0_4_8
							elif "huge" in model_name:
														lowerCamelCase_					=       1_2_8_0
														lowerCamelCase_					=       5_1_2_0
														lowerCamelCase_					=       3_2
														lowerCamelCase_					=       1_6
														lowerCamelCase_					=       1_2
														lowerCamelCase_					=       8
														lowerCamelCase_					=       6_4_0
														lowerCamelCase_					=       2_5_6_0
							elif "base" not in model_name:
														raise ValueError("""Model name should include either \"small\", \"base\", \"large\", or \"huge\""""							)



def   __lowerCAmelCase    (				UpperCAmelCase__					:   Dict							)	->     Dict:
							if "encoder." in name:
														lowerCamelCase_					=       name.replace("""encoder."""      ,					""""""							)
							if "cls_token" in name:
														lowerCamelCase_					=       name.replace("""cls_token"""      ,					"""videomae.embeddings.cls_token"""							)
							if "decoder_pos_embed" in name:
														lowerCamelCase_					=       name.replace("""decoder_pos_embed"""      ,					"""decoder.decoder_pos_embed"""							)
							if "pos_embed" in name and "decoder" not in name:
														lowerCamelCase_					=       name.replace("""pos_embed"""      ,					"""videomae.embeddings.position_embeddings"""							)
							if "patch_embed.proj" in name:
														lowerCamelCase_					=       name.replace("""patch_embed.proj"""      ,					"""videomae.embeddings.patch_embeddings.projection"""							)
							if "patch_embed.norm" in name:
														lowerCamelCase_					=       name.replace("""patch_embed.norm"""      ,					"""videomae.embeddings.norm"""							)
							if "decoder.blocks" in name:
														lowerCamelCase_					=       name.replace("""decoder.blocks"""      ,					"""decoder.decoder_layers"""							)
							if "blocks" in name:
														lowerCamelCase_					=       name.replace("""blocks"""      ,					"""videomae.encoder.layer"""							)
							if "attn.proj" in name:
														lowerCamelCase_					=       name.replace("""attn.proj"""      ,					"""attention.output.dense"""							)
							if "attn" in name and "bias" not in name:
														lowerCamelCase_					=       name.replace("""attn"""      ,					"""attention.self"""							)
							if "attn" in name:
														lowerCamelCase_					=       name.replace("""attn"""      ,					"""attention.attention"""							)
							if "norm1" in name:
														lowerCamelCase_					=       name.replace("""norm1"""      ,					"""layernorm_before"""							)
							if "norm2" in name:
														lowerCamelCase_					=       name.replace("""norm2"""      ,					"""layernorm_after"""							)
							if "mlp.fc1" in name:
														lowerCamelCase_					=       name.replace("""mlp.fc1"""      ,					"""intermediate.dense"""							)
							if "mlp.fc2" in name:
														lowerCamelCase_					=       name.replace("""mlp.fc2"""      ,					"""output.dense"""							)
							if "decoder_embed" in name:
														lowerCamelCase_					=       name.replace("""decoder_embed"""      ,					"""decoder.decoder_embed"""							)
							if "decoder_norm" in name:
														lowerCamelCase_					=       name.replace("""decoder_norm"""      ,					"""decoder.decoder_norm"""							)
							if "decoder_pred" in name:
														lowerCamelCase_					=       name.replace("""decoder_pred"""      ,					"""decoder.decoder_pred"""							)
							if "norm.weight" in name and "decoder" not in name and "fc" not in name:
														lowerCamelCase_					=       name.replace("""norm.weight"""      ,					"""videomae.layernorm.weight"""							)
							if "norm.bias" in name and "decoder" not in name and "fc" not in name:
														lowerCamelCase_					=       name.replace("""norm.bias"""      ,					"""videomae.layernorm.bias"""							)
							if "head" in name and "decoder" not in name:
														lowerCamelCase_					=       name.replace("""head"""      ,					"""classifier"""							)

							return name



def   __lowerCAmelCase    (				UpperCAmelCase__					:   Optional[int]      ,					UpperCAmelCase__					:   List[Any]							)	->     Union[str, Any]:
							for key in orig_state_dict.copy().keys():
														lowerCamelCase_					=       orig_state_dict.pop(UpperCAmelCase__							)

														if key.startswith("""encoder."""							):
																					lowerCamelCase_					=       key.replace("""encoder."""      ,					""""""							)

														if "qkv" in key:
																					lowerCamelCase_					=       key.split("""."""							)
																					if key.startswith("""decoder.blocks"""							):
																												lowerCamelCase_					=       config.decoder_hidden_size
																												lowerCamelCase_					=       int(key_split[2]							)
																												lowerCamelCase_					=       """decoder.decoder_layers."""
																												if "weight" in key:
																																			lowerCamelCase_					=       val[:dim, :]
																																			lowerCamelCase_					=       val[dim : dim * 2, :]
																																			lowerCamelCase_					=       val[-dim:, :]
																					else:
																												lowerCamelCase_					=       config.hidden_size
																												lowerCamelCase_					=       int(key_split[1]							)
																												lowerCamelCase_					=       """videomae.encoder.layer."""
																												if "weight" in key:
																																			lowerCamelCase_					=       val[:dim, :]
																																			lowerCamelCase_					=       val[dim : dim * 2, :]
																																			lowerCamelCase_					=       val[-dim:, :]
														else:
																					lowerCamelCase_					=       val

							return orig_state_dict



def   __lowerCAmelCase    (				)	->     int:
							lowerCamelCase_					=       hf_hub_download(
							    repo_id="""hf-internal-testing/spaghetti-video"""      ,					filename="""eating_spaghetti.npy"""      ,					repo_type="""dataset"""							)
							lowerCamelCase_					=       np.load(UpperCAmelCase__							)
							return list(UpperCAmelCase__							)



def   __lowerCAmelCase    (				UpperCAmelCase__					:   Union[str, Any]      ,					UpperCAmelCase__					:   int      ,					UpperCAmelCase__					:   int      ,					UpperCAmelCase__					:   Any							)	->     Union[str, Any]:
							lowerCamelCase_					=       get_videomae_config(UpperCAmelCase__							)

							if "finetuned" in model_name:
														lowerCamelCase_					=       VideoMAEForVideoClassification(UpperCAmelCase__							)
							else:
														lowerCamelCase_					=       VideoMAEForPreTraining(UpperCAmelCase__							)

							# download original checkpoint, hosted on Google Drive
							lowerCamelCase_					=       """pytorch_model.bin"""
							gdown.cached_download(UpperCAmelCase__      ,					UpperCAmelCase__      ,					quiet=UpperCAmelCase__							)
							lowerCamelCase_					=       torch.load(UpperCAmelCase__      ,					map_location="""cpu"""							)
							if "model" in files:
														lowerCamelCase_					=       files["""model"""]
							else:
														lowerCamelCase_					=       files["""module"""]
							lowerCamelCase_					=       convert_state_dict(UpperCAmelCase__      ,					UpperCAmelCase__							)

							model.load_state_dict(UpperCAmelCase__							)
							model.eval()

							# verify model on basic input
							lowerCamelCase_					=       VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]      ,					image_std=[0.5, 0.5, 0.5]							)
							lowerCamelCase_					=       prepare_video()
							lowerCamelCase_					=       image_processor(UpperCAmelCase__      ,					return_tensors="""pt"""							)

							if "finetuned" not in model_name:
														lowerCamelCase_					=       hf_hub_download(repo_id="""hf-internal-testing/bool-masked-pos"""      ,					filename="""bool_masked_pos.pt"""							)
														lowerCamelCase_					=       torch.load(UpperCAmelCase__							)

							lowerCamelCase_					=       model(**UpperCAmelCase__							)
							lowerCamelCase_					=       outputs.logits

							lowerCamelCase_					=       [
							    """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_					=       torch.Size([1, 4_0_0]							)
														lowerCamelCase_					=       torch.tensor([-0.9_2_9_1, -0.4_0_6_1, -0.9_3_0_7]							)
							elif model_name == "videomae-small-finetuned-ssv2":
														lowerCamelCase_					=       torch.Size([1, 1_7_4]							)
														lowerCamelCase_					=       torch.tensor([0.2_6_7_1, -0.4_6_8_9, -0.8_2_3_5]							)
							elif model_name == "videomae-base":
														lowerCamelCase_					=       torch.Size([1, 1_4_0_8, 1_5_3_6]							)
														lowerCamelCase_					=       torch.tensor([[0.7_7_3_9, 0.7_9_6_8, 0.7_0_8_9], [0.6_7_0_1, 0.7_4_8_7, 0.6_2_0_9], [0.4_2_8_7, 0.5_1_5_8, 0.4_7_7_3]]							)
							elif model_name == "videomae-base-short":
														lowerCamelCase_					=       torch.Size([1, 1_4_0_8, 1_5_3_6]							)
														lowerCamelCase_					=       torch.tensor([[0.7_9_9_4, 0.9_6_1_2, 0.8_5_0_8], [0.7_4_0_1, 0.8_9_5_8, 0.8_3_0_2], [0.5_8_6_2, 0.7_4_6_8, 0.7_3_2_5]]							)
														# we verified the loss both for normalized and unnormalized targets for this one
														lowerCamelCase_					=       torch.tensor([0.5_1_4_2]							) if config.norm_pix_loss else torch.tensor([0.6_4_6_9]							)
							elif model_name == "videomae-large":
														lowerCamelCase_					=       torch.Size([1, 1_4_0_8, 1_5_3_6]							)
														lowerCamelCase_					=       torch.tensor([[0.7_1_4_9, 0.7_9_9_7, 0.6_9_6_6], [0.6_7_6_8, 0.7_8_6_9, 0.6_9_4_8], [0.5_1_3_9, 0.6_2_2_1, 0.5_6_0_5]]							)
							elif model_name == "videomae-large-finetuned-kinetics":
														lowerCamelCase_					=       torch.Size([1, 4_0_0]							)
														lowerCamelCase_					=       torch.tensor([0.0_7_7_1, 0.0_0_1_1, -0.3_6_2_5]							)
							elif model_name == "videomae-huge-finetuned-kinetics":
														lowerCamelCase_					=       torch.Size([1, 4_0_0]							)
														lowerCamelCase_					=       torch.tensor([0.2_4_3_3, 0.1_6_3_2, -0.4_8_9_4]							)
							elif model_name == "videomae-base-short-finetuned-kinetics":
														lowerCamelCase_					=       torch.Size([1, 4_0_0]							)
														lowerCamelCase_					=       torch.tensor([0.6_5_8_8, 0.0_9_9_0, -0.2_4_9_3]							)
							elif model_name == "videomae-base-finetuned-kinetics":
														lowerCamelCase_					=       torch.Size([1, 4_0_0]							)
														lowerCamelCase_					=       torch.tensor([0.3_6_6_9, -0.0_6_8_8, -0.2_4_2_1]							)
							elif model_name == "videomae-base-short-ssv2":
														lowerCamelCase_					=       torch.Size([1, 1_4_0_8, 1_5_3_6]							)
														lowerCamelCase_					=       torch.tensor([[0.4_7_1_2, 0.5_2_9_6, 0.5_7_8_6], [0.2_2_7_8, 0.2_7_2_9, 0.4_0_2_6], [0.0_3_5_2, 0.0_7_3_0, 0.2_5_0_6]]							)
							elif model_name == "videomae-base-short-finetuned-ssv2":
														lowerCamelCase_					=       torch.Size([1, 1_7_4]							)
														lowerCamelCase_					=       torch.tensor([-0.0_5_3_7, -0.1_5_3_9, -0.3_2_6_6]							)
							elif model_name == "videomae-base-ssv2":
														lowerCamelCase_					=       torch.Size([1, 1_4_0_8, 1_5_3_6]							)
														lowerCamelCase_					=       torch.tensor([[0.8_1_3_1, 0.8_7_2_7, 0.8_5_4_6], [0.7_3_6_6, 0.9_3_7_7, 0.8_8_7_0], [0.5_9_3_5, 0.8_8_7_4, 0.8_5_6_4]]							)
							elif model_name == "videomae-base-finetuned-ssv2":
														lowerCamelCase_					=       torch.Size([1, 1_7_4]							)
														lowerCamelCase_					=       torch.tensor([0.1_9_6_1, -0.8_3_3_7, -0.6_3_8_9]							)
							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]      ,					UpperCAmelCase__      ,					atol=1e-4							)
							else:
														print("""Logits:"""      ,					logits[0, :3, :3]							)
														assert torch.allclose(logits[0, :3, :3]      ,					UpperCAmelCase__      ,					atol=1e-4							)
							print("""Logits ok!"""							)

							# verify loss, if applicable
							if model_name == "videomae-base-short":
														lowerCamelCase_					=       outputs.loss
														assert torch.allclose(UpperCAmelCase__      ,					UpperCAmelCase__      ,					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(UpperCAmelCase__							)
														model.save_pretrained(UpperCAmelCase__							)

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


if __name__ == "__main__":
		lowercase   =	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.'''
		)

		lowercase   =	parser.parse_args()
		convert_videomae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)







from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowercase   =	logging.get_logger(__name__)

lowercase   =	{
    '''abeja/gpt-neox-japanese-2.7b''': '''https://huggingface.co/abeja/gpt-neox-japanese-2.7b/resolve/main/config.json''',
}





class    __A(					UpperCAmelCase							):
					SCREAMING_SNAKE_CASE           =     '''gpt_neox_japanese'''

					def __init__(						self       :  Union[str, Any] ,       __UpperCamelCase       :  str=3_2_0_0_0 ,       __UpperCamelCase       :  List[Any]=2_5_6_0 ,       __UpperCamelCase       :  Any=3_2 ,       __UpperCamelCase       :  List[str]=3_2 ,       __UpperCamelCase       :  List[str]=4 ,       __UpperCamelCase       :  Dict="gelu" ,       __UpperCamelCase       :  List[Any]=1.00 ,       __UpperCamelCase       :  Any=1_0_0_0_0 ,       __UpperCamelCase       :  Optional[Any]=2_0_4_8 ,       __UpperCamelCase       :  Tuple=0.02 ,       __UpperCamelCase       :  List[str]=1E-5 ,       __UpperCamelCase       :  str=True ,       __UpperCamelCase       :  str=3_1_9_9_6 ,       __UpperCamelCase       :  int=3_1_9_9_9 ,       __UpperCamelCase       :  Optional[Any]=0.1 ,       __UpperCamelCase       :  Tuple=0.0 ,       **__UpperCamelCase       :  List[str] ,       ):
												super().__init__(bos_token_id=__UpperCamelCase ,       eos_token_id=__UpperCamelCase ,       **__UpperCamelCase )
												lowerCamelCase_					=       vocab_size
												lowerCamelCase_					=       max_position_embeddings
												lowerCamelCase_					=       hidden_size
												lowerCamelCase_					=       num_hidden_layers
												lowerCamelCase_					=       num_attention_heads
												lowerCamelCase_					=       intermediate_multiple_size
												lowerCamelCase_					=       hidden_act
												lowerCamelCase_					=       rotary_pct
												lowerCamelCase_					=       rotary_emb_base
												lowerCamelCase_					=       initializer_range
												lowerCamelCase_					=       layer_norm_eps
												lowerCamelCase_					=       use_cache
												lowerCamelCase_					=       attention_dropout
												lowerCamelCase_					=       hidden_dropout





