





import logging
import os
import sys
from dataclasses import dataclass, field
from typing import Optional

import numpy as np
import torch
from datasets import load_dataset
from torchvision.transforms import Compose, Lambda, Normalize, RandomHorizontalFlip, RandomResizedCrop, ToTensor

import transformers
from transformers import (
    CONFIG_MAPPING,
    IMAGE_PROCESSOR_MAPPING,
    MODEL_FOR_MASKED_IMAGE_MODELING_MAPPING,
    AutoConfig,
    AutoImageProcessor,
    AutoModelForMaskedImageModeling,
    HfArgumentParser,
    Trainer,
    TrainingArguments,
)
from transformers.trainer_utils import get_last_checkpoint
from transformers.utils import check_min_version, send_example_telemetry
from transformers.utils.versions import require_version
UpperCAmelCase_ 							=       logging.getLogger(__name__)

# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
check_min_version("""4.31.0""")

require_version("""datasets>=1.8.0""", """To fix: pip install -r examples/pytorch/image-pretraining/requirements.txt""")

UpperCAmelCase_ 							=       list(MODEL_FOR_MASKED_IMAGE_MODELING_MAPPING.keys())
UpperCAmelCase_ 							=       tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
@dataclass
class      UpperCamelCase_						:
					lowerCAmelCase_							=		field(
					    default='''cifar10'''							,					metadata={'''help''': '''Name of a dataset from the datasets package'''}   )
					lowerCAmelCase_							=		field(
					    default=_lowerCamelCase							,					metadata={'''help''': '''The configuration name of the dataset to use (via the datasets library).'''}   )
					lowerCAmelCase_							=		field(
					    default=_lowerCamelCase							,					metadata={'''help''': '''The column name of the images in the files. If not set, will try to use \'image\' or \'img\'.'''}							,					)
					lowerCAmelCase_							=		field(default=_lowerCamelCase							,					metadata={'''help''': '''A folder containing the training data.'''}   )
					lowerCAmelCase_							=		field(default=_lowerCamelCase							,					metadata={'''help''': '''A folder containing the validation data.'''}   )
					lowerCAmelCase_							=		field(
					    default=0.15							,					metadata={'''help''': '''Percent to split off of train for validation.'''}   )
					lowerCAmelCase_							=		field(default=32							,					metadata={'''help''': '''The size of the square patches to use for masking.'''}   )
					lowerCAmelCase_							=		field(
					    default=0.6							,					metadata={'''help''': '''Percentage of patches to mask.'''}							,					)
					lowerCAmelCase_							=		field(
					    default=_lowerCamelCase							,					metadata={
					        '''help''': (
					            '''For debugging purposes or quicker training, truncate the number of training examples to this '''
					            '''value if set.'''
					        )
					    }							,					)
					lowerCAmelCase_							=		field(
					    default=_lowerCamelCase							,					metadata={
					        '''help''': (
					            '''For debugging purposes or quicker training, truncate the number of evaluation examples to this '''
					            '''value if set.'''
					        )
					    }							,					)




					def 		lowerCAmelCase     (						self					)						->							Any:
										_snake_case													=			{}
										if self.train_dir is not None:
															_snake_case													=			self.train_dir
										if self.validation_dir is not None:
															_snake_case													=			self.validation_dir
										_snake_case													=			data_files if data_files else None





