



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 (       lowerCamelCase__						):

				"""simple docstring"""


				lowercase__					:   Dict													=					VideoMAEConfig()

				set_architecture_configs(lowerCamelCase__					,			lowerCamelCase__						)

				if "finetuned" not in model_name:
								lowercase__					:   str													=					False

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

				return config


def  __lowerCamelCase (       lowerCamelCase__					,			lowerCamelCase__						):

				"""simple docstring"""


				if "small" in model_name:
								lowercase__					:   List[str]													=					384
								lowercase__					:   List[Any]													=					1_536
								lowercase__					:   Tuple													=					12
								lowercase__					:   Optional[Any]													=					16
								lowercase__					:   Union[str, Any]													=					12
								lowercase__					:   List[Any]													=					3
								lowercase__					:   List[Any]													=					192
								lowercase__					:   Optional[Any]													=					768
				elif "large" in model_name:
								lowercase__					:   Dict													=					1_024
								lowercase__					:   Tuple													=					4_096
								lowercase__					:   Any													=					24
								lowercase__					:   Union[str, Any]													=					16
								lowercase__					:   str													=					12
								lowercase__					:   str													=					8
								lowercase__					:   str													=					512
								lowercase__					:   Union[str, Any]													=					2_048
				elif "huge" in model_name:
								lowercase__					:   Union[str, Any]													=					1_280
								lowercase__					:   Optional[Any]													=					5_120
								lowercase__					:   Any													=					32
								lowercase__					:   Dict													=					16
								lowercase__					:   int													=					12
								lowercase__					:   Optional[int]													=					8
								lowercase__					:   Optional[Any]													=					640
								lowercase__					:   Optional[Any]													=					2_560
				elif "base" not in model_name:
								raise ValueError("Model name should include either \"small\", \"base\", \"large\", or \"huge\""						)


def  __lowerCamelCase (       lowerCamelCase__						):

				"""simple docstring"""


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

				return name


def  __lowerCamelCase (       lowerCamelCase__					,			lowerCamelCase__						):

				"""simple docstring"""


				for key in orig_state_dict.copy().keys():
								lowercase__					:   Optional[Any]													=					orig_state_dict.pop(lowerCamelCase__						)

								if key.startswith("encoder."						):
												lowercase__					:   List[Any]													=					key.replace("encoder."					,			""						)

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

				return orig_state_dict


def  __lowerCamelCase (       ):

				"""simple docstring"""


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


