





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__  (      __a				):
	snake_case_						:     Optional[Any]     =		VideoMAEConfig()

	set_architecture_configs(__a				,	__a				)

	if "finetuned" not in model_name:
		snake_case_						:     Union[str, Any]     =		False

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

	return config


def 							SCREAMING_SNAKE_CASE__  (      __a				,	__a				):
	if "small" in model_name:
		snake_case_						:     Any     =		3_84
		snake_case_						:     Any     =		15_36
		snake_case_						:     Optional[Any]     =		12
		snake_case_						:     List[Any]     =		16
		snake_case_						:     Optional[int]     =		12
		snake_case_						:     str     =		3
		snake_case_						:     Dict     =		1_92
		snake_case_						:     Tuple     =		7_68
	elif "large" in model_name:
		snake_case_						:     int     =		10_24
		snake_case_						:     Optional[Any]     =		40_96
		snake_case_						:     Any     =		24
		snake_case_						:     str     =		16
		snake_case_						:     Tuple     =		12
		snake_case_						:     List[str]     =		8
		snake_case_						:     str     =		5_12
		snake_case_						:     Union[str, Any]     =		20_48
	elif "huge" in model_name:
		snake_case_						:     Optional[int]     =		12_80
		snake_case_						:     Tuple     =		51_20
		snake_case_						:     List[Any]     =		32
		snake_case_						:     Optional[int]     =		16
		snake_case_						:     str     =		12
		snake_case_						:     Optional[Any]     =		8
		snake_case_						:     str     =		6_40
		snake_case_						:     Tuple     =		25_60
	elif "base" not in model_name:
		raise ValueError('Model name should include either "small", "base", "large", or "huge"'				)


def 							SCREAMING_SNAKE_CASE__  (      __a				):
	if "encoder." in name:
		snake_case_						:     Optional[int]     =		name.replace('encoder.'				,	''				)
	if "cls_token" in name:
		snake_case_						:     Optional[int]     =		name.replace('cls_token'				,	'videomae.embeddings.cls_token'				)
	if "decoder_pos_embed" in name:
		snake_case_						:     List[Any]     =		name.replace('decoder_pos_embed'				,	'decoder.decoder_pos_embed'				)
	if "pos_embed" in name and "decoder" not in name:
		snake_case_						:     List[str]     =		name.replace('pos_embed'				,	'videomae.embeddings.position_embeddings'				)
	if "patch_embed.proj" in name:
		snake_case_						:     List[Any]     =		name.replace('patch_embed.proj'				,	'videomae.embeddings.patch_embeddings.projection'				)
	if "patch_embed.norm" in name:
		snake_case_						:     Dict     =		name.replace('patch_embed.norm'				,	'videomae.embeddings.norm'				)
	if "decoder.blocks" in name:
		snake_case_						:     Optional[Any]     =		name.replace('decoder.blocks'				,	'decoder.decoder_layers'				)
	if "blocks" in name:
		snake_case_						:     Dict     =		name.replace('blocks'				,	'videomae.encoder.layer'				)
	if "attn.proj" in name:
		snake_case_						:     Optional[Any]     =		name.replace('attn.proj'				,	'attention.output.dense'				)
	if "attn" in name and "bias" not in name:
		snake_case_						:     int     =		name.replace('attn'				,	'attention.self'				)
	if "attn" in name:
		snake_case_						:     Tuple     =		name.replace('attn'				,	'attention.attention'				)
	if "norm1" in name:
		snake_case_						:     Any     =		name.replace('norm1'				,	'layernorm_before'				)
	if "norm2" in name:
		snake_case_						:     Dict     =		name.replace('norm2'				,	'layernorm_after'				)
	if "mlp.fc1" in name:
		snake_case_						:     str     =		name.replace('mlp.fc1'				,	'intermediate.dense'				)
	if "mlp.fc2" in name:
		snake_case_						:     Tuple     =		name.replace('mlp.fc2'				,	'output.dense'				)
	if "decoder_embed" in name:
		snake_case_						:     List[Any]     =		name.replace('decoder_embed'				,	'decoder.decoder_embed'				)
	if "decoder_norm" in name:
		snake_case_						:     int     =		name.replace('decoder_norm'				,	'decoder.decoder_norm'				)
	if "decoder_pred" in name:
		snake_case_						:     int     =		name.replace('decoder_pred'				,	'decoder.decoder_pred'				)
	if "norm.weight" in name and "decoder" not in name and "fc" not in name:
		snake_case_						:     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_						:     Optional[Any]     =		name.replace('norm.bias'				,	'videomae.layernorm.bias'				)
	if "head" in name and "decoder" not in name:
		snake_case_						:     List[str]     =		name.replace('head'				,	'classifier'				)

	return name


def 							SCREAMING_SNAKE_CASE__  (      __a				,	__a				):
	for key in orig_state_dict.copy().keys():
		snake_case_						:     Union[str, Any]     =		orig_state_dict.pop(__a				)

		if key.startswith('encoder.'				):
			snake_case_						:     str     =		key.replace('encoder.'				,	''				)

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

	return orig_state_dict