@dataclass
class      UpperCamelCase_						:
					lowerCAmelCase_							=		field(
					    default=_lowerCamelCase							,					metadata={
					        '''help''': (
					            '''The model checkpoint for weights initialization. Can be a local path to a pytorch_model.bin or a '''
					            '''checkpoint identifier on the hub. '''
					            '''Don\'t set if you want to train a model from scratch.'''
					        )
					    }							,					)
					lowerCAmelCase_							=		field(
					    default=_lowerCamelCase							,					metadata={'''help''': '''If training from scratch, pass a model type from the list: ''' + ''', '''.join(_lowerCamelCase   )}							,					)
					lowerCAmelCase_							=		field(
					    default=_lowerCamelCase							,					metadata={'''help''': '''Pretrained config name or path if not the same as model_name'''}   )
					lowerCAmelCase_							=		field(
					    default=_lowerCamelCase							,					metadata={
					        '''help''': (
					            '''Override some existing default config settings when a model is trained from scratch. Example: '''
					            '''n_embd=10,resid_pdrop=0.2,scale_attn_weights=false,summary_type=cls_index'''
					        )
					    }							,					)
					lowerCAmelCase_							=		field(
					    default=_lowerCamelCase							,					metadata={'''help''': '''Where do you want to store (cache) the pretrained models/datasets downloaded from the hub'''}							,					)
					lowerCAmelCase_							=		field(
					    default='''main'''							,					metadata={'''help''': '''The specific model version to use (can be a branch name, tag name or commit id).'''}							,					)
					lowerCAmelCase_							=		field(default=_lowerCamelCase							,					metadata={'''help''': '''Name or path of preprocessor config.'''}   )
					lowerCAmelCase_							=		field(
					    default=_lowerCamelCase							,					metadata={
					        '''help''': (
					            '''Will use the token generated when running `huggingface-cli login` (necessary to use this script '''
					            '''with private models).'''
					        )
					    }							,					)
					lowerCAmelCase_							=		field(
					    default=_lowerCamelCase							,					metadata={
					        '''help''': (
					            '''The size (resolution) of each image. If not specified, will use `image_size` of the configuration.'''
					        )
					    }							,					)
					lowerCAmelCase_							=		field(
					    default=_lowerCamelCase							,					metadata={
					        '''help''': (
					            '''The size (resolution) of each patch. If not specified, will use `patch_size` of the configuration.'''
					        )
					    }							,					)
					lowerCAmelCase_							=		field(
					    default=_lowerCamelCase							,					metadata={'''help''': '''Stride to use for the encoder.'''}							,					)





class      UpperCamelCase_						:




					def __init__(						self			,  lowerCAmelCase_=192			,  lowerCAmelCase_=32			,  lowerCAmelCase_=4			,  lowerCAmelCase_=0.6					)						->							Tuple:
										_snake_case													=			input_size
										_snake_case													=			mask_patch_size
										_snake_case													=			model_patch_size
										_snake_case													=			mask_ratio

										if self.input_size % self.mask_patch_size != 0:
															raise ValueError('Input size must be divisible by mask patch size'					)
										if self.mask_patch_size % self.model_patch_size != 0:
															raise ValueError('Mask patch size must be divisible by model patch size'					)

										_snake_case													=			self.input_size // self.mask_patch_size
										_snake_case													=			self.mask_patch_size // self.model_patch_size

										_snake_case													=			self.rand_size**2
										_snake_case													=			int(np.ceil(self.token_count * self.mask_ratio					)					)






					def __call__(						self					)						->							Any:
										_snake_case													=			np.random.permutation(self.token_count					)[: self.mask_count]
										_snake_case													=			np.zeros(self.token_count			,  dtype=lowerCAmelCase_					)
										_snake_case													=			1

										_snake_case													=			mask.reshape((self.rand_size, self.rand_size)					)
										_snake_case													=			mask.repeat(self.scale			,  axis=0					).repeat(self.scale			,  axis=1					)

										return torch.tensor(mask.flatten()					)





def   lowerCamelCase__       (       UpperCamelCase__     :    Union[str, Any]  )     ->	Tuple:



					'''simple docstring'''




					_snake_case													=			torch.stack([example['pixel_values'] for example in examples]  )
					_snake_case													=			torch.stack([example['mask'] for example in examples]  )
					return {"pixel_values": pixel_values, "bool_masked_pos": mask}





