



"""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			(							_a    ):
    snake_case_ :     List[str]				=	VideoMAEConfig()

    set_architecture_configs(lowerCAmelCase__		,			lowerCAmelCase__    )

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

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

    return config




def __lowercase			(							_a		,			_a    ):
    if "small" in model_name:
        snake_case_ :     Any				=	384
        snake_case_ :     List[Any]				=	1_536
        snake_case_ :     List[str]				=	12
        snake_case_ :     List[Any]				=	16
        snake_case_ :     List[str]				=	12
        snake_case_ :     Optional[Any]				=	3
        snake_case_ :     Union[str, Any]				=	192
        snake_case_ :     Optional[int]				=	768
    elif "large" in model_name:
        snake_case_ :     List[Any]				=	1_024
        snake_case_ :     str				=	4_096
        snake_case_ :     Optional[Any]				=	24
        snake_case_ :     List[str]				=	16
        snake_case_ :     List[Any]				=	12
        snake_case_ :     List[str]				=	8
        snake_case_ :     Tuple				=	512
        snake_case_ :     Optional[int]				=	2_048
    elif "huge" in model_name:
        snake_case_ :     Union[str, Any]				=	1_280
        snake_case_ :     str				=	5_120
        snake_case_ :     Dict				=	32
        snake_case_ :     List[str]				=	16
        snake_case_ :     int				=	12
        snake_case_ :     str				=	8
        snake_case_ :     Optional[Any]				=	640
        snake_case_ :     str				=	2_560
    elif "base" not in model_name:
        raise ValueError('''Model name should include either "small", "base", "large", or "huge"'''    )




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

    return name




def __lowercase			(							_a		,			_a    ):
    for key in orig_state_dict.copy().keys():
        snake_case_ :     Any				=	orig_state_dict.pop(lowerCAmelCase__    )

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

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

    return orig_state_dict




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




def __lowercase			(							_a		,			_a		,			_a		,			_a    ):
    snake_case_ :     int				=	get_videomae_config(lowerCAmelCase__    )

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

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

    model.load_state_dict(lowerCAmelCase__    )
    model.eval()

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

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

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

    snake_case_ :     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":
        snake_case_ :     str				=	torch.Size([1, 400]    )
        snake_case_ :     Any				=	torch.tensor([-0.9291, -0.4061, -0.9307]    )
    elif model_name == "videomae-small-finetuned-ssv2":
        snake_case_ :     Any				=	torch.Size([1, 174]    )
        snake_case_ :     Dict				=	torch.tensor([0.2671, -0.4689, -0.8235]    )
    elif model_name == "videomae-base":
        snake_case_ :     Tuple				=	torch.Size([1, 1_408, 1_536]    )
        snake_case_ :     List[str]				=	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_ :     Union[str, Any]				=	torch.Size([1, 1_408, 1_536]    )
        snake_case_ :     int				=	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_ :     int				=	torch.tensor([0.5142]    ) if config.norm_pix_loss else torch.tensor([0.6469]    )
    elif model_name == "videomae-large":
        snake_case_ :     List[str]				=	torch.Size([1, 1_408, 1_536]    )
        snake_case_ :     Tuple				=	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_ :     List[Any]				=	torch.Size([1, 400]    )
        snake_case_ :     Any				=	torch.tensor([0.0771, 0.0011, -0.3625]    )
    elif model_name == "videomae-huge-finetuned-kinetics":
        snake_case_ :     Optional[int]				=	torch.Size([1, 400]    )
        snake_case_ :     Tuple				=	torch.tensor([0.2433, 0.1632, -0.4894]    )
    elif model_name == "videomae-base-short-finetuned-kinetics":
        snake_case_ :     Tuple				=	torch.Size([1, 400]    )
        snake_case_ :     Union[str, Any]				=	torch.tensor([0.6588, 0.0990, -0.2493]    )
    elif model_name == "videomae-base-finetuned-kinetics":
        snake_case_ :     Optional[Any]				=	torch.Size([1, 400]    )
        snake_case_ :     Optional[Any]				=	torch.tensor([0.3669, -0.0688, -0.2421]    )
    elif model_name == "videomae-base-short-ssv2":
        snake_case_ :     Union[str, Any]				=	torch.Size([1, 1_408, 1_536]    )
        snake_case_ :     List[Any]				=	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_ :     List[str]				=	torch.Size([1, 174]    )
        snake_case_ :     int				=	torch.tensor([-0.0537, -0.1539, -0.3266]    )
    elif model_name == "videomae-base-ssv2":
        snake_case_ :     int				=	torch.Size([1, 1_408, 1_536]    )
        snake_case_ :     Any				=	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_ :     Union[str, Any]				=	torch.Size([1, 174]    )
        snake_case_ :     Any				=	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]		,			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":
        snake_case_ :     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__":
     lowercase__    :  List[Any]							    =	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.'''
     )

     lowercase__    :  Optional[Any]							    =	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_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCamelCase__    					=      {
    '''configuration_distilbert''': [
        '''DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''DistilBertConfig''',
        '''DistilBertOnnxConfig''',
    ],
    '''tokenization_distilbert''': ['''DistilBertTokenizer'''],
}

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

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCamelCase__    					=      [
      '''DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST''',
      '''DistilBertForMaskedLM''',
      '''DistilBertForMultipleChoice''',
      '''DistilBertForQuestionAnswering''',
      '''DistilBertForSequenceClassification''',
      '''DistilBertForTokenClassification''',
      '''DistilBertModel''',
      '''DistilBertPreTrainedModel''',
  ]

try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCamelCase__    					=      [
      '''TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST''',
      '''TFDistilBertForMaskedLM''',
      '''TFDistilBertForMultipleChoice''',
      '''TFDistilBertForQuestionAnswering''',
      '''TFDistilBertForSequenceClassification''',
      '''TFDistilBertForTokenClassification''',
      '''TFDistilBertMainLayer''',
      '''TFDistilBertModel''',
      '''TFDistilBertPreTrainedModel''',
  ]

try:
  if not is_flax_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCamelCase__    					=      [
      '''FlaxDistilBertForMaskedLM''',
      '''FlaxDistilBertForMultipleChoice''',
      '''FlaxDistilBertForQuestionAnswering''',
      '''FlaxDistilBertForSequenceClassification''',
      '''FlaxDistilBertForTokenClassification''',
      '''FlaxDistilBertModel''',
      '''FlaxDistilBertPreTrainedModel''',
  ]


if TYPE_CHECKING:
  from .configuration_distilbert import (
      DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
      DistilBertConfig,
      DistilBertOnnxConfig,
  )
  from .tokenization_distilbert import DistilBertTokenizer

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_distilbert_fast import DistilBertTokenizerFast

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_distilbert import (
        DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        DistilBertForMaskedLM,
        DistilBertForMultipleChoice,
        DistilBertForQuestionAnswering,
        DistilBertForSequenceClassification,
        DistilBertForTokenClassification,
        DistilBertModel,
        DistilBertPreTrainedModel,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_distilbert import (
        TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFDistilBertForMaskedLM,
        TFDistilBertForMultipleChoice,
        TFDistilBertForQuestionAnswering,
        TFDistilBertForSequenceClassification,
        TFDistilBertForTokenClassification,
        TFDistilBertMainLayer,
        TFDistilBertModel,
        TFDistilBertPreTrainedModel,
    )

  try:
    if not is_flax_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_flax_distilbert import (
        FlaxDistilBertForMaskedLM,
        FlaxDistilBertForMultipleChoice,
        FlaxDistilBertForQuestionAnswering,
        FlaxDistilBertForSequenceClassification,
        FlaxDistilBertForTokenClassification,
        FlaxDistilBertModel,
        FlaxDistilBertPreTrainedModel,
    )

else:
  import sys

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