



'''simple docstring'''



import argparse
import json

import gdown
import numpy as np
import torch
from huggingface_hub import hf_hub_download

from transformers import (
    VideoMAEConfig,
    VideoMAEForPreTraining,
    VideoMAEForVideoClassification,
    VideoMAEImageProcessor,
)



def   _A       (   _lowerCAmelCase					):



							"""simple docstring"""



							__lowercase       	=VideoMAEConfig()

							set_architecture_configs(__lowercase  ,   __lowercase					)

							if "finetuned" not in model_name:
														__lowercase       	=False

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

							return config



def   _A       (   _lowerCAmelCase  ,   _lowerCAmelCase					):



							"""simple docstring"""



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



def   _A       (   _lowerCAmelCase					):



							"""simple docstring"""



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

							return name



def   _A       (   _lowerCAmelCase  ,   _lowerCAmelCase					):



							"""simple docstring"""



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

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

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

							return orig_state_dict



def   _A       (   ):



							"""simple docstring"""



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



def   _A       (   _lowerCAmelCase  ,   _lowerCAmelCase  ,   _lowerCAmelCase  ,   _lowerCAmelCase					):



							"""simple docstring"""



							__lowercase       	=get_videomae_config(__lowercase					)

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

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

							model.load_state_dict(__lowercase					)
							model.eval()

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

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

							__lowercase       	=model(**__lowercase					)
							__lowercase       	=outputs.logits

							__lowercase       	=[
							    '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       	=torch.Size([1, 400]					)
														__lowercase       	=torch.tensor([-0.92_91, -0.40_61, -0.93_07]					)
							elif model_name == "videomae-small-finetuned-ssv2":
														__lowercase       	=torch.Size([1, 174]					)
														__lowercase       	=torch.tensor([0.26_71, -0.46_89, -0.82_35]					)
							elif model_name == "videomae-base":
														__lowercase       	=torch.Size([1, 1_408, 1_536]					)
														__lowercase       	=torch.tensor([[0.77_39, 0.79_68, 0.70_89], [0.67_01, 0.74_87, 0.62_09], [0.42_87, 0.51_58, 0.47_73]]					)
							elif model_name == "videomae-base-short":
														__lowercase       	=torch.Size([1, 1_408, 1_536]					)
														__lowercase       	=torch.tensor([[0.79_94, 0.96_12, 0.85_08], [0.74_01, 0.89_58, 0.83_02], [0.58_62, 0.74_68, 0.73_25]]					)
														# we verified the loss both for normalized and unnormalized targets for this one
														__lowercase       	=torch.tensor([0.51_42]					) if config.norm_pix_loss else torch.tensor([0.64_69]					)
							elif model_name == "videomae-large":
														__lowercase       	=torch.Size([1, 1_408, 1_536]					)
														__lowercase       	=torch.tensor([[0.71_49, 0.79_97, 0.69_66], [0.67_68, 0.78_69, 0.69_48], [0.51_39, 0.62_21, 0.56_05]]					)
							elif model_name == "videomae-large-finetuned-kinetics":
														__lowercase       	=torch.Size([1, 400]					)
														__lowercase       	=torch.tensor([0.07_71, 0.00_11, -0.36_25]					)
							elif model_name == "videomae-huge-finetuned-kinetics":
														__lowercase       	=torch.Size([1, 400]					)
														__lowercase       	=torch.tensor([0.24_33, 0.16_32, -0.48_94]					)
							elif model_name == "videomae-base-short-finetuned-kinetics":
														__lowercase       	=torch.Size([1, 400]					)
														__lowercase       	=torch.tensor([0.65_88, 0.09_90, -0.24_93]					)
							elif model_name == "videomae-base-finetuned-kinetics":
														__lowercase       	=torch.Size([1, 400]					)
														__lowercase       	=torch.tensor([0.36_69, -0.06_88, -0.24_21]					)
							elif model_name == "videomae-base-short-ssv2":
														__lowercase       	=torch.Size([1, 1_408, 1_536]					)
														__lowercase       	=torch.tensor([[0.47_12, 0.52_96, 0.57_86], [0.22_78, 0.27_29, 0.40_26], [0.03_52, 0.07_30, 0.25_06]]					)
							elif model_name == "videomae-base-short-finetuned-ssv2":
														__lowercase       	=torch.Size([1, 174]					)
														__lowercase       	=torch.tensor([-0.05_37, -0.15_39, -0.32_66]					)
							elif model_name == "videomae-base-ssv2":
														__lowercase       	=torch.Size([1, 1_408, 1_536]					)
														__lowercase       	=torch.tensor([[0.81_31, 0.87_27, 0.85_46], [0.73_66, 0.93_77, 0.88_70], [0.59_35, 0.88_74, 0.85_64]]					)
							elif model_name == "videomae-base-finetuned-ssv2":
														__lowercase       	=torch.Size([1, 174]					)
														__lowercase       	=torch.tensor([0.19_61, -0.83_37, -0.63_89]					)
							else:
														raise ValueError(f"""Model name not supported. Should be one of {model_names}"""					)

							# verify logits
							assert logits.shape == expected_shape
							if "finetuned" in model_name:
														assert torch.allclose(logits[0, :3]  ,   __lowercase  ,   atol=1e-4					)
							else:
														print('Logits:'  ,   logits[0, :3, :3]					)
														assert torch.allclose(logits[0, :3, :3]  ,   __lowercase  ,   atol=1e-4					)
							print('Logits ok!'					)

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

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


if __name__ == "__main__":
		lowerCamelCase	 =    argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    """--checkpoint_url""",
		    default="""https://drive.google.com/u/1/uc?id=1tEhLyskjb755TJ65ptsrafUG2llSwQE1&amp;export=download&amp;confirm=t&amp;uuid=aa3276eb-fb7e-482a-adec-dc7171df14c4""",
		    type=str,
		    help=(
		        """URL of the original PyTorch checkpoint (on Google Drive) you'd like to convert. Should be a direct"""
		        """ download link."""
		    ),
		)
		parser.add_argument(
		    """--pytorch_dump_folder_path""",
		    default="""/Users/nielsrogge/Documents/VideoMAE/Test""",
		    type=str,
		    help="""Path to the output PyTorch model directory.""",
		)
		parser.add_argument("""--model_name""", default="""videomae-base""", type=str, help="""Name of the model.""")
		parser.add_argument(
		    """--push_to_hub""", action="""store_true""", help="""Whether or not to push the converted model to the 🤗 hub."""
		)

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








