"""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__( __SCREAMING_SNAKE_CASE							:  List[str]   ):
     lowercase_			:       List[Any]									=			VideoMAEConfig()

     set_architecture_configs(__SCREAMING_SNAKE_CASE  ,    __SCREAMING_SNAKE_CASE   )

     if "finetuned" not in model_name:
          lowercase_			:       str									=			False

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

     return config



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



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

     return name



def     lowercase__( __SCREAMING_SNAKE_CASE							:  str  ,    __SCREAMING_SNAKE_CASE							:  int   ):
     for key in orig_state_dict.copy().keys():
          lowercase_			:       Optional[Any]									=			orig_state_dict.pop(__SCREAMING_SNAKE_CASE   )

          if key.startswith('encoder.'   ):
               lowercase_			:       Union[str, Any]									=			key.replace('encoder.'  ,    ''   )

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

     return orig_state_dict



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



def     lowercase__( __SCREAMING_SNAKE_CASE							:  int  ,    __SCREAMING_SNAKE_CASE							:  Optional[int]  ,    __SCREAMING_SNAKE_CASE							:  Optional[int]  ,    __SCREAMING_SNAKE_CASE							:  str   ):
     lowercase_			:       Optional[Any]									=			get_videomae_config(__SCREAMING_SNAKE_CASE   )

     if "finetuned" in model_name:
          lowercase_			:       Tuple									=			VideoMAEForVideoClassification(__SCREAMING_SNAKE_CASE   )
     else:
          lowercase_			:       Tuple									=			VideoMAEForPreTraining(__SCREAMING_SNAKE_CASE   )

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

     model.load_state_dict(__SCREAMING_SNAKE_CASE   )
     model.eval()

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

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

     lowercase_			:       List[Any]									=			model(**__SCREAMING_SNAKE_CASE   )
     lowercase_			:       Optional[Any]									=			outputs.logits

     lowercase_			:       Optional[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":
          lowercase_			:       List[str]									=			torch.Size([1, 4_00]   )
          lowercase_			:       Tuple									=			torch.tensor([-0.9291, -0.4061, -0.9307]   )
     elif model_name == "videomae-small-finetuned-ssv2":
          lowercase_			:       List[Any]									=			torch.Size([1, 1_74]   )
          lowercase_			:       Union[str, Any]									=			torch.tensor([0.2671, -0.4689, -0.8235]   )
     elif model_name == "videomae-base":
          lowercase_			:       Dict									=			torch.Size([1, 14_08, 15_36]   )
          lowercase_			:       Optional[Any]									=			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":
          lowercase_			:       Dict									=			torch.Size([1, 14_08, 15_36]   )
          lowercase_			:       Union[str, Any]									=			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
          lowercase_			:       Dict									=			torch.tensor([0.5142]   ) if config.norm_pix_loss else torch.tensor([0.6469]   )
     elif model_name == "videomae-large":
          lowercase_			:       Tuple									=			torch.Size([1, 14_08, 15_36]   )
          lowercase_			:       str									=			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":
          lowercase_			:       Optional[int]									=			torch.Size([1, 4_00]   )
          lowercase_			:       Dict									=			torch.tensor([0.0771, 0.0011, -0.3625]   )
     elif model_name == "videomae-huge-finetuned-kinetics":
          lowercase_			:       str									=			torch.Size([1, 4_00]   )
          lowercase_			:       List[Any]									=			torch.tensor([0.2433, 0.1632, -0.4894]   )
     elif model_name == "videomae-base-short-finetuned-kinetics":
          lowercase_			:       Optional[int]									=			torch.Size([1, 4_00]   )
          lowercase_			:       int									=			torch.tensor([0.6588, 0.0990, -0.2493]   )
     elif model_name == "videomae-base-finetuned-kinetics":
          lowercase_			:       Tuple									=			torch.Size([1, 4_00]   )
          lowercase_			:       Dict									=			torch.tensor([0.3669, -0.0688, -0.2421]   )
     elif model_name == "videomae-base-short-ssv2":
          lowercase_			:       List[str]									=			torch.Size([1, 14_08, 15_36]   )
          lowercase_			:       str									=			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":
          lowercase_			:       Optional[int]									=			torch.Size([1, 1_74]   )
          lowercase_			:       Tuple									=			torch.tensor([-0.0537, -0.1539, -0.3266]   )
     elif model_name == "videomae-base-ssv2":
          lowercase_			:       Union[str, Any]									=			torch.Size([1, 14_08, 15_36]   )
          lowercase_			:       Optional[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":
          lowercase_			:       List[str]									=			torch.Size([1, 1_74]   )
          lowercase_			:       Union[str, 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]  ,    __SCREAMING_SNAKE_CASE  ,    atol=1E-4   )
     else:
          print('Logits:'  ,    logits[0, :3, :3]   )
          assert torch.allclose(logits[0, :3, :3]  ,    __SCREAMING_SNAKE_CASE  ,    atol=1E-4   )
     print('Logits ok!'   )

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

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


if __name__ == "__main__":
       __SCREAMING_SNAKE_CASE								=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."
       )

       __SCREAMING_SNAKE_CASE								=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 torch

from transformers import (
    UniSpeechSatConfig,
    UniSpeechSatForAudioFrameClassification,
    UniSpeechSatForSequenceClassification,
    UniSpeechSatForXVector,
    WavaVecaFeatureExtractor,
    logging,
)


logging.set_verbosity_info()
__SCREAMING_SNAKE_CASE								=logging.get_logger(__name__)



def     lowercase__( __SCREAMING_SNAKE_CASE							:  List[str]  ,    __SCREAMING_SNAKE_CASE							:  List[str]  ,    __SCREAMING_SNAKE_CASE							:  int   ):
     lowercase_			:       List[str]									=			UniSpeechSatForSequenceClassification.from_pretrained(__SCREAMING_SNAKE_CASE  ,    config=__SCREAMING_SNAKE_CASE   )
     lowercase_			:       List[Any]									=			downstream_dict['projector.weight']
     lowercase_			:       Optional[Any]									=			downstream_dict['projector.bias']
     lowercase_			:       Optional[int]									=			downstream_dict['model.post_net.linear.weight']
     lowercase_			:       Tuple									=			downstream_dict['model.post_net.linear.bias']
     return model



def     lowercase__( __SCREAMING_SNAKE_CASE							:  List[Any]  ,    __SCREAMING_SNAKE_CASE							:  str  ,    __SCREAMING_SNAKE_CASE							:  Tuple   ):
     lowercase_			:       List[Any]									=			UniSpeechSatForAudioFrameClassification.from_pretrained(__SCREAMING_SNAKE_CASE  ,    config=__SCREAMING_SNAKE_CASE   )
     lowercase_			:       Dict									=			downstream_dict['model.linear.weight']
     lowercase_			:       Union[str, Any]									=			downstream_dict['model.linear.bias']
     return model



def     lowercase__( __SCREAMING_SNAKE_CASE							:  Optional[Any]  ,    __SCREAMING_SNAKE_CASE							:  int  ,    __SCREAMING_SNAKE_CASE							:  Optional[Any]   ):
     lowercase_			:       Any									=			UniSpeechSatForXVector.from_pretrained(__SCREAMING_SNAKE_CASE  ,    config=__SCREAMING_SNAKE_CASE   )
     lowercase_			:       int									=			downstream_dict['connector.weight']
     lowercase_			:       List[str]									=			downstream_dict['connector.bias']
     for i, kernel_size in enumerate(hf_config.tdnn_kernel   ):
          lowercase_			:       Dict									=			downstream_dict[
              F'''model.framelevel_feature_extractor.module.{i}.kernel.weight'''
          ]
          lowercase_			:       List[Any]									=			downstream_dict[F'''model.framelevel_feature_extractor.module.{i}.kernel.bias''']

     lowercase_			:       Optional[Any]									=			downstream_dict['model.utterancelevel_feature_extractor.linear1.weight']
     lowercase_			:       int									=			downstream_dict['model.utterancelevel_feature_extractor.linear1.bias']
     lowercase_			:       Dict									=			downstream_dict['model.utterancelevel_feature_extractor.linear2.weight']
     lowercase_			:       int									=			downstream_dict['model.utterancelevel_feature_extractor.linear2.bias']
     lowercase_			:       Dict									=			downstream_dict['objective.W']
     return model



@torch.no_grad()
def     lowercase__( __SCREAMING_SNAKE_CASE							:  str  ,    __SCREAMING_SNAKE_CASE							:  Optional[int]  ,    __SCREAMING_SNAKE_CASE							:  List[Any]  ,    __SCREAMING_SNAKE_CASE							:  Optional[Any]   ):
     lowercase_			:       int									=			torch.load(__SCREAMING_SNAKE_CASE  ,    map_location='cpu'   )

     lowercase_			:       str									=			checkpoint['Downstream']

     lowercase_			:       str									=			UniSpeechSatConfig.from_pretrained(__SCREAMING_SNAKE_CASE   )
     lowercase_			:       Tuple									=			WavaVecaFeatureExtractor.from_pretrained(
         __SCREAMING_SNAKE_CASE  ,    return_attention_mask=__SCREAMING_SNAKE_CASE  ,    do_normalize=__SCREAMING_SNAKE_CASE   )

     lowercase_			:       List[str]									=			hf_config.architectures[0]
     if arch.endswith('ForSequenceClassification'   ):
          lowercase_			:       List[Any]									=			convert_classification(__SCREAMING_SNAKE_CASE  ,    __SCREAMING_SNAKE_CASE  ,    __SCREAMING_SNAKE_CASE   )
     elif arch.endswith('ForAudioFrameClassification'   ):
          lowercase_			:       List[str]									=			convert_diarization(__SCREAMING_SNAKE_CASE  ,    __SCREAMING_SNAKE_CASE  ,    __SCREAMING_SNAKE_CASE   )
     elif arch.endswith('ForXVector'   ):
          lowercase_			:       Optional[Any]									=			convert_xvector(__SCREAMING_SNAKE_CASE  ,    __SCREAMING_SNAKE_CASE  ,    __SCREAMING_SNAKE_CASE   )
     else:
          raise NotImplementedError(F'''S3PRL weights conversion is not supported for {arch}'''   )

     if hf_config.use_weighted_layer_sum:
          lowercase_			:       Dict									=			checkpoint['Featurizer']['weights']

     hf_feature_extractor.save_pretrained(__SCREAMING_SNAKE_CASE   )
     hf_model.save_pretrained(__SCREAMING_SNAKE_CASE   )


if __name__ == "__main__":
       __SCREAMING_SNAKE_CASE								=argparse.ArgumentParser()
       parser.add_argument(
           "--base_model_name", default=None, type=str, help="Name of the huggingface pretrained base model."
       )
       parser.add_argument("--config_path", default=None, type=str, help="Path to the huggingface classifier config.")
       parser.add_argument("--checkpoint_path", default=None, type=str, help="Path to the s3prl checkpoint.")
       parser.add_argument("--model_dump_path", default=None, type=str, help="Path to the final converted model.")
       __SCREAMING_SNAKE_CASE								=parser.parse_args()
       convert_saprl_checkpoint(args.base_model_name, args.config_path, args.checkpoint_path, args.model_dump_path)