def 							SCREAMING_SNAKE_CASE__  (      ):
	snake_case_						:     List[Any]     =		hf_hub_download(
	    repo_id='hf-internal-testing/spaghetti-video'				,	filename='eating_spaghetti.npy'				,	repo_type='dataset'				)
	snake_case_						:     Any     =		np.load(__a				)
	return list(__a				)


def 							SCREAMING_SNAKE_CASE__  (      __a				,	__a				,	__a				,	__a				):
	snake_case_						:     Tuple     =		get_videomae_config(__a				)

	if "finetuned" in model_name:
		snake_case_						:     Optional[int]     =		VideoMAEForVideoClassification(__a				)
	else:
		snake_case_						:     List[Any]     =		VideoMAEForPreTraining(__a				)

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

	model.load_state_dict(__a				)
	model.eval()

	# verify model on basic input
	snake_case_						:     Tuple     =		VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]				,	image_std=[0.5, 0.5, 0.5]				)
	snake_case_						:     Optional[Any]     =		prepare_video()
	snake_case_						:     Dict     =		image_processor(__a				,	return_tensors='pt'				)

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

	snake_case_						:     List[str]     =		model(**__a				)
	snake_case_						:     Tuple     =		outputs.logits

	snake_case_						:     Optional[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_						:     Dict     =		torch.Size([1, 4_00]				)
		snake_case_						:     str     =		torch.tensor([-0.9291, -0.4061, -0.9307]				)
	elif model_name == "videomae-small-finetuned-ssv2":
		snake_case_						:     List[Any]     =		torch.Size([1, 1_74]				)
		snake_case_						:     str     =		torch.tensor([0.2671, -0.4689, -0.8235]				)
	elif model_name == "videomae-base":
		snake_case_						:     int     =		torch.Size([1, 14_08, 15_36]				)
		snake_case_						:     List[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":
		snake_case_						:     Any     =		torch.Size([1, 14_08, 15_36]				)
		snake_case_						:     Dict     =		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_						:     Any     =		torch.tensor([0.5142]				) if config.norm_pix_loss else torch.tensor([0.6469]				)
	elif model_name == "videomae-large":
		snake_case_						:     Optional[Any]     =		torch.Size([1, 14_08, 15_36]				)
		snake_case_						:     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":
		snake_case_						:     Optional[int]     =		torch.Size([1, 4_00]				)
		snake_case_						:     Union[str, Any]     =		torch.tensor([0.0771, 0.0011, -0.3625]				)
	elif model_name == "videomae-huge-finetuned-kinetics":
		snake_case_						:     List[Any]     =		torch.Size([1, 4_00]				)
		snake_case_						:     int     =		torch.tensor([0.2433, 0.1632, -0.4894]				)
	elif model_name == "videomae-base-short-finetuned-kinetics":
		snake_case_						:     Any     =		torch.Size([1, 4_00]				)
		snake_case_						:     int     =		torch.tensor([0.6588, 0.0990, -0.2493]				)
	elif model_name == "videomae-base-finetuned-kinetics":
		snake_case_						:     Tuple     =		torch.Size([1, 4_00]				)
		snake_case_						:     List[Any]     =		torch.tensor([0.3669, -0.0688, -0.2421]				)
	elif model_name == "videomae-base-short-ssv2":
		snake_case_						:     Dict     =		torch.Size([1, 14_08, 15_36]				)
		snake_case_						:     str     =		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_						:     List[Any]     =		torch.Size([1, 1_74]				)
		snake_case_						:     List[str]     =		torch.tensor([-0.0537, -0.1539, -0.3266]				)
	elif model_name == "videomae-base-ssv2":
		snake_case_						:     Any     =		torch.Size([1, 14_08, 15_36]				)
		snake_case_						:     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":
		snake_case_						:     Any     =		torch.Size([1, 1_74]				)
		snake_case_						:     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]				,	__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_						:     Dict     =		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__":
						_SCREAMING_SNAKE_CASE 		=						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."""
						)

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








import re
import string

import numpy as np

import datasets


_SCREAMING_SNAKE_CASE 		=						"""
Returns the rate at which the input predicted strings exactly match their references, ignoring any strings input as part of the regexes_to_ignore list.
"""

_SCREAMING_SNAKE_CASE 		=						"""
Args:
    predictions: List of predicted texts.
    references: List of reference texts.
    regexes_to_ignore: List, defaults to None. Regex expressions of characters to
        ignore when calculating the exact matches. Note: these regexes are removed
        from the input data before the changes based on the options below (e.g. ignore_case,
        ignore_punctuation, ignore_numbers) are applied.
    ignore_case: Boolean, defaults to False. If true, turns everything
        to lowercase so that capitalization differences are ignored.
    ignore_punctuation: Boolean, defaults to False. If true, removes all punctuation before
        comparing predictions and references.
    ignore_numbers: Boolean, defaults to False. If true, removes all punctuation before
        comparing predictions and references.
Returns:
    exact_match: Dictionary containing exact_match rate. Possible values are between 0.0 and 100.0, inclusive.
Examples:
    >>> exact_match = datasets.load_metric(\"exact_match\")
    >>> refs = [\"the cat\", \"theater\", \"YELLING\", \"agent007\"]
    >>> preds = [\"cat?\", \"theater\", \"yelling\", \"agent\"]
    >>> results = exact_match.compute(references=refs, predictions=preds)
    >>> print(round(results[\"exact_match\"], 1))
    25.0

    >>> exact_match = datasets.load_metric(\"exact_match\")
    >>> refs = [\"the cat\", \"theater\", \"YELLING\", \"agent007\"]
    >>> preds = [\"cat?\", \"theater\", \"yelling\", \"agent\"]
    >>> results = exact_match.compute(references=refs, predictions=preds, regexes_to_ignore=[\"the \", \"yell\"], ignore_case=True, ignore_punctuation=True)
    >>> print(round(results[\"exact_match\"], 1))
    50.0


    >>> exact_match = datasets.load_metric(\"exact_match\")
    >>> refs = [\"the cat\", \"theater\", \"YELLING\", \"agent007\"]
    >>> preds = [\"cat?\", \"theater\", \"yelling\", \"agent\"]
    >>> results = exact_match.compute(references=refs, predictions=preds, regexes_to_ignore=[\"the \", \"yell\", \"YELL\"], ignore_case=True, ignore_punctuation=True)
    >>> print(round(results[\"exact_match\"], 1))
    75.0

    >>> exact_match = datasets.load_metric(\"exact_match\")
    >>> refs = [\"the cat\", \"theater\", \"YELLING\", \"agent007\"]
    >>> preds = [\"cat?\", \"theater\", \"yelling\", \"agent\"]
    >>> results = exact_match.compute(references=refs, predictions=preds, regexes_to_ignore=[\"the \", \"yell\", \"YELL\"], ignore_case=True, ignore_punctuation=True, ignore_numbers=True)
    >>> print(round(results[\"exact_match\"], 1))
    100.0

    >>> exact_match = datasets.load_metric(\"exact_match\")
    >>> refs = [\"The cat sat on the mat.\", \"Theaters are great.\", \"It's like comparing oranges and apples.\"]
    >>> preds = [\"The cat sat on the mat?\", \"Theaters are great.\", \"It's like comparing apples and oranges.\"]
    >>> results = exact_match.compute(references=refs, predictions=preds)
    >>> print(round(results[\"exact_match\"], 1))
    33.3

"""

_SCREAMING_SNAKE_CASE 		=						"""
"""
@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION    ,			_KWARGS_DESCRIPTION   )
class SCREAMING_SNAKE_CASE_  (   datasets.Metric   ):






						def UpperCAmelCase_			( self : Union[str, Any]		)							->      Optional[Any]:







							"""simple docstring"""


							return datasets.MetricInfo(
							    description=_DESCRIPTION						,				citation=_CITATION						,				inputs_description=_KWARGS_DESCRIPTION						,				features=datasets.Features(
							        {
							            'predictions': datasets.Value('string'						,				id='sequence'		),
							            'references': datasets.Value('string'						,				id='sequence'		),
							        }		)						,				reference_urls=[]						,				)






						def UpperCAmelCase_			( self : List[Any]						,				_A : Optional[Any]						,				_A : Optional[int]						,				_A : Optional[int]=None						,				_A : Dict=False						,				_A : Dict=False						,				_A : Optional[Any]=False						,				)							->      List[str]:







							"""simple docstring"""


							if regexes_to_ignore is not None:
								for s in regexes_to_ignore:
									snake_case_						:     List[str]     =		np.array([re.sub(_A						,				''						,				_A		) for x in predictions]		)
									snake_case_						:     int     =		np.array([re.sub(_A						,				''						,				_A		) for x in references]		)
							else:
								snake_case_						:     Optional[Any]     =		np.asarray(_A		)
								snake_case_						:     Optional[Any]     =		np.asarray(_A		)

							if ignore_case:
								snake_case_						:     int     =		np.char.lower(_A		)
								snake_case_						:     List[str]     =		np.char.lower(_A		)

							if ignore_punctuation:
								snake_case_						:     str     =		string.punctuation.maketrans(''						,				''						,				string.punctuation		)
								snake_case_						:     str     =		np.char.translate(_A						,				table=_A		)
								snake_case_						:     Any     =		np.char.translate(_A						,				table=_A		)

							if ignore_numbers:
								snake_case_						:     int     =		string.digits.maketrans(''						,				''						,				string.digits		)
								snake_case_						:     Tuple     =		np.char.translate(_A						,				table=_A		)
								snake_case_						:     Optional[Any]     =		np.char.translate(_A						,				table=_A		)

							snake_case_						:     Optional[Any]     =		predictions == references

							return {"exact_match": np.mean(_A		) * 100}