'''simple docstring'''



import argparse
import os

import evaluate
import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType
from accelerate.local_sgd import LocalSGD


########################################################################
# This is a fully working simple example to use Accelerate
# with LocalSGD, which is a method to synchronize model
# parameters every K batches. It is different, but complementary
# to gradient accumulation.
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################


lowerCamelCase	 =       16
lowerCamelCase	 =       32



def   _A       (   _lowerCAmelCase  ,   _lowerCAmelCase = 16					):



							"""simple docstring"""



							__lowercase       	=AutoTokenizer.from_pretrained('bert-base-cased'					)
							__lowercase       	=load_dataset('glue'  ,   'mrpc'					)

							def tokenize_function(_lowerCAmelCase					):
														# max_length=None => use the model max length (it's actually the default)
														__lowercase       	=tokenizer(examples['sentence1']  ,   examples['sentence2']  ,   truncation=_lowerCAmelCase  ,   max_length=_lowerCAmelCase					)
														return outputs

							# Apply the method we just defined to all the examples in all the splits of the dataset
							# starting with the main process first:
							with accelerator.main_process_first():
														__lowercase       	=datasets.map(
														    _lowerCAmelCase  ,   batched=_lowerCAmelCase  ,   remove_columns=['idx', 'sentence1', 'sentence2']  ,   )

							# We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
							# transformers library
							__lowercase       	=tokenized_datasets.rename_column('label'  ,   'labels'					)

							def collate_fn(_lowerCAmelCase					):
														# On TPU it's best to pad everything to the same length or training will be very slow.
														__lowercase       	=128 if accelerator.distributed_type == DistributedType.TPU else None
														# When using mixed precision we want round multiples of 8/16
														if accelerator.mixed_precision == "fp8":
																					__lowercase       	=16
														elif accelerator.mixed_precision != "no":
																					__lowercase       	=8
														else:
																					__lowercase       	=None

														return tokenizer.pad(
														    _lowerCAmelCase  ,   padding='longest'  ,   max_length=_lowerCAmelCase  ,   pad_to_multiple_of=_lowerCAmelCase  ,   return_tensors='pt'  ,   )

							# Instantiate dataloaders.
							__lowercase       	=DataLoader(
							    tokenized_datasets['train']  ,   shuffle=_lowerCAmelCase  ,   collate_fn=_lowerCAmelCase  ,   batch_size=_lowerCAmelCase					)
							__lowercase       	=DataLoader(
							    tokenized_datasets['validation']  ,   shuffle=_lowerCAmelCase  ,   collate_fn=_lowerCAmelCase  ,   batch_size=_lowerCAmelCase					)

							return train_dataloader, eval_dataloader


# For testing only
if os.environ.get("""TESTING_MOCKED_DATALOADERS""", None) == "1":
		from accelerate.test_utils.training import mocked_dataloaders

		lowerCamelCase	 =       mocked_dataloaders  # noqa: F811