def   lowerCamelCase__       (       )     ->	Union[str, Any]:



					'''simple docstring'''




					_snake_case													=			HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)  )
					if len(sys.argv  ) == 2 and sys.argv[1].endswith('.json'  ):
										# If we pass only one argument to the script and it's the path to a json file,
										# let's parse it to get our arguments.
										_snake_case							,       _snake_case							,       _snake_case													=			parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]  )  )
					else:
										_snake_case							,       _snake_case							,       _snake_case													=			parser.parse_args_into_dataclasses()

					# Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
					# information sent is the one passed as arguments along with your Python/PyTorch versions.
					send_example_telemetry('run_mim'    ,      UpperCamelCase__    ,      UpperCamelCase__  )

					# Setup logging
					logging.basicConfig(
					    format='%(asctime)s - %(levelname)s - %(name)s - %(message)s'    ,      datefmt='%m/%d/%Y %H:%M:%S'    ,      handlers=[logging.StreamHandler(sys.stdout  )]    ,      )

					if training_args.should_log:
										# The default of training_args.log_level is passive, so we set log level at info here to have that default.
										transformers.utils.logging.set_verbosity_info()

					_snake_case													=			training_args.get_process_log_level()
					logger.setLevel(UpperCamelCase__  )
					transformers.utils.logging.set_verbosity(UpperCamelCase__  )
					transformers.utils.logging.enable_default_handler()
					transformers.utils.logging.enable_explicit_format()

					# Log on each process the small summary:
					logger.warning(
					    F'''Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}'''
					    + F'''distributed training: {bool(training_args.local_rank != -1  )}, 16-bits training: {training_args.fpaa}'''  )
					logger.info(F'''Training/evaluation parameters {training_args}'''  )

					# Detecting last checkpoint.
					_snake_case													=			None
					if os.path.isdir(training_args.output_dir  ) and training_args.do_train and not training_args.overwrite_output_dir:
										_snake_case													=			get_last_checkpoint(training_args.output_dir  )
										if last_checkpoint is None and len(os.listdir(training_args.output_dir  )  ) > 0:
															raise ValueError(
															    F'''Output directory ({training_args.output_dir}) already exists and is not empty. '''
															    'Use --overwrite_output_dir to overcome.'  )
										elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
															logger.info(
															    F'''Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change '''
															    'the `--output_dir` or add `--overwrite_output_dir` to train from scratch.'  )

    # Initialize our dataset.
					_snake_case													=			load_dataset(
					    data_args.dataset_name    ,      data_args.dataset_config_name    ,      data_files=data_args.data_files    ,      cache_dir=model_args.cache_dir    ,      use_auth_token=True if model_args.use_auth_token else None    ,      )

					# If we don't have a validation split, split off a percentage of train as validation.
					_snake_case													=			None if 'validation' in ds.keys() else data_args.train_val_split
					if isinstance(data_args.train_val_split    ,      UpperCamelCase__  ) and data_args.train_val_split > 0.0:
										_snake_case													=			ds['train'].train_test_split(data_args.train_val_split  )
										_snake_case													=			split['train']
										_snake_case													=			split['test']

					# Create config
					# Distributed training:
					# The .from_pretrained methods guarantee that only one local process can concurrently
					# download model & vocab.
					_snake_case													=			{
					    'cache_dir': model_args.cache_dir,
					    'revision': model_args.model_revision,
					    'use_auth_token': True if model_args.use_auth_token else None,
					}
					if model_args.config_name_or_path:
										_snake_case													=			AutoConfig.from_pretrained(model_args.config_name_or_path    ,      **UpperCamelCase__  )
					elif model_args.model_name_or_path:
										_snake_case													=			AutoConfig.from_pretrained(model_args.model_name_or_path    ,      **UpperCamelCase__  )
					else:
										_snake_case													=			CONFIG_MAPPING[model_args.model_type]()
										logger.warning('You are instantiating a new config instance from scratch.'  )
										if model_args.config_overrides is not None:
															logger.info(F'''Overriding config: {model_args.config_overrides}'''  )
															config.update_from_string(model_args.config_overrides  )
															logger.info(F'''New config: {config}'''  )

    # make sure the decoder_type is "simmim" (only relevant for BEiT)
					if hasattr(UpperCamelCase__    ,      'decoder_type'  ):
										_snake_case													=			'simmim'

					# adapt config
					_snake_case													=			model_args.image_size if model_args.image_size is not None else config.image_size
					_snake_case													=			model_args.patch_size if model_args.patch_size is not None else config.patch_size
					_snake_case													=			(
					    model_args.encoder_stride if model_args.encoder_stride is not None else config.encoder_stride
					)

					config.update(
					    {
					        'image_size': model_args.image_size,
					        'patch_size': model_args.patch_size,
					        'encoder_stride': model_args.encoder_stride,
					    }  )

					# create image processor
					if model_args.image_processor_name:
										_snake_case													=			AutoImageProcessor.from_pretrained(model_args.image_processor_name    ,      **UpperCamelCase__  )
					elif model_args.model_name_or_path:
										_snake_case													=			AutoImageProcessor.from_pretrained(model_args.model_name_or_path    ,      **UpperCamelCase__  )
					else:
										_snake_case													=			{
										    conf.model_type: image_processor_class for conf, image_processor_class in IMAGE_PROCESSOR_MAPPING.items()
										}
										_snake_case													=			IMAGE_PROCESSOR_TYPES[model_args.model_type]()

					# create model
					if model_args.model_name_or_path:
										_snake_case													=			AutoModelForMaskedImageModeling.from_pretrained(
										    model_args.model_name_or_path    ,      from_tf=bool('.ckpt' in model_args.model_name_or_path  )    ,      config=UpperCamelCase__    ,      cache_dir=model_args.cache_dir    ,      revision=model_args.model_revision    ,      use_auth_token=True if model_args.use_auth_token else None    ,      )
					else:
										logger.info('Training new model from scratch'  )
										_snake_case													=			AutoModelForMaskedImageModeling.from_config(UpperCamelCase__  )

					if training_args.do_train:
										_snake_case													=			ds['train'].column_names
					else:
										_snake_case													=			ds['validation'].column_names

					if data_args.image_column_name is not None:
										_snake_case													=			data_args.image_column_name
					elif "image" in column_names:
										_snake_case													=			'image'
					elif "img" in column_names:
										_snake_case													=			'img'
					else:
										_snake_case													=			column_names[0]

					# transformations as done in original SimMIM paper
					# source: https://github.com/microsoft/SimMIM/blob/main/data/data_simmim.py
					_snake_case													=			Compose(
					    [
					        Lambda(lambda UpperCamelCase__  : img.convert('RGB'  ) if img.mode != "RGB" else img  ),
					        RandomResizedCrop(model_args.image_size    ,      scale=(0.67, 1.0)    ,      ratio=(3.0 / 4.0, 4.0 / 3.0)  ),
					        RandomHorizontalFlip(),
					        ToTensor(),
					        Normalize(mean=image_processor.image_mean    ,      std=image_processor.image_std  ),
					    ]  )

					# create mask generator
					_snake_case													=			MaskGenerator(
					    input_size=model_args.image_size    ,      mask_patch_size=data_args.mask_patch_size    ,      model_patch_size=model_args.patch_size    ,      mask_ratio=data_args.mask_ratio    ,      )

					def preprocess_images(UpperCamelCase__     :    Dict  ):
										_snake_case													=			[transforms(UpperCamelCase__  ) for image in examples[image_column_name]]
										_snake_case													=			[mask_generator() for i in range(len(examples[image_column_name]  )  )]

										return examples

					if training_args.do_train:
										if "train" not in ds:
															raise ValueError('--do_train requires a train dataset'  )
										if data_args.max_train_samples is not None:
															_snake_case													=			ds['train'].shuffle(seed=training_args.seed  ).select(range(data_args.max_train_samples  )  )
										# Set the training transforms
										ds["train"].set_transform(UpperCamelCase__  )

					if training_args.do_eval:
										if "validation" not in ds:
															raise ValueError('--do_eval requires a validation dataset'  )
										if data_args.max_eval_samples is not None:
															_snake_case													=			(
															    ds['validation'].shuffle(seed=training_args.seed  ).select(range(data_args.max_eval_samples  )  )
															)
										# Set the validation transforms
										ds["validation"].set_transform(UpperCamelCase__  )

					# Initialize our trainer
					_snake_case													=			Trainer(
					    model=UpperCamelCase__    ,      args=UpperCamelCase__    ,      train_dataset=ds['train'] if training_args.do_train else None    ,      eval_dataset=ds['validation'] if training_args.do_eval else None    ,      tokenizer=UpperCamelCase__    ,      data_collator=UpperCamelCase__    ,      )

					# Training
					if training_args.do_train:
										_snake_case													=			None
										if training_args.resume_from_checkpoint is not None:
															_snake_case													=			training_args.resume_from_checkpoint
										elif last_checkpoint is not None:
															_snake_case													=			last_checkpoint
										_snake_case													=			trainer.train(resume_from_checkpoint=UpperCamelCase__  )
										trainer.save_model()
										trainer.log_metrics('train'    ,      train_result.metrics  )
										trainer.save_metrics('train'    ,      train_result.metrics  )
										trainer.save_state()

					# Evaluation
					if training_args.do_eval:
										_snake_case													=			trainer.evaluate()
										trainer.log_metrics('eval'    ,      UpperCamelCase__  )
										trainer.save_metrics('eval'    ,      UpperCamelCase__  )

					# Write model card and (optionally) push to hub
					_snake_case													=			{
					    'finetuned_from': model_args.model_name_or_path,
					    'tasks': 'masked-image-modeling',
					    'dataset': data_args.dataset_name,
					    'tags': ['masked-image-modeling'],
					}
					if training_args.push_to_hub:
										trainer.push_to_hub(**UpperCamelCase__  )
					else:
										trainer.create_model_card(**UpperCamelCase__  )


if __name__ == "__main__":
					main()










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__       (       UpperCamelCase__     :    int  )     ->	List[str]:



					'''simple docstring'''




					_snake_case													=			VideoMAEConfig()

					set_architecture_configs(UpperCamelCase__    ,      UpperCamelCase__  )

					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(UpperCamelCase__    ,      UpperCamelCase__    ,      repo_type='dataset'  )    ,      'r'  )  )
										_snake_case													=			{int(UpperCamelCase__  ): v for k, v in idalabel.items()}
										_snake_case													=			idalabel
										_snake_case													=			{v: k for k, v in idalabel.items()}

					return config





def   lowerCamelCase__       (       UpperCamelCase__     :    str    ,      UpperCamelCase__     :    Dict  )     ->	int:



					'''simple docstring'''




					if "small" in model_name:
										_snake_case													=			384
										_snake_case													=			1_536
										_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													=			1_024
										_snake_case													=			4_096
										_snake_case													=			24
										_snake_case													=			16
										_snake_case													=			12
										_snake_case													=			8
										_snake_case													=			512
										_snake_case													=			2_048
					elif "huge" in model_name:
										_snake_case													=			1_280
										_snake_case													=			5_120
										_snake_case													=			32
										_snake_case													=			16
										_snake_case													=			12
										_snake_case													=			8
										_snake_case													=			640
										_snake_case													=			2_560
					elif "base" not in model_name:
										raise ValueError('Model name should include either "small", "base", "large", or "huge"'  )