def  __lowerCamelCase (       lowerCamelCase__					,			lowerCamelCase__					,			lowerCamelCase__					,			lowerCamelCase__						):

				"""simple docstring"""


				lowercase__					:   Dict													=					get_videomae_config(lowerCamelCase__						)

				if "finetuned" in model_name:
								lowercase__					:   Any													=					VideoMAEForVideoClassification(lowerCamelCase__						)
				else:
								lowercase__					:   Union[str, Any]													=					VideoMAEForPreTraining(lowerCamelCase__						)

				# download original checkpoint, hosted on Google Drive
				lowercase__					:   Dict													=					"pytorch_model.bin"
				gdown.cached_download(lowerCamelCase__					,			lowerCamelCase__					,			quiet=lowerCamelCase__						)
				lowercase__					:   Dict													=					torch.load(lowerCamelCase__					,			map_location="cpu"						)
				if "model" in files:
								lowercase__					:   List[Any]													=					files["model"]
				else:
								lowercase__					:   Optional[int]													=					files["module"]
				lowercase__					:   Dict													=					convert_state_dict(lowerCamelCase__					,			lowerCamelCase__						)

				model.load_state_dict(lowerCamelCase__						)
				model.eval()

				# verify model on basic input
				lowercase__					:   List[Any]													=					VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]					,			image_std=[0.5, 0.5, 0.5]						)
				lowercase__					:   Optional[int]													=					prepare_video()
				lowercase__					:   Tuple													=					image_processor(lowerCamelCase__					,			return_tensors="pt"						)

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

				lowercase__					:   int													=					model(**lowerCamelCase__						)
				lowercase__					:   List[Any]													=					outputs.logits

				lowercase__					:   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":
								lowercase__					:   Union[str, Any]													=					torch.Size([1, 400]						)
								lowercase__					:   Optional[int]													=					torch.tensor([-0.9291, -0.4061, -0.9307]						)
				elif model_name == "videomae-small-finetuned-ssv2":
								lowercase__					:   Optional[Any]													=					torch.Size([1, 174]						)
								lowercase__					:   Dict													=					torch.tensor([0.2671, -0.4689, -0.8235]						)
				elif model_name == "videomae-base":
								lowercase__					:   Dict													=					torch.Size([1, 1_408, 1_536]						)
								lowercase__					:   str													=					torch.tensor([[0.7739, 0.7968, 0.7089], [0.6701, 0.7487, 0.6209], [0.4287, 0.5158, 0.4773]]						)
				elif model_name == "videomae-base-short":
								lowercase__					:   Optional[Any]													=					torch.Size([1, 1_408, 1_536]						)
								lowercase__					:   str													=					torch.tensor([[0.7994, 0.9612, 0.8508], [0.7401, 0.8958, 0.8302], [0.5862, 0.7468, 0.7325]]						)
								# we verified the loss both for normalized and unnormalized targets for this one
								lowercase__					:   Dict													=					torch.tensor([0.5142]						) if config.norm_pix_loss else torch.tensor([0.6469]						)
				elif model_name == "videomae-large":
								lowercase__					:   Optional[Any]													=					torch.Size([1, 1_408, 1_536]						)
								lowercase__					:   Dict													=					torch.tensor([[0.7149, 0.7997, 0.6966], [0.6768, 0.7869, 0.6948], [0.5139, 0.6221, 0.5605]]						)
				elif model_name == "videomae-large-finetuned-kinetics":
								lowercase__					:   Any													=					torch.Size([1, 400]						)
								lowercase__					:   str													=					torch.tensor([0.0771, 0.0011, -0.3625]						)
				elif model_name == "videomae-huge-finetuned-kinetics":
								lowercase__					:   str													=					torch.Size([1, 400]						)
								lowercase__					:   Union[str, Any]													=					torch.tensor([0.2433, 0.1632, -0.4894]						)
				elif model_name == "videomae-base-short-finetuned-kinetics":
								lowercase__					:   Any													=					torch.Size([1, 400]						)
								lowercase__					:   List[Any]													=					torch.tensor([0.6588, 0.0990, -0.2493]						)
				elif model_name == "videomae-base-finetuned-kinetics":
								lowercase__					:   Dict													=					torch.Size([1, 400]						)
								lowercase__					:   List[Any]													=					torch.tensor([0.3669, -0.0688, -0.2421]						)
				elif model_name == "videomae-base-short-ssv2":
								lowercase__					:   List[Any]													=					torch.Size([1, 1_408, 1_536]						)
								lowercase__					:   Tuple													=					torch.tensor([[0.4712, 0.5296, 0.5786], [0.2278, 0.2729, 0.4026], [0.0352, 0.0730, 0.2506]]						)
				elif model_name == "videomae-base-short-finetuned-ssv2":
								lowercase__					:   Dict													=					torch.Size([1, 174]						)
								lowercase__					:   List[str]													=					torch.tensor([-0.0537, -0.1539, -0.3266]						)
				elif model_name == "videomae-base-ssv2":
								lowercase__					:   List[str]													=					torch.Size([1, 1_408, 1_536]						)
								lowercase__					:   Optional[int]													=					torch.tensor([[0.8131, 0.8727, 0.8546], [0.7366, 0.9377, 0.8870], [0.5935, 0.8874, 0.8564]]						)
				elif model_name == "videomae-base-finetuned-ssv2":
								lowercase__					:   Optional[Any]													=					torch.Size([1, 174]						)
								lowercase__					:   Optional[int]													=					torch.tensor([0.1961, -0.8337, -0.6389]						)
				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":
								lowercase__					:   Optional[int]													=					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__					  =  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)









from collections import Counter
from timeit import timeit


def  __lowerCamelCase (       lowerCamelCase__ = ""					,			):

				"""simple docstring"""


				return sum(c % 2 for c in Counter(input_str.replace(" "					,			""						).lower()						).values()						) < 2


def  __lowerCamelCase (       lowerCamelCase__ = ""						):

				"""simple docstring"""


				if len(lowerCamelCase__						) == 0:
								return True
				lowercase__					:   str													=					input_str.replace(" "					,			""						).lower()
				# character_freq_dict: Stores the frequency of every character in the input string
				lowercase__					:   dict[str, int]													=					{}

				for character in lower_case_input_str:
								lowercase__					:   List[Any]													=					character_freq_dict.get(lowerCamelCase__					,			0						) + 1
				lowercase__					:   List[Any]													=					0

				for character_count in character_freq_dict.values():
								if character_count % 2:
												odd_char += 1
				if odd_char > 1:
								return False
				return True


def  __lowerCamelCase (       lowerCamelCase__ = ""						):

				"""simple docstring"""


				print("\nFor string = "					,			lowerCamelCase__					,			":"						)
				print(
				    "> can_string_be_rearranged_as_palindrome_counter()"					,			"\tans ="					,			can_string_be_rearranged_as_palindrome_counter(lowerCamelCase__						)					,			"\ttime ="					,			timeit(
				        "z.can_string_be_rearranged_as_palindrome_counter(z.check_str)"					,			setup="import __main__ as z"					,			)					,			"seconds"					,			)
				print(
				    "> can_string_be_rearranged_as_palindrome()"					,			"\tans ="					,			can_string_be_rearranged_as_palindrome(lowerCamelCase__						)					,			"\ttime ="					,			timeit(
				        "z.can_string_be_rearranged_as_palindrome(z.check_str)"					,			setup="import __main__ as z"					,			)					,			"seconds"					,			)


if __name__ == "__main__":
				lowerCAmelCase__					  =  input(
				    '''Enter string to determine if it can be rearranged as a palindrome or not: '''
				).strip()
				benchmark(check_str)
				lowerCAmelCase__					  =  can_string_be_rearranged_as_palindrome_counter(check_str)
				print(f'''{check_str} can {"" if status else "not "}be rearranged as a palindrome''')




