




"""simple docstring"""





import argparse
import json

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

from transformers import TimesformerConfig, TimesformerForVideoClassification, VideoMAEImageProcessor

def 						lowerCAmelCase							(	lowerCamelCase						):
						'''simple docstring'''


						snake_case__    :   Union[str, Any]							     =   TimesformerConfig()

						if "large" in model_name:
												snake_case__    :   List[Any]							     =   96

						if "hr" in model_name:
												snake_case__    :   Dict							     =   16
												snake_case__    :   Dict							     =   448

						snake_case__    :   Dict							     =   '''huggingface/label-files'''
						if "k400" in model_name:
												snake_case__    :   List[str]							     =   400
												snake_case__    :   str							     =   '''kinetics400-id2label.json'''
						elif "k600" in model_name:
												snake_case__    :   Optional[Any]							     =   600
												snake_case__    :   List[str]							     =   '''kinetics600-id2label.json'''
						elif "ssv2" in model_name:
												snake_case__    :   Optional[int]							     =   174
												snake_case__    :   Tuple							     =   '''something-something-v2-id2label.json'''
						else:
												raise ValueError('''Model name should either contain \'k400\', \'k600\' or \'ssv2\'.'''						)
						snake_case__    :   List[str]							     =   json.load(open(hf_hub_download(lowerCamelCase ,     lowerCamelCase ,     repo_type='''dataset'''						) ,     '''r'''						)						)
						snake_case__    :   Optional[int]							     =   {int(lowerCamelCase						): v for k, v in idalabel.items()}
						snake_case__    :   List[str]							     =   idalabel
						snake_case__    :   Optional[int]							     =   {v: k for k, v in idalabel.items()}

						return config

def 						lowerCAmelCase							(	lowerCamelCase						):
						'''simple docstring'''


						if "encoder." in name:
												snake_case__    :   Dict							     =   name.replace('''encoder.''' ,     ''''''						)
						if "cls_token" in name:
												snake_case__    :   Tuple							     =   name.replace('''cls_token''' ,     '''timesformer.embeddings.cls_token'''						)
						if "pos_embed" in name:
												snake_case__    :   Tuple							     =   name.replace('''pos_embed''' ,     '''timesformer.embeddings.position_embeddings'''						)
						if "time_embed" in name:
												snake_case__    :   Optional[int]							     =   name.replace('''time_embed''' ,     '''timesformer.embeddings.time_embeddings'''						)
						if "patch_embed.proj" in name:
												snake_case__    :   Optional[int]							     =   name.replace('''patch_embed.proj''' ,     '''timesformer.embeddings.patch_embeddings.projection'''						)
						if "patch_embed.norm" in name:
												snake_case__    :   List[str]							     =   name.replace('''patch_embed.norm''' ,     '''timesformer.embeddings.norm'''						)
						if "blocks" in name:
												snake_case__    :   Dict							     =   name.replace('''blocks''' ,     '''timesformer.encoder.layer'''						)
						if "attn.proj" in name:
												snake_case__    :   int							     =   name.replace('''attn.proj''' ,     '''attention.output.dense'''						)
						if "attn" in name and "bias" not in name and "temporal" not in name:
												snake_case__    :   str							     =   name.replace('''attn''' ,     '''attention.self'''						)
						if "attn" in name and "temporal" not in name:
												snake_case__    :   List[Any]							     =   name.replace('''attn''' ,     '''attention.attention'''						)
						if "temporal_norm1" in name:
												snake_case__    :   Any							     =   name.replace('''temporal_norm1''' ,     '''temporal_layernorm'''						)
						if "temporal_attn.proj" in name:
												snake_case__    :   Any							     =   name.replace('''temporal_attn''' ,     '''temporal_attention.output.dense'''						)
						if "temporal_fc" in name:
												snake_case__    :   Tuple							     =   name.replace('''temporal_fc''' ,     '''temporal_dense'''						)
						if "norm1" in name and "temporal" not in name:
												snake_case__    :   List[Any]							     =   name.replace('''norm1''' ,     '''layernorm_before'''						)
						if "norm2" in name:
												snake_case__    :   str							     =   name.replace('''norm2''' ,     '''layernorm_after'''						)
						if "mlp.fc1" in name:
												snake_case__    :   List[Any]							     =   name.replace('''mlp.fc1''' ,     '''intermediate.dense'''						)
						if "mlp.fc2" in name:
												snake_case__    :   List[Any]							     =   name.replace('''mlp.fc2''' ,     '''output.dense'''						)
						if "norm.weight" in name and "fc" not in name and "temporal" not in name:
												snake_case__    :   int							     =   name.replace('''norm.weight''' ,     '''timesformer.layernorm.weight'''						)
						if "norm.bias" in name and "fc" not in name and "temporal" not in name:
												snake_case__    :   Union[str, Any]							     =   name.replace('''norm.bias''' ,     '''timesformer.layernorm.bias'''						)
						if "head" in name:
												snake_case__    :   Dict							     =   name.replace('''head''' ,     '''classifier'''						)

						return name

def 						lowerCAmelCase							(	lowerCamelCase ,     lowerCamelCase						):
						'''simple docstring'''


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

												if key.startswith('''model.'''						):
																		snake_case__    :   Dict							     =   key.replace('''model.''' ,     ''''''						)

												if "qkv" in key:
																		snake_case__    :   Optional[int]							     =   key.split('''.'''						)
																		snake_case__    :   Any							     =   int(key_split[1]						)
																		snake_case__    :   Union[str, Any]							     =   '''timesformer.encoder.layer.'''
																		if "temporal" in key:
																								snake_case__    :   Any							     =   '''.temporal_attention.attention.qkv.'''
																		else:
																								snake_case__    :   int							     =   '''.attention.attention.qkv.'''
																		if "weight" in key:
																								snake_case__    :   Tuple							     =   val
																		else:
																								snake_case__    :   Optional[Any]							     =   val
												else:
																		snake_case__    :   Optional[int]							     =   val

						return orig_state_dict

def 						lowerCAmelCase							(	):
						'''simple docstring'''


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

def 						lowerCAmelCase							(	lowerCamelCase ,     lowerCamelCase ,     lowerCamelCase ,     lowerCamelCase						):
						'''simple docstring'''


						snake_case__    :   List[Any]							     =   get_timesformer_config(lowerCamelCase						)

						snake_case__    :   List[Any]							     =   TimesformerForVideoClassification(lowerCamelCase						)

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

						model.load_state_dict(lowerCamelCase						)
						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__    :   int							     =   prepare_video()
						snake_case__    :   Optional[int]							     =   image_processor(video[:8] ,     return_tensors='''pt'''						)

						snake_case__    :   Optional[int]							     =   model(**lowerCamelCase						)
						snake_case__    :   Dict							     =   outputs.logits

						snake_case__    :   Optional[int]							     =   [
						    # Kinetics-400 checkpoints (hr = high resolution input of 448px instead of 224px)
						    '''timesformer-base-finetuned-k400''',
						    '''timesformer-large-finetuned-k400''',
						    '''timesformer-hr-finetuned-k400''',
						    # Kinetics-600 checkpoints (hr = high resolution input of 448px instead of 224px)
						    '''timesformer-base-finetuned-k600''',
						    '''timesformer-large-finetuned-k600''',
						    '''timesformer-hr-finetuned-k600''',
						    # Something-Something-v2 checkpoints (hr = high resolution input of 448px instead of 224px)
						    '''timesformer-base-finetuned-ssv2''',
						    '''timesformer-large-finetuned-ssv2''',
						    '''timesformer-hr-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 == "timesformer-base-finetuned-k400":
												snake_case__    :   Optional[Any]							     =   torch.Size([1, 400]						)
												snake_case__    :   List[str]							     =   torch.tensor([-0.3_016, -0.7_713, -0.4_205]						)
						elif model_name == "timesformer-base-finetuned-k600":
												snake_case__    :   Any							     =   torch.Size([1, 600]						)
												snake_case__    :   int							     =   torch.tensor([-0.7_267, -0.7_466, 3.2_404]						)
						elif model_name == "timesformer-base-finetuned-ssv2":
												snake_case__    :   List[Any]							     =   torch.Size([1, 174]						)
												snake_case__    :   List[str]							     =   torch.tensor([-0.9_059, 0.6_433, -3.1_457]						)
						elif model_name == "timesformer-large-finetuned-k400":
												snake_case__    :   List[str]							     =   torch.Size([1, 400]						)
												snake_case__    :   Optional[Any]							     =   torch.tensor([0, 0, 0]						)
						elif model_name == "timesformer-large-finetuned-k600":
												snake_case__    :   Dict							     =   torch.Size([1, 600]						)
												snake_case__    :   str							     =   torch.tensor([0, 0, 0]						)
						elif model_name == "timesformer-large-finetuned-ssv2":
												snake_case__    :   Tuple							     =   torch.Size([1, 174]						)
												snake_case__    :   Tuple							     =   torch.tensor([0, 0, 0]						)
						elif model_name == "timesformer-hr-finetuned-k400":
												snake_case__    :   int							     =   torch.Size([1, 400]						)
												snake_case__    :   List[str]							     =   torch.tensor([-0.9_617, -3.7_311, -3.7_708]						)
						elif model_name == "timesformer-hr-finetuned-k600":
												snake_case__    :   Tuple							     =   torch.Size([1, 600]						)
												snake_case__    :   Dict							     =   torch.tensor([2.5_273, 0.7_127, 1.8_848]						)
						elif model_name == "timesformer-hr-finetuned-ssv2":
												snake_case__    :   List[Any]							     =   torch.Size([1, 174]						)
												snake_case__    :   Tuple							     =   torch.tensor([-3.6_756, -0.7_513, 0.7_180]						)
						else:
												raise ValueError(f'''Model name not supported. Should be one of {model_names}'''						)

						# verify logits
						assert logits.shape == expected_shape
						assert torch.allclose(logits[0, :3] ,     lowerCamelCase ,     atol=1e-4						)
						print('''Logits ok!'''						)

						if pytorch_dump_folder_path is not None:
												print(f'''Saving model and image processor to {pytorch_dump_folder_path}'''						)
												image_processor.save_pretrained(lowerCamelCase						)
												model.save_pretrained(lowerCamelCase						)

						if push_to_hub:
												print('''Pushing to the hub...'''						)
												model.push_to_hub(f'''fcakyon/{model_name}'''						)


if __name__ == "__main__":
							_lowercase		 =      argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    """--checkpoint_url""",
							    default="""https://drive.google.com/u/1/uc?id=17yvuYp9L4mn-HpIcK5Zo6K3UoOy1kA5l&export=download""",
							    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="""""",
							    type=str,
							    help="""Path to the output PyTorch model directory.""",
							)
							parser.add_argument("""--model_name""", default="""timesformer-base-finetuned-k400""", 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."""
							)

							_lowercase		 =      parser.parse_args()
							convert_timesformer_checkpoint(
							    args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub
							)





"""simple docstring"""





import argparse

import torch

from transformers import MraConfig, MraForMaskedLM

def 						lowerCAmelCase							(	lowerCamelCase						):
						'''simple docstring'''


						if "model" in orig_key:
												snake_case__    :   Union[str, Any]							     =   orig_key.replace('''model.''' ,     ''''''						)
						if "norm1" in orig_key:
												snake_case__    :   Union[str, Any]							     =   orig_key.replace('''norm1''' ,     '''attention.output.LayerNorm'''						)
						if "norm2" in orig_key:
												snake_case__    :   Dict							     =   orig_key.replace('''norm2''' ,     '''output.LayerNorm'''						)
						if "norm" in orig_key:
												snake_case__    :   int							     =   orig_key.replace('''norm''' ,     '''LayerNorm'''						)
						if "transformer" in orig_key:
												snake_case__    :   Optional[Any]							     =   orig_key.split('''.'''						)[0].split('''_'''						)[-1]
												snake_case__    :   Optional[Any]							     =   orig_key.replace(f'''transformer_{layer_num}''' ,     f'''encoder.layer.{layer_num}'''						)
						if "mha.attn" in orig_key:
												snake_case__    :   int							     =   orig_key.replace('''mha.attn''' ,     '''attention.self'''						)
						if "mha" in orig_key:
												snake_case__    :   Union[str, Any]							     =   orig_key.replace('''mha''' ,     '''attention'''						)
						if "W_q" in orig_key:
												snake_case__    :   Optional[int]							     =   orig_key.replace('''W_q''' ,     '''self.query'''						)
						if "W_k" in orig_key:
												snake_case__    :   Tuple							     =   orig_key.replace('''W_k''' ,     '''self.key'''						)
						if "W_v" in orig_key:
												snake_case__    :   Dict							     =   orig_key.replace('''W_v''' ,     '''self.value'''						)
						if "ff.0" in orig_key:
												snake_case__    :   str							     =   orig_key.replace('''ff.0''' ,     '''intermediate.dense'''						)
						if "ff.2" in orig_key:
												snake_case__    :   int							     =   orig_key.replace('''ff.2''' ,     '''output.dense'''						)
						if "ff" in orig_key:
												snake_case__    :   Union[str, Any]							     =   orig_key.replace('''ff''' ,     '''output.dense'''						)
						if "mlm_class" in orig_key:
												snake_case__    :   str							     =   orig_key.replace('''mlm.mlm_class''' ,     '''cls.predictions.decoder'''						)
						if "mlm" in orig_key:
												snake_case__    :   List[Any]							     =   orig_key.replace('''mlm''' ,     '''cls.predictions.transform'''						)
						if "backbone.backbone.encoders" in orig_key:
												snake_case__    :   Optional[Any]							     =   orig_key.replace('''backbone.backbone.encoders''' ,     '''encoder.layer'''						)
						if "cls" not in orig_key:
												snake_case__    :   Union[str, Any]							     =   '''mra.''' + orig_key

						return orig_key

def 						lowerCAmelCase							(	lowerCamelCase ,     lowerCamelCase						):
						'''simple docstring'''


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

												if ("pooler" in key) or ("sen_class" in key):
																		continue
												else:
																		snake_case__    :   Optional[int]							     =   val

						snake_case__    :   Union[str, Any]							     =   orig_state_dict['''cls.predictions.decoder.bias''']
						snake_case__    :   Any							     =   torch.arange(lowerCamelCase						).expand((1, -1)						) + 2

						return orig_state_dict

def 						lowerCAmelCase							(	lowerCamelCase ,     lowerCamelCase ,     lowerCamelCase						):
						'''simple docstring'''


						snake_case__    :   List[Any]							     =   torch.load(lowerCamelCase ,     map_location='''cpu'''						)['''model_state_dict''']
						snake_case__    :   Optional[Any]							     =   MraConfig.from_json_file(lowerCamelCase						)
						snake_case__    :   List[Any]							     =   MraForMaskedLM(lowerCamelCase						)

						snake_case__    :   List[Any]							     =   convert_checkpoint_helper(config.max_position_embeddings ,     lowerCamelCase						)

						print(model.load_state_dict(lowerCamelCase						)						)
						model.eval()
						model.save_pretrained(lowerCamelCase						)

						print(f'''Checkpoint successfuly converted. Model saved at {pytorch_dump_path}'''						)


if __name__ == "__main__":
							_lowercase		 =      argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    """--pytorch_model_path""", default=None, type=str, required=True, help="""Path to Mra pytorch checkpoint."""
							)
							parser.add_argument(
							    """--config_file""",
							    default=None,
							    type=str,
							    required=True,
							    help="""The json file for Mra model config.""",
							)
							parser.add_argument(
							    """--pytorch_dump_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
							)
							_lowercase		 =      parser.parse_args()
							convert_mra_checkpoint(args.pytorch_model_path, args.config_file, args.pytorch_dump_path)