def   lowerCamelCase__       (       UpperCamelCase__     :    Any  )     ->	Tuple:



					'''simple docstring'''




					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   lowerCamelCase__       (       UpperCamelCase__     :    Any    ,      UpperCamelCase__     :    List[Any]  )     ->	Union[str, Any]:



					'''simple docstring'''




					for key in orig_state_dict.copy().keys():
										_snake_case													=			orig_state_dict.pop(UpperCamelCase__  )

										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   lowerCamelCase__       (       )     ->	Union[str, Any]:



					'''simple docstring'''




					_snake_case													=			hf_hub_download(
					    repo_id='hf-internal-testing/spaghetti-video'    ,      filename='eating_spaghetti.npy'    ,      repo_type='dataset'  )
					_snake_case													=			np.load(UpperCamelCase__  )
					return list(UpperCamelCase__  )





def   lowerCamelCase__       (       UpperCamelCase__     :    List[str]    ,      UpperCamelCase__     :    Union[str, Any]    ,      UpperCamelCase__     :    Optional[int]    ,      UpperCamelCase__     :    Optional[int]  )     ->	List[Any]:



					'''simple docstring'''




					_snake_case													=			get_videomae_config(UpperCamelCase__  )

					if "finetuned" in model_name:
										_snake_case													=			VideoMAEForVideoClassification(UpperCamelCase__  )
					else:
										_snake_case													=			VideoMAEForPreTraining(UpperCamelCase__  )

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

					model.load_state_dict(UpperCamelCase__  )
					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(UpperCamelCase__    ,      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(UpperCamelCase__  )

					_snake_case													=			model(**UpperCamelCase__  )
					_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, 1_408, 1_536]  )
										_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, 1_408, 1_536]  )
										_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, 1_408, 1_536]  )
										_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, 1_408, 1_536]  )
										_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, 1_408, 1_536]  )
										_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]    ,      UpperCamelCase__    ,      atol=1e-4  )
					else:
										print('Logits:'    ,      logits[0, :3, :3]  )
										assert torch.allclose(logits[0, :3, :3]    ,      UpperCamelCase__    ,      atol=1e-4  )
					print('Logits ok!'  )

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

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


if __name__ == "__main__":
					UpperCAmelCase_ 							=       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."""
					)

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



