





"""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       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	VideoMAEConfig()

		set_architecture_configs(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    )

		if "finetuned" not in model_name:
				UpperCamelCase   =	False

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

		return config


def       a__					(					_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		if "small" in model_name:
				UpperCamelCase   =	384
				UpperCamelCase   =	1_536
				UpperCamelCase   =	12
				UpperCamelCase   =	16
				UpperCamelCase   =	12
				UpperCamelCase   =	3
				UpperCamelCase   =	192
				UpperCamelCase   =	768
		elif "large" in model_name:
				UpperCamelCase   =	1_024
				UpperCamelCase   =	4_096
				UpperCamelCase   =	24
				UpperCamelCase   =	16
				UpperCamelCase   =	12
				UpperCamelCase   =	8
				UpperCamelCase   =	512
				UpperCamelCase   =	2_048
		elif "huge" in model_name:
				UpperCamelCase   =	1_280
				UpperCamelCase   =	5_120
				UpperCamelCase   =	32
				UpperCamelCase   =	16
				UpperCamelCase   =	12
				UpperCamelCase   =	8
				UpperCamelCase   =	640
				UpperCamelCase   =	2_560
		elif "base" not in model_name:
				raise ValueError("Model name should include either \"small\", \"base\", \"large\", or \"huge\""    )


def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


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

		return name


def       a__					(					_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		for key in orig_state_dict.copy().keys():
				UpperCamelCase   =	orig_state_dict.pop(_SCREAMING_SNAKE_CASE    )

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

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

		return orig_state_dict


def       a__					(					):




		"""simple docstring"""


		UpperCamelCase   =	hf_hub_download(
		    repo_id="hf-internal-testing/spaghetti-video"  ,	filename="eating_spaghetti.npy"  ,	repo_type="dataset"    )
		UpperCamelCase   =	np.load(_SCREAMING_SNAKE_CASE    )
		return list(_SCREAMING_SNAKE_CASE    )


def       a__					(					_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		UpperCamelCase   =	get_videomae_config(_SCREAMING_SNAKE_CASE    )

		if "finetuned" in model_name:
				UpperCamelCase   =	VideoMAEForVideoClassification(_SCREAMING_SNAKE_CASE    )
		else:
				UpperCamelCase   =	VideoMAEForPreTraining(_SCREAMING_SNAKE_CASE    )

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

		model.load_state_dict(_SCREAMING_SNAKE_CASE    )
		model.eval()

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

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

		UpperCamelCase   =	model(**_SCREAMING_SNAKE_CASE    )
		UpperCamelCase   =	outputs.logits

		UpperCamelCase   =	[
		    "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":
				UpperCamelCase   =	torch.Size([1, 400]    )
				UpperCamelCase   =	torch.tensor([-0.92_91, -0.40_61, -0.93_07]    )
		elif model_name == "videomae-small-finetuned-ssv2":
				UpperCamelCase   =	torch.Size([1, 174]    )
				UpperCamelCase   =	torch.tensor([0.26_71, -0.46_89, -0.82_35]    )
		elif model_name == "videomae-base":
				UpperCamelCase   =	torch.Size([1, 1_408, 1_536]    )
				UpperCamelCase   =	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":
				UpperCamelCase   =	torch.Size([1, 1_408, 1_536]    )
				UpperCamelCase   =	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
				UpperCamelCase   =	torch.tensor([0.51_42]    ) if config.norm_pix_loss else torch.tensor([0.64_69]    )
		elif model_name == "videomae-large":
				UpperCamelCase   =	torch.Size([1, 1_408, 1_536]    )
				UpperCamelCase   =	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":
				UpperCamelCase   =	torch.Size([1, 400]    )
				UpperCamelCase   =	torch.tensor([0.07_71, 0.00_11, -0.36_25]    )
		elif model_name == "videomae-huge-finetuned-kinetics":
				UpperCamelCase   =	torch.Size([1, 400]    )
				UpperCamelCase   =	torch.tensor([0.24_33, 0.16_32, -0.48_94]    )
		elif model_name == "videomae-base-short-finetuned-kinetics":
				UpperCamelCase   =	torch.Size([1, 400]    )
				UpperCamelCase   =	torch.tensor([0.65_88, 0.09_90, -0.24_93]    )
		elif model_name == "videomae-base-finetuned-kinetics":
				UpperCamelCase   =	torch.Size([1, 400]    )
				UpperCamelCase   =	torch.tensor([0.36_69, -0.06_88, -0.24_21]    )
		elif model_name == "videomae-base-short-ssv2":
				UpperCamelCase   =	torch.Size([1, 1_408, 1_536]    )
				UpperCamelCase   =	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":
				UpperCamelCase   =	torch.Size([1, 174]    )
				UpperCamelCase   =	torch.tensor([-0.05_37, -0.15_39, -0.32_66]    )
		elif model_name == "videomae-base-ssv2":
				UpperCamelCase   =	torch.Size([1, 1_408, 1_536]    )
				UpperCamelCase   =	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":
				UpperCamelCase   =	torch.Size([1, 174]    )
				UpperCamelCase   =	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]  ,	_SCREAMING_SNAKE_CASE  ,	atol=1e-4    )
		else:
				print("Logits:"  ,	logits[0, :3, :3]    )
				assert torch.allclose(logits[0, :3, :3]  ,	_SCREAMING_SNAKE_CASE  ,	atol=1e-4    )
		print("Logits ok!"    )

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

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


if __name__ == "__main__":
						lowerCAmelCase__     			=			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__     			=			parser.parse_args()
						convert_videomae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)













"""simple docstring"""





import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device


if is_torch_available():
						from transformers import AutoModelForSeqaSeqLM, AutoTokenizer





@require_torch
@require_sentencepiece
@require_tokenizers
class        _lowerCamelCase      (       unittest.TestCase ):


						@slow
						def     snake_case_   (self		) ->   List[str]:
								UpperCamelCase   =	AutoModelForSeqaSeqLM.from_pretrained("google/mt5-small"   ,							return_dict=__a		).to(__a		)
								UpperCamelCase   =	AutoTokenizer.from_pretrained("google/mt5-small"		)

								UpperCamelCase   =	tokenizer("Hello there"   ,							return_tensors="pt"		).input_ids
								UpperCamelCase   =	tokenizer("Hi I am"   ,							return_tensors="pt"		).input_ids

								UpperCamelCase   =	model(input_ids.to(__a		)   ,							labels=labels.to(__a		)		).loss
								UpperCamelCase   =	-(labels.shape[-1] * loss.item())

								UpperCamelCase   =	-84.9127
								self.assertTrue(abs(mtf_score - EXPECTED_SCORE		) < 1e-4		)






