




'''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   lowercase__(			__UpperCamelCase:					Optional[int]    ):

      """simple docstring"""






      SCREAMING_SNAKE_CASE   :      Any    					=						VideoMAEConfig()

      set_architecture_configs(__UpperCamelCase      ,__UpperCamelCase    )

      if "finetuned" not in model_name:
            SCREAMING_SNAKE_CASE   :      Optional[int]    					=						False

      if "finetuned" in model_name:
            SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						'huggingface/label-files'
            if "kinetics" in model_name:
                  SCREAMING_SNAKE_CASE   :      Tuple    					=						4_00
                  SCREAMING_SNAKE_CASE   :      str    					=						'kinetics400-id2label.json'
            elif "ssv2" in model_name:
                  SCREAMING_SNAKE_CASE   :      Union[str, Any]    					=						1_74
                  SCREAMING_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.'    )
            SCREAMING_SNAKE_CASE   :      Union[str, Any]    					=						json.load(open(hf_hub_download(__UpperCamelCase      ,__UpperCamelCase      ,repo_type='dataset'    )      ,'r'    )    )
            SCREAMING_SNAKE_CASE   :      Any    					=						{int(__UpperCamelCase    ): v for k, v in idalabel.items()}
            SCREAMING_SNAKE_CASE   :      Optional[int]    					=						idalabel
            SCREAMING_SNAKE_CASE   :      str    					=						{v: k for k, v in idalabel.items()}

      return config




def   lowercase__(			__UpperCamelCase:					List[str]      ,__UpperCamelCase:					Any    ):

      """simple docstring"""






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




def   lowercase__(			__UpperCamelCase:					str    ):

      """simple docstring"""






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

      return name




def   lowercase__(			__UpperCamelCase:					str      ,__UpperCamelCase:					str    ):

      """simple docstring"""






      for key in orig_state_dict.copy().keys():
            SCREAMING_SNAKE_CASE   :      int    					=						orig_state_dict.pop(__UpperCamelCase    )

            if key.startswith('encoder.'    ):
                  SCREAMING_SNAKE_CASE   :      List[str]    					=						key.replace('encoder.'      ,''    )

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

      return orig_state_dict




def   lowercase__(			):

      """simple docstring"""






      SCREAMING_SNAKE_CASE   :      Any    					=						hf_hub_download(
          repo_id='hf-internal-testing/spaghetti-video'      ,filename='eating_spaghetti.npy'      ,repo_type='dataset'    )
      SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						np.load(__UpperCamelCase    )
      return list(__UpperCamelCase    )




def   lowercase__(			__UpperCamelCase:					Optional[int]      ,__UpperCamelCase:					List[str]      ,__UpperCamelCase:					Union[str, Any]      ,__UpperCamelCase:					Dict    ):

      """simple docstring"""






      SCREAMING_SNAKE_CASE   :      Tuple    					=						get_videomae_config(__UpperCamelCase    )

      if "finetuned" in model_name:
            SCREAMING_SNAKE_CASE   :      Tuple    					=						VideoMAEForVideoClassification(__UpperCamelCase    )
      else:
            SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						VideoMAEForPreTraining(__UpperCamelCase    )

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

      model.load_state_dict(__UpperCamelCase    )
      model.eval()

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

      if "finetuned" not in model_name:
            SCREAMING_SNAKE_CASE   :      int    					=						hf_hub_download(repo_id='hf-internal-testing/bool-masked-pos'      ,filename='bool_masked_pos.pt'    )
            SCREAMING_SNAKE_CASE   :      List[Any]    					=						torch.load(__UpperCamelCase    )

      SCREAMING_SNAKE_CASE   :      int    					=						model(**__UpperCamelCase    )
      SCREAMING_SNAKE_CASE   :      Optional[int]    					=						outputs.logits

      SCREAMING_SNAKE_CASE   :      int    					=						[
          '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":
            SCREAMING_SNAKE_CASE   :      Optional[int]    					=						torch.Size([1, 4_00]    )
            SCREAMING_SNAKE_CASE   :      int    					=						torch.tensor([-0.9_2_9_1, -0.4_0_6_1, -0.9_3_0_7]    )
      elif model_name == "videomae-small-finetuned-ssv2":
            SCREAMING_SNAKE_CASE   :      Any    					=						torch.Size([1, 1_74]    )
            SCREAMING_SNAKE_CASE   :      str    					=						torch.tensor([0.2_6_7_1, -0.4_6_8_9, -0.8_2_3_5]    )
      elif model_name == "videomae-base":
            SCREAMING_SNAKE_CASE   :      Tuple    					=						torch.Size([1, 14_08, 15_36]    )
            SCREAMING_SNAKE_CASE   :      Any    					=						torch.tensor([[0.7_7_3_9, 0.7_9_6_8, 0.7_0_8_9], [0.6_7_0_1, 0.7_4_8_7, 0.6_2_0_9], [0.4_2_8_7, 0.5_1_5_8, 0.4_7_7_3]]    )
      elif model_name == "videomae-base-short":
            SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						torch.Size([1, 14_08, 15_36]    )
            SCREAMING_SNAKE_CASE   :      List[Any]    					=						torch.tensor([[0.7_9_9_4, 0.9_6_1_2, 0.8_5_0_8], [0.7_4_0_1, 0.8_9_5_8, 0.8_3_0_2], [0.5_8_6_2, 0.7_4_6_8, 0.7_3_2_5]]    )
            # we verified the loss both for normalized and unnormalized targets for this one
            SCREAMING_SNAKE_CASE   :      str    					=						torch.tensor([0.5_1_4_2]    ) if config.norm_pix_loss else torch.tensor([0.6_4_6_9]    )
      elif model_name == "videomae-large":
            SCREAMING_SNAKE_CASE   :      List[str]    					=						torch.Size([1, 14_08, 15_36]    )
            SCREAMING_SNAKE_CASE   :      Optional[int]    					=						torch.tensor([[0.7_1_4_9, 0.7_9_9_7, 0.6_9_6_6], [0.6_7_6_8, 0.7_8_6_9, 0.6_9_4_8], [0.5_1_3_9, 0.6_2_2_1, 0.5_6_0_5]]    )
      elif model_name == "videomae-large-finetuned-kinetics":
            SCREAMING_SNAKE_CASE   :      List[str]    					=						torch.Size([1, 4_00]    )
            SCREAMING_SNAKE_CASE   :      List[Any]    					=						torch.tensor([0.0_7_7_1, 0.0_0_1_1, -0.3_6_2_5]    )
      elif model_name == "videomae-huge-finetuned-kinetics":
            SCREAMING_SNAKE_CASE   :      int    					=						torch.Size([1, 4_00]    )
            SCREAMING_SNAKE_CASE   :      List[str]    					=						torch.tensor([0.2_4_3_3, 0.1_6_3_2, -0.4_8_9_4]    )
      elif model_name == "videomae-base-short-finetuned-kinetics":
            SCREAMING_SNAKE_CASE   :      Any    					=						torch.Size([1, 4_00]    )
            SCREAMING_SNAKE_CASE   :      List[str]    					=						torch.tensor([0.6_5_8_8, 0.0_9_9_0, -0.2_4_9_3]    )
      elif model_name == "videomae-base-finetuned-kinetics":
            SCREAMING_SNAKE_CASE   :      Tuple    					=						torch.Size([1, 4_00]    )
            SCREAMING_SNAKE_CASE   :      List[Any]    					=						torch.tensor([0.3_6_6_9, -0.0_6_8_8, -0.2_4_2_1]    )
      elif model_name == "videomae-base-short-ssv2":
            SCREAMING_SNAKE_CASE   :      Tuple    					=						torch.Size([1, 14_08, 15_36]    )
            SCREAMING_SNAKE_CASE   :      Dict    					=						torch.tensor([[0.4_7_1_2, 0.5_2_9_6, 0.5_7_8_6], [0.2_2_7_8, 0.2_7_2_9, 0.4_0_2_6], [0.0_3_5_2, 0.0_7_3_0, 0.2_5_0_6]]    )
      elif model_name == "videomae-base-short-finetuned-ssv2":
            SCREAMING_SNAKE_CASE   :      Any    					=						torch.Size([1, 1_74]    )
            SCREAMING_SNAKE_CASE   :      Any    					=						torch.tensor([-0.0_5_3_7, -0.1_5_3_9, -0.3_2_6_6]    )
      elif model_name == "videomae-base-ssv2":
            SCREAMING_SNAKE_CASE   :      Optional[int]    					=						torch.Size([1, 14_08, 15_36]    )
            SCREAMING_SNAKE_CASE   :      str    					=						torch.tensor([[0.8_1_3_1, 0.8_7_2_7, 0.8_5_4_6], [0.7_3_6_6, 0.9_3_7_7, 0.8_8_7_0], [0.5_9_3_5, 0.8_8_7_4, 0.8_5_6_4]]    )
      elif model_name == "videomae-base-finetuned-ssv2":
            SCREAMING_SNAKE_CASE   :      List[str]    					=						torch.Size([1, 1_74]    )
            SCREAMING_SNAKE_CASE   :      int    					=						torch.tensor([0.1_9_6_1, -0.8_3_3_7, -0.6_3_8_9]    )
      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":
            SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						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)




"""simple docstring"""
from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


SCREAMING_SNAKE_CASE_					     =	{
    '''configuration_rag''': ['''RagConfig'''],
    '''retrieval_rag''': ['''RagRetriever'''],
    '''tokenization_rag''': ['''RagTokenizer'''],
}

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			SCREAMING_SNAKE_CASE_					     =	[
			    '''RagModel''',
			    '''RagPreTrainedModel''',
			    '''RagSequenceForGeneration''',
			    '''RagTokenForGeneration''',
			]

try:
			if not is_tf_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			SCREAMING_SNAKE_CASE_					     =	[
			    '''TFRagModel''',
			    '''TFRagPreTrainedModel''',
			    '''TFRagSequenceForGeneration''',
			    '''TFRagTokenForGeneration''',
			]


if TYPE_CHECKING:
			from .configuration_rag import RagConfig
			from .retrieval_rag import RagRetriever
			from .tokenization_rag import RagTokenizer

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_rag import RagModel, RagPreTrainedModel, RagSequenceForGeneration, RagTokenForGeneration

			try:
						if not is_tf_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_tf_rag import (
						    TFRagModel,
						    TFRagPreTrainedModel,
						    TFRagSequenceForGeneration,
						    TFRagTokenForGeneration,
						)

else:
			import sys

			SCREAMING_SNAKE_CASE_					     =	_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)
