






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						:	Tuple ):
							snake_case__	  =       VideoMAEConfig()

							set_architecture_configs(a  ,					a )

							if "finetuned" not in model_name:
														snake_case__	  =       False

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

							return config

def _UpperCAmelCase   (	a						:	str  ,					a						:	str ):







							if "small" in model_name:
														snake_case__	  =       384
														snake_case__	  =       1536
														snake_case__	  =       12
														snake_case__	  =       16
														snake_case__	  =       12
														snake_case__	  =       3
														snake_case__	  =       192
														snake_case__	  =       768
							elif "large" in model_name:
														snake_case__	  =       1024
														snake_case__	  =       4096
														snake_case__	  =       24
														snake_case__	  =       16
														snake_case__	  =       12
														snake_case__	  =       8
														snake_case__	  =       512
														snake_case__	  =       2048
							elif "huge" in model_name:
														snake_case__	  =       1280
														snake_case__	  =       5120
														snake_case__	  =       32
														snake_case__	  =       16
														snake_case__	  =       12
														snake_case__	  =       8
														snake_case__	  =       640
														snake_case__	  =       2560
							elif "base" not in model_name:
														raise ValueError("""Model name should include either \"small\", \"base\", \"large\", or \"huge\"""" )

def _UpperCAmelCase   (	a						:	Dict ):







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

							return name

def _UpperCAmelCase   (	a						:	List[str]  ,					a						:	Union[str, Any] ):







							for key in orig_state_dict.copy().keys():
														snake_case__	  =       orig_state_dict.pop(a )

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

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

							return orig_state_dict

def _UpperCAmelCase   (	):
							snake_case__	  =       hf_hub_download(
							    repo_id="""hf-internal-testing/spaghetti-video"""  ,					filename="""eating_spaghetti.npy"""  ,					repo_type="""dataset""" )
							snake_case__	  =       np.load(a )
							return list(a )

def _UpperCAmelCase   (	a						:	Optional[Any]  ,					a						:	List[str]  ,					a						:	Optional[Any]  ,					a						:	Dict ):
							snake_case__	  =       get_videomae_config(a )

							if "finetuned" in model_name:
														snake_case__	  =       VideoMAEForVideoClassification(a )
							else:
														snake_case__	  =       VideoMAEForPreTraining(a )

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

							model.load_state_dict(a )
							model.eval()

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

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

							snake_case__	  =       model(**a )
							snake_case__	  =       outputs.logits

							snake_case__	  =       [
							    """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__	  =       torch.Size([1, 400] )
														snake_case__	  =       torch.tensor([-0.9291, -0.4061, -0.9307] )
							elif model_name == "videomae-small-finetuned-ssv2":
														snake_case__	  =       torch.Size([1, 174] )
														snake_case__	  =       torch.tensor([0.2671, -0.4689, -0.8235] )
							elif model_name == "videomae-base":
														snake_case__	  =       torch.Size([1, 1408, 1536] )
														snake_case__	  =       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":
														snake_case__	  =       torch.Size([1, 1408, 1536] )
														snake_case__	  =       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
														snake_case__	  =       torch.tensor([0.5142] ) if config.norm_pix_loss else torch.tensor([0.6469] )
							elif model_name == "videomae-large":
														snake_case__	  =       torch.Size([1, 1408, 1536] )
														snake_case__	  =       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":
														snake_case__	  =       torch.Size([1, 400] )
														snake_case__	  =       torch.tensor([0.0771, 0.0011, -0.3625] )
							elif model_name == "videomae-huge-finetuned-kinetics":
														snake_case__	  =       torch.Size([1, 400] )
														snake_case__	  =       torch.tensor([0.2433, 0.1632, -0.4894] )
							elif model_name == "videomae-base-short-finetuned-kinetics":
														snake_case__	  =       torch.Size([1, 400] )
														snake_case__	  =       torch.tensor([0.6588, 0.0990, -0.2493] )
							elif model_name == "videomae-base-finetuned-kinetics":
														snake_case__	  =       torch.Size([1, 400] )
														snake_case__	  =       torch.tensor([0.3669, -0.0688, -0.2421] )
							elif model_name == "videomae-base-short-ssv2":
														snake_case__	  =       torch.Size([1, 1408, 1536] )
														snake_case__	  =       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":
														snake_case__	  =       torch.Size([1, 174] )
														snake_case__	  =       torch.tensor([-0.0537, -0.1539, -0.3266] )
							elif model_name == "videomae-base-ssv2":
														snake_case__	  =       torch.Size([1, 1408, 1536] )
														snake_case__	  =       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":
														snake_case__	  =       torch.Size([1, 174] )
														snake_case__	  =       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]  ,					a  ,					atol=1e-4 )
							else:
														print("""Logits:"""  ,					logits[0, :3, :3] )
														assert torch.allclose(logits[0, :3, :3]  ,					a  ,					atol=1e-4 )
							print("""Logits ok!""" )

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

							if push_to_hub:
														print("""Pushing to the hub...""" )
														model.push_to_hub(a  ,					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)








import math
from collections.abc import Iterator
from itertools import takewhile

def _UpperCAmelCase   (	a						:	int ):

							if 1 < number < 4:
														# 2 and 3 are primes
														return True
							elif number < 2 or number % 2 == 0 or number % 3 == 0:
														# Negatives, 0, 1, all even numbers, all multiples of 3 are not primes
														return False

							# All primes number are in format of 6k +/- 1
							for i in range(5  ,					int(math.sqrt(a ) + 1 )  ,					6 ):
														if number % i == 0 or number % (i + 2) == 0:
																					return False
							return True

def _UpperCAmelCase   (	):
							snake_case__	  =       2
							while True:
														if is_prime(a ):
																					yield num
														num += 1

def _UpperCAmelCase   (	a						:	int = 200_0000 ):
							return sum(takewhile(lambda a : x < n  ,					prime_generator() ) )


if __name__ == "__main__":
	print(F'''{solution() = }''')