def   _A       (   _lowerCAmelCase  ,   _lowerCAmelCase					):



							"""simple docstring"""



							if os.environ.get('TESTING_MOCKED_DATALOADERS'  ,   _lowerCAmelCase					) == "1":
														__lowercase       	=2
							# New Code #
							__lowercase       	=int(args.gradient_accumulation_steps					)
							__lowercase       	=int(args.local_sgd_steps					)
							# Initialize accelerator
							__lowercase       	=Accelerator(
							    cpu=args.cpu  ,   mixed_precision=args.mixed_precision  ,   gradient_accumulation_steps=_lowerCAmelCase					)
							if accelerator.distributed_type not in [DistributedType.NO, DistributedType.MULTI_CPU, DistributedType.MULTI_GPU]:
														raise NotImplementedError('LocalSGD is supported only for CPUs and GPUs (no DeepSpeed or MegatronLM)'					)
							# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
							__lowercase       	=config['lr']
							__lowercase       	=int(config['num_epochs']					)
							__lowercase       	=int(config['seed']					)
							__lowercase       	=int(config['batch_size']					)

							__lowercase       	=evaluate.load('glue'  ,   'mrpc'					)

							set_seed(_lowerCAmelCase					)
							__lowercase	,				__lowercase       	=get_dataloaders(_lowerCAmelCase  ,   _lowerCAmelCase					)
							# Instantiate the model (we build the model here so that the seed also control new weights initialization)
							__lowercase       	=AutoModelForSequenceClassification.from_pretrained('bert-base-cased'  ,   return_dict=_lowerCAmelCase					)

							# We could avoid this line since the accelerator is set with `device_placement=True` (default value).
							# Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
							# creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
							__lowercase       	=model.to(accelerator.device					)

							# Instantiate optimizer
							__lowercase       	=AdamW(params=model.parameters()  ,   lr=_lowerCAmelCase					)

							# Instantiate scheduler
							__lowercase       	=get_linear_schedule_with_warmup(
							    optimizer=_lowerCAmelCase  ,   num_warmup_steps=100  ,   num_training_steps=(len(_lowerCAmelCase					) * num_epochs)  ,   )

							# Prepare everything
							# There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
							# prepare method.
							__lowercase	,				__lowercase	,				__lowercase	,				__lowercase	,				__lowercase       	=accelerator.prepare(
							    _lowerCAmelCase  ,   _lowerCAmelCase  ,   _lowerCAmelCase  ,   _lowerCAmelCase  ,   _lowerCAmelCase					)

							# Now we train the model
							for epoch in range(_lowerCAmelCase					):
														model.train()
														with LocalSGD(
														    accelerator=_lowerCAmelCase  ,   model=_lowerCAmelCase  ,   local_sgd_steps=_lowerCAmelCase  ,   enabled=local_sgd_steps is not None					) as local_sgd:
																					for step, batch in enumerate(_lowerCAmelCase					):
																												# We could avoid this line since we set the accelerator with `device_placement=True`.
																												batch.to(accelerator.device					)
																												# New code #
																												# We use the new `accumulate` context manager to perform gradient accumulation
																												# We also currently do not support TPUs nor advise it as bugs were found on the XLA side when running our tests.
																												with accelerator.accumulate(_lowerCAmelCase					):
																																			__lowercase       	=model(**_lowerCAmelCase					)
																																			__lowercase       	=output.loss
																																			accelerator.backward(_lowerCAmelCase					)
																																			optimizer.step()
																																			lr_scheduler.step()
																																			optimizer.zero_grad()
																																			# LocalSGD-specific line
																																			local_sgd.step()

														model.eval()
														for step, batch in enumerate(_lowerCAmelCase					):
																					# We could avoid this line since we set the accelerator with `device_placement=True`.
																					batch.to(accelerator.device					)
																					with torch.no_grad():
																												__lowercase       	=model(**_lowerCAmelCase					)
																					__lowercase       	=outputs.logits.argmax(dim=-1					)
																					__lowercase	,				__lowercase       	=accelerator.gather_for_metrics((predictions, batch['labels'])					)
																					metric.add_batch(
																					    predictions=_lowerCAmelCase  ,   references=_lowerCAmelCase  ,   )

														__lowercase       	=metric.compute()
														# Use accelerator.print to print only on the main process.
														accelerator.print(f"""epoch {epoch}:"""  ,   _lowerCAmelCase					)



def   _A       (   ):



							"""simple docstring"""



							__lowercase       	=argparse.ArgumentParser(description='Simple example of training script.'					)
							parser.add_argument(
							    '--mixed_precision'  ,   type=_lowerCAmelCase  ,   default=_lowerCAmelCase  ,   choices=['no', 'fp16', 'bf16', 'fp8']  ,   help='Whether to use mixed precision. Choose'
							    'between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.'
							    'and an Nvidia Ampere GPU.'  ,   )
							# New Code #
							parser.add_argument(
							    '--gradient_accumulation_steps'  ,   type=_lowerCAmelCase  ,   default=1  ,   help='The number of minibatches to be ran before gradients are accumulated.'  ,   )
							parser.add_argument(
							    '--local_sgd_steps'  ,   type=_lowerCAmelCase  ,   default=8  ,   help='Number of local SGD steps or None to disable local SGD'					)
							parser.add_argument('--cpu'  ,   action='store_true'  ,   help='If passed, will train on the CPU.'					)
							__lowercase       	=parser.parse_args()
							__lowercase       	={'lr': 2e-5, 'num_epochs': 3, 'seed': 42, 'batch_size': 16}
							training_function(_lowerCAmelCase  ,   _lowerCAmelCase					)


if __name__ == "__main__":
		main()



