


'''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__						)    ->	Optional[int]:
    UpperCAmelCase__			:		List[str]													=				VideoMAEConfig()

    set_architecture_configs(lowerCAmelCase__  ,					lowerCAmelCase__						)

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

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

    return config





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





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

    return name





def 					a__       (					lowerCAmelCase__  ,					lowerCAmelCase__						)    ->	Optional[int]:
    for key in orig_state_dict.copy().keys():
        UpperCAmelCase__			:		Any													=				orig_state_dict.pop(lowerCAmelCase__						)

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

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

    return orig_state_dict





def 					a__       (					)    ->	Any:
    UpperCAmelCase__			:		List[Any]													=				hf_hub_download(
        repo_id='''hf-internal-testing/spaghetti-video'''  ,					filename='''eating_spaghetti.npy'''  ,					repo_type='''dataset'''						)
    UpperCAmelCase__			:		List[str]													=				np.load(lowerCAmelCase__						)
    return list(lowerCAmelCase__						)





def 					a__       (					lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__						)    ->	int:
    UpperCAmelCase__			:		int													=				get_videomae_config(lowerCAmelCase__						)

    if "finetuned" in model_name:
        UpperCAmelCase__			:		Any													=				VideoMAEForVideoClassification(lowerCAmelCase__						)
    else:
        UpperCAmelCase__			:		List[Any]													=				VideoMAEForPreTraining(lowerCAmelCase__						)

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

    model.load_state_dict(lowerCAmelCase__						)
    model.eval()

    # verify model on basic input
    UpperCAmelCase__			:		List[Any]													=				VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]  ,					image_std=[0.5, 0.5, 0.5]						)
    UpperCAmelCase__			:		List[str]													=				prepare_video()
    UpperCAmelCase__			:		Union[str, Any]													=				image_processor(lowerCAmelCase__  ,					return_tensors='''pt'''						)

    if "finetuned" not in model_name:
        UpperCAmelCase__			:		List[str]													=				hf_hub_download(repo_id='''hf-internal-testing/bool-masked-pos'''  ,					filename='''bool_masked_pos.pt'''						)
        UpperCAmelCase__			:		Dict													=				torch.load(lowerCAmelCase__						)

    UpperCAmelCase__			:		Any													=				model(**lowerCAmelCase__						)
    UpperCAmelCase__			:		List[Any]													=				outputs.logits

    UpperCAmelCase__			:		Tuple													=				[
        '''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":
        UpperCAmelCase__			:		str													=				torch.Size([1, 4_00]						)
        UpperCAmelCase__			:		Any													=				torch.tensor([-0.9_2_9_1, -0.4_0_6_1, -0.9_3_0_7]						)
    elif model_name == "videomae-small-finetuned-ssv2":
        UpperCAmelCase__			:		Any													=				torch.Size([1, 1_74]						)
        UpperCAmelCase__			:		Dict													=				torch.tensor([0.2_6_7_1, -0.4_6_8_9, -0.8_2_3_5]						)
    elif model_name == "videomae-base":
        UpperCAmelCase__			:		Tuple													=				torch.Size([1, 14_08, 15_36]						)
        UpperCAmelCase__			:		List[str]													=				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":
        UpperCAmelCase__			:		Union[str, Any]													=				torch.Size([1, 14_08, 15_36]						)
        UpperCAmelCase__			:		int													=				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
        UpperCAmelCase__			:		int													=				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":
        UpperCAmelCase__			:		List[str]													=				torch.Size([1, 14_08, 15_36]						)
        UpperCAmelCase__			:		Tuple													=				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":
        UpperCAmelCase__			:		List[Any]													=				torch.Size([1, 4_00]						)
        UpperCAmelCase__			:		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":
        UpperCAmelCase__			:		Optional[int]													=				torch.Size([1, 4_00]						)
        UpperCAmelCase__			:		Tuple													=				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":
        UpperCAmelCase__			:		Tuple													=				torch.Size([1, 4_00]						)
        UpperCAmelCase__			:		Union[str, Any]													=				torch.tensor([0.6_5_8_8, 0.0_9_9_0, -0.2_4_9_3]						)
    elif model_name == "videomae-base-finetuned-kinetics":
        UpperCAmelCase__			:		Optional[Any]													=				torch.Size([1, 4_00]						)
        UpperCAmelCase__			:		Optional[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":
        UpperCAmelCase__			:		Union[str, Any]													=				torch.Size([1, 14_08, 15_36]						)
        UpperCAmelCase__			:		List[Any]													=				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":
        UpperCAmelCase__			:		List[str]													=				torch.Size([1, 1_74]						)
        UpperCAmelCase__			:		int													=				torch.tensor([-0.0_5_3_7, -0.1_5_3_9, -0.3_2_6_6]						)
    elif model_name == "videomae-base-ssv2":
        UpperCAmelCase__			:		int													=				torch.Size([1, 14_08, 15_36]						)
        UpperCAmelCase__			:		Any													=				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":
        UpperCAmelCase__			:		Union[str, Any]													=				torch.Size([1, 1_74]						)
        UpperCAmelCase__			:		Any													=				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]  ,					lowerCAmelCase__  ,					atol=1E-4						)
    else:
        print('''Logits:'''  ,					logits[0, :3, :3]						)
        assert torch.allclose(logits[0, :3, :3]  ,					lowerCAmelCase__  ,					atol=1E-4						)
    print('''Logits ok!'''						)

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

    if push_to_hub:
        print('''Pushing to the hub...'''						)
        model.push_to_hub(lowerCAmelCase__  ,					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_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCamelCase__    					=      {
    '''configuration_rembert''': ['''REMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''RemBertConfig''', '''RemBertOnnxConfig''']
}

try:
  if not is_sentencepiece_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCamelCase__    					=      ['''RemBertTokenizer''']

try:
  if not is_tokenizers_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCamelCase__    					=      ['''RemBertTokenizerFast''']

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCamelCase__    					=      [
      '''REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST''',
      '''RemBertForCausalLM''',
      '''RemBertForMaskedLM''',
      '''RemBertForMultipleChoice''',
      '''RemBertForQuestionAnswering''',
      '''RemBertForSequenceClassification''',
      '''RemBertForTokenClassification''',
      '''RemBertLayer''',
      '''RemBertModel''',
      '''RemBertPreTrainedModel''',
      '''load_tf_weights_in_rembert''',
  ]


try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCamelCase__    					=      [
      '''TF_REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST''',
      '''TFRemBertForCausalLM''',
      '''TFRemBertForMaskedLM''',
      '''TFRemBertForMultipleChoice''',
      '''TFRemBertForQuestionAnswering''',
      '''TFRemBertForSequenceClassification''',
      '''TFRemBertForTokenClassification''',
      '''TFRemBertLayer''',
      '''TFRemBertModel''',
      '''TFRemBertPreTrainedModel''',
  ]


if TYPE_CHECKING:
  from .configuration_rembert import REMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, RemBertConfig, RemBertOnnxConfig

  try:
    if not is_sentencepiece_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_rembert import RemBertTokenizer

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_rembert_fast import RemBertTokenizerFast

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_rembert import (
        REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        RemBertForCausalLM,
        RemBertForMaskedLM,
        RemBertForMultipleChoice,
        RemBertForQuestionAnswering,
        RemBertForSequenceClassification,
        RemBertForTokenClassification,
        RemBertLayer,
        RemBertModel,
        RemBertPreTrainedModel,
        load_tf_weights_in_rembert,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_rembert import (
        TF_REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFRemBertForCausalLM,
        TFRemBertForMaskedLM,
        TFRemBertForMultipleChoice,
        TFRemBertForQuestionAnswering,
        TFRemBertForSequenceClassification,
        TFRemBertForTokenClassification,
        TFRemBertLayer,
        TFRemBertModel,
        TFRemBertPreTrainedModel,
    )


else:
  import sys

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