



"""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  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							Dict ):
						"""simple docstring"""






						snake_case   :			Optional[int]						    = VideoMAEConfig()

						set_architecture_configs(lowerCamelCase_     ,		lowerCamelCase_ )

						if "finetuned" not in model_name:
												snake_case   :			Tuple						    = False

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

						return config

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							Optional[Any]     ,		lowerCamelCase_:							int ):
						"""simple docstring"""






						if "small" in model_name:
												snake_case   :			Any						    = 3_8_4
												snake_case   :			int						    = 1_5_3_6
												snake_case   :			List[str]						    = 1_2
												snake_case   :			List[Any]						    = 1_6
												snake_case   :			Dict						    = 1_2
												snake_case   :			List[str]						    = 3
												snake_case   :			Optional[int]						    = 1_9_2
												snake_case   :			List[Any]						    = 7_6_8
						elif "large" in model_name:
												snake_case   :			Union[str, Any]						    = 1_0_2_4
												snake_case   :			List[Any]						    = 4_0_9_6
												snake_case   :			str						    = 2_4
												snake_case   :			Optional[Any]						    = 1_6
												snake_case   :			Optional[Any]						    = 1_2
												snake_case   :			Dict						    = 8
												snake_case   :			Dict						    = 5_1_2
												snake_case   :			List[str]						    = 2_0_4_8
						elif "huge" in model_name:
												snake_case   :			List[Any]						    = 1_2_8_0
												snake_case   :			Optional[int]						    = 5_1_2_0
												snake_case   :			Tuple						    = 3_2
												snake_case   :			List[str]						    = 1_6
												snake_case   :			Optional[Any]						    = 1_2
												snake_case   :			Tuple						    = 8
												snake_case   :			Union[str, Any]						    = 6_4_0
												snake_case   :			List[str]						    = 2_5_6_0
						elif "base" not in model_name:
												raise ValueError("Model name should include either \"small\", \"base\", \"large\", or \"huge\"" )

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							Tuple ):
						"""simple docstring"""






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

						return name

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							Any     ,		lowerCamelCase_:							str ):
						"""simple docstring"""






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

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

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

						return orig_state_dict

def  __SCREAMING_SNAKE_CASE		(			):
						"""simple docstring"""






						snake_case   :			List[Any]						    = hf_hub_download(
						    repo_id="hf-internal-testing/spaghetti-video"     ,		filename="eating_spaghetti.npy"     ,		repo_type="dataset" )
						snake_case   :			Optional[Any]						    = np.load(lowerCamelCase_ )
						return list(lowerCamelCase_ )

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							Optional[int]     ,		lowerCamelCase_:							str     ,		lowerCamelCase_:							Optional[Any]     ,		lowerCamelCase_:							Optional[Any] ):
						"""simple docstring"""






						snake_case   :			Tuple						    = get_videomae_config(lowerCamelCase_ )

						if "finetuned" in model_name:
												snake_case   :			Optional[Any]						    = VideoMAEForVideoClassification(lowerCamelCase_ )
						else:
												snake_case   :			List[str]						    = VideoMAEForPreTraining(lowerCamelCase_ )

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

						model.load_state_dict(lowerCamelCase_ )
						model.eval()

						# verify model on basic input
						snake_case   :			Any						    = VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]     ,		image_std=[0.5, 0.5, 0.5] )
						snake_case   :			List[str]						    = prepare_video()
						snake_case   :			str						    = image_processor(lowerCamelCase_     ,		return_tensors="pt" )

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

						snake_case   :			int						    = model(**lowerCamelCase_ )
						snake_case   :			str						    = outputs.logits

						snake_case   :			List[Any]						    = [
						    "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":
												snake_case   :			Any						    = torch.Size([1, 4_0_0] )
												snake_case   :			int						    = torch.tensor([-0.9_2_9_1, -0.4_0_6_1, -0.9_3_0_7] )
						elif model_name == "videomae-small-finetuned-ssv2":
												snake_case   :			List[Any]						    = torch.Size([1, 1_7_4] )
												snake_case   :			Optional[Any]						    = torch.tensor([0.2_6_7_1, -0.4_6_8_9, -0.8_2_3_5] )
						elif model_name == "videomae-base":
												snake_case   :			List[Any]						    = torch.Size([1, 1_4_0_8, 1_5_3_6] )
												snake_case   :			Union[str, Any]						    = 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":
												snake_case   :			Any						    = torch.Size([1, 1_4_0_8, 1_5_3_6] )
												snake_case   :			Optional[int]						    = 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
												snake_case   :			int						    = 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":
												snake_case   :			List[str]						    = torch.Size([1, 1_4_0_8, 1_5_3_6] )
												snake_case   :			Any						    = 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":
												snake_case   :			List[str]						    = torch.Size([1, 4_0_0] )
												snake_case   :			Optional[Any]						    = torch.tensor([0.0_7_7_1, 0.0_0_1_1, -0.3_6_2_5] )
						elif model_name == "videomae-huge-finetuned-kinetics":
												snake_case   :			Optional[Any]						    = torch.Size([1, 4_0_0] )
												snake_case   :			List[Any]						    = 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":
												snake_case   :			Optional[int]						    = torch.Size([1, 4_0_0] )
												snake_case   :			Optional[int]						    = torch.tensor([0.6_5_8_8, 0.0_9_9_0, -0.2_4_9_3] )
						elif model_name == "videomae-base-finetuned-kinetics":
												snake_case   :			List[str]						    = torch.Size([1, 4_0_0] )
												snake_case   :			Optional[Any]						    = torch.tensor([0.3_6_6_9, -0.0_6_8_8, -0.2_4_2_1] )
						elif model_name == "videomae-base-short-ssv2":
												snake_case   :			Optional[Any]						    = torch.Size([1, 1_4_0_8, 1_5_3_6] )
												snake_case   :			int						    = 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":
												snake_case   :			Dict						    = torch.Size([1, 1_7_4] )
												snake_case   :			Any						    = torch.tensor([-0.0_5_3_7, -0.1_5_3_9, -0.3_2_6_6] )
						elif model_name == "videomae-base-ssv2":
												snake_case   :			List[str]						    = torch.Size([1, 1_4_0_8, 1_5_3_6] )
												snake_case   :			Union[str, Any]						    = 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":
												snake_case   :			Tuple						    = torch.Size([1, 1_7_4] )
												snake_case   :			Dict						    = 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]     ,		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":
												snake_case   :			Any						    = 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__":
			A		=						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.'
			)

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








"""simple docstring"""







import contextlib
import importlib
import io
import unittest

import transformers

# Try to import everything from transformers to ensure every object can be loaded.
from transformers import *  # noqa F406
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, require_flax, require_tf, require_torch
from transformers.utils import ContextManagers, find_labels, is_flax_available, is_tf_available, is_torch_available


if is_torch_available():
			from transformers import BertForPreTraining, BertForQuestionAnswering, BertForSequenceClassification

if is_tf_available():
			from transformers import TFBertForPreTraining, TFBertForQuestionAnswering, TFBertForSequenceClassification

if is_flax_available():
			from transformers import FlaxBertForPreTraining, FlaxBertForQuestionAnswering, FlaxBertForSequenceClassification


A		=						DUMMY_UNKNOWN_IDENTIFIER
# An actual model hosted on huggingface.co

A		=						'main'
# Default branch name
A		=						'f2c752cfc5c0ab6f4bdec59acea69eefbee381c2'
# One particular commit (not the top of `main`)
A		=						'aaaaaaa'
# This commit does not exist, so we should 404.

A		=						'd9e9f15bc825e4b2c9249e9578f884bbcb5e3684'
# Sha-1 of config.json on the top of `main`, for checking purposes
A		=						'4b243c475af8d0a7754e87d7d096c92e5199ec2fe168a2ee7998e3b8e9bcb1d3'

@contextlib.contextmanager
def  __SCREAMING_SNAKE_CASE		(			):
						"""simple docstring"""






						print("Welcome!" )
						yield
						print("Bye!" )

@contextlib.contextmanager
def  __SCREAMING_SNAKE_CASE		(			):
						"""simple docstring"""






						print("Bonjour!" )
						yield
						print("Au revoir!" )

class       _a  ( unittest.TestCase):



			def 				__lowercase					(			self							:				Dict		)							->					Optional[int]:
									# If the spec is missing, importlib would not be able to import the module dynamically.
									assert transformers.__spec__ is not None
									assert importlib.util.find_spec("transformers"		) is not None



class       _a  ( unittest.TestCase):



			@unittest.mock.patch("sys.stdout"   ,	new_callable=io.StringIO		)
			def 				__lowercase					(			self							:				Union[str, Any]   ,	_lowercase							:				Optional[int]		)							->					Optional[Any]:
									with ContextManagers([]		):
															print("Transformers are awesome!"		)
									# The print statement adds a new line at the end of the output
									self.assertEqual(mock_stdout.getvalue()   ,	"Transformers are awesome!\n"		)



			@unittest.mock.patch("sys.stdout"   ,	new_callable=io.StringIO		)
			def 				__lowercase					(			self							:				Any   ,	_lowercase							:				Any		)							->					Dict:
									with ContextManagers([context_en()]		):
															print("Transformers are awesome!"		)
									# The output should be wrapped with an English welcome and goodbye
									self.assertEqual(mock_stdout.getvalue()   ,	"Welcome!\nTransformers are awesome!\nBye!\n"		)



			@unittest.mock.patch("sys.stdout"   ,	new_callable=io.StringIO		)
			def 				__lowercase					(			self							:				str   ,	_lowercase							:				List[Any]		)							->					Optional[int]:
									with ContextManagers([context_fr(), context_en()]		):
															print("Transformers are awesome!"		)
									# The output should be wrapped with an English and French welcome and goodbye
									self.assertEqual(mock_stdout.getvalue()   ,	"Bonjour!\nWelcome!\nTransformers are awesome!\nBye!\nAu revoir!\n"		)



			@require_torch
			def 				__lowercase					(			self							:				Any		)							->					Union[str, Any]:
									self.assertEqual(find_labels(_lowercase		)   ,	["labels"]		)
									self.assertEqual(find_labels(_lowercase		)   ,	["labels", "next_sentence_label"]		)
									self.assertEqual(find_labels(_lowercase		)   ,	["start_positions", "end_positions"]		)

									class       _a  ( SCREAMING_SNAKE_CASE__):
												pass



									self.assertEqual(find_labels(_lowercase		)   ,	["labels"]		)



			@require_tf
			def 				__lowercase					(			self							:				Optional[Any]		)							->					int:
									self.assertEqual(find_labels(_lowercase		)   ,	["labels"]		)
									self.assertEqual(find_labels(_lowercase		)   ,	["labels", "next_sentence_label"]		)
									self.assertEqual(find_labels(_lowercase		)   ,	["start_positions", "end_positions"]		)

									class       _a  ( SCREAMING_SNAKE_CASE__):
												pass



									self.assertEqual(find_labels(_lowercase		)   ,	["labels"]		)



			@require_flax
			def 				__lowercase					(			self							:				Tuple		)							->					List[Any]:
									# Flax models don't have labels
									self.assertEqual(find_labels(_lowercase		)   ,	[]		)
									self.assertEqual(find_labels(_lowercase		)   ,	[]		)
									self.assertEqual(find_labels(_lowercase		)   ,	[]		)

									class       _a  ( SCREAMING_SNAKE_CASE__):
												pass



									self.assertEqual(find_labels(_lowercase		)   ,	[]		)



