'''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 __lowerCAmelCase					(				UpperCamelCase__							)      ->       Union[str, Any]:
    __lowerCamelCase            =    VideoMAEConfig()

    set_architecture_configs(UpperCamelCase__							,    UpperCamelCase__							)

    if "finetuned" not in model_name:
        __lowerCamelCase            =    False

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

    return config






def __lowerCAmelCase					(				UpperCamelCase__							,    UpperCamelCase__							)      ->       Any:
    if "small" in model_name:
        __lowerCamelCase            =    3_84
        __lowerCamelCase            =    15_36
        __lowerCamelCase            =    12
        __lowerCamelCase            =    16
        __lowerCamelCase            =    12
        __lowerCamelCase            =    3
        __lowerCamelCase            =    1_92
        __lowerCamelCase            =    7_68
    elif "large" in model_name:
        __lowerCamelCase            =    10_24
        __lowerCamelCase            =    40_96
        __lowerCamelCase            =    24
        __lowerCamelCase            =    16
        __lowerCamelCase            =    12
        __lowerCamelCase            =    8
        __lowerCamelCase            =    5_12
        __lowerCamelCase            =    20_48
    elif "huge" in model_name:
        __lowerCamelCase            =    12_80
        __lowerCamelCase            =    51_20
        __lowerCamelCase            =    32
        __lowerCamelCase            =    16
        __lowerCamelCase            =    12
        __lowerCamelCase            =    8
        __lowerCamelCase            =    6_40
        __lowerCamelCase            =    25_60
    elif "base" not in model_name:
        raise ValueError('''Model name should include either "small", "base", "large", or "huge"'''							)






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

    return name






def __lowerCAmelCase					(				UpperCamelCase__							,    UpperCamelCase__							)      ->       str:
    for key in orig_state_dict.copy().keys():
        __lowerCamelCase            =    orig_state_dict.pop(UpperCamelCase__							)

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

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

    return orig_state_dict






def __lowerCAmelCase					(				)      ->       Dict:
    __lowerCamelCase            =    hf_hub_download(
        repo_id='''hf-internal-testing/spaghetti-video'''							,    filename='''eating_spaghetti.npy'''							,    repo_type='''dataset'''							)
    __lowerCamelCase            =    np.load(UpperCamelCase__							)
    return list(UpperCamelCase__							)






def __lowerCAmelCase					(				UpperCamelCase__							,    UpperCamelCase__							,    UpperCamelCase__							,    UpperCamelCase__							)      ->       Any:
    __lowerCamelCase            =    get_videomae_config(UpperCamelCase__							)

    if "finetuned" in model_name:
        __lowerCamelCase            =    VideoMAEForVideoClassification(UpperCamelCase__							)
    else:
        __lowerCamelCase            =    VideoMAEForPreTraining(UpperCamelCase__							)

    # download original checkpoint, hosted on Google Drive
    __lowerCamelCase            =    '''pytorch_model.bin'''
    gdown.cached_download(UpperCamelCase__							,    UpperCamelCase__							,    quiet=UpperCamelCase__							)
    __lowerCamelCase            =    torch.load(UpperCamelCase__							,    map_location='''cpu'''							)
    if "model" in files:
        __lowerCamelCase            =    files['''model''']
    else:
        __lowerCamelCase            =    files['''module''']
    __lowerCamelCase            =    convert_state_dict(UpperCamelCase__							,    UpperCamelCase__							)

    model.load_state_dict(UpperCamelCase__							)
    model.eval()

    # verify model on basic input
    __lowerCamelCase            =    VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]							,    image_std=[0.5, 0.5, 0.5]							)
    __lowerCamelCase            =    prepare_video()
    __lowerCamelCase            =    image_processor(UpperCamelCase__							,    return_tensors='''pt'''							)

    if "finetuned" not in model_name:
        __lowerCamelCase            =    hf_hub_download(repo_id='''hf-internal-testing/bool-masked-pos'''							,    filename='''bool_masked_pos.pt'''							)
        __lowerCamelCase            =    torch.load(UpperCamelCase__							)

    __lowerCamelCase            =    model(**UpperCamelCase__							)
    __lowerCamelCase            =    outputs.logits

    __lowerCamelCase            =    [
        '''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":
        __lowerCamelCase            =    torch.Size([1, 4_00]							)
        __lowerCamelCase            =    torch.tensor([-0.9_2_9_1, -0.4_0_6_1, -0.9_3_0_7]							)
    elif model_name == "videomae-small-finetuned-ssv2":
        __lowerCamelCase            =    torch.Size([1, 1_74]							)
        __lowerCamelCase            =    torch.tensor([0.2_6_7_1, -0.4_6_8_9, -0.8_2_3_5]							)
    elif model_name == "videomae-base":
        __lowerCamelCase            =    torch.Size([1, 14_08, 15_36]							)
        __lowerCamelCase            =    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":
        __lowerCamelCase            =    torch.Size([1, 14_08, 15_36]							)
        __lowerCamelCase            =    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
        __lowerCamelCase            =    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":
        __lowerCamelCase            =    torch.Size([1, 14_08, 15_36]							)
        __lowerCamelCase            =    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":
        __lowerCamelCase            =    torch.Size([1, 4_00]							)
        __lowerCamelCase            =    torch.tensor([0.0_7_7_1, 0.0_0_1_1, -0.3_6_2_5]							)
    elif model_name == "videomae-huge-finetuned-kinetics":
        __lowerCamelCase            =    torch.Size([1, 4_00]							)
        __lowerCamelCase            =    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":
        __lowerCamelCase            =    torch.Size([1, 4_00]							)
        __lowerCamelCase            =    torch.tensor([0.6_5_8_8, 0.0_9_9_0, -0.2_4_9_3]							)
    elif model_name == "videomae-base-finetuned-kinetics":
        __lowerCamelCase            =    torch.Size([1, 4_00]							)
        __lowerCamelCase            =    torch.tensor([0.3_6_6_9, -0.0_6_8_8, -0.2_4_2_1]							)
    elif model_name == "videomae-base-short-ssv2":
        __lowerCamelCase            =    torch.Size([1, 14_08, 15_36]							)
        __lowerCamelCase            =    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":
        __lowerCamelCase            =    torch.Size([1, 1_74]							)
        __lowerCamelCase            =    torch.tensor([-0.0_5_3_7, -0.1_5_3_9, -0.3_2_6_6]							)
    elif model_name == "videomae-base-ssv2":
        __lowerCamelCase            =    torch.Size([1, 14_08, 15_36]							)
        __lowerCamelCase            =    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":
        __lowerCamelCase            =    torch.Size([1, 1_74]							)
        __lowerCamelCase            =    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":
        __lowerCamelCase            =    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'''




import argparse

import torch

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


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






def __lowerCAmelCase					(				UpperCamelCase__							,    UpperCamelCase__							,    UpperCamelCase__							)      ->       Optional[Any]:
    __lowerCamelCase            =    UniSpeechSatForSequenceClassification.from_pretrained(UpperCamelCase__							,    config=UpperCamelCase__							)
    __lowerCamelCase            =    downstream_dict['''projector.weight''']
    __lowerCamelCase            =    downstream_dict['''projector.bias''']
    __lowerCamelCase            =    downstream_dict['''model.post_net.linear.weight''']
    __lowerCamelCase            =    downstream_dict['''model.post_net.linear.bias''']
    return model






def __lowerCAmelCase					(				UpperCamelCase__							,    UpperCamelCase__							,    UpperCamelCase__							)      ->       Union[str, Any]:
    __lowerCamelCase            =    UniSpeechSatForAudioFrameClassification.from_pretrained(UpperCamelCase__							,    config=UpperCamelCase__							)
    __lowerCamelCase            =    downstream_dict['''model.linear.weight''']
    __lowerCamelCase            =    downstream_dict['''model.linear.bias''']
    return model






def __lowerCAmelCase					(				UpperCamelCase__							,    UpperCamelCase__							,    UpperCamelCase__							)      ->       List[Any]:
    __lowerCamelCase            =    UniSpeechSatForXVector.from_pretrained(UpperCamelCase__							,    config=UpperCamelCase__							)
    __lowerCamelCase            =    downstream_dict['''connector.weight''']
    __lowerCamelCase            =    downstream_dict['''connector.bias''']
    for i, kernel_size in enumerate(hf_config.tdnn_kernel							):
        __lowerCamelCase            =    downstream_dict[
            f"""model.framelevel_feature_extractor.module.{i}.kernel.weight"""
        ]
        __lowerCamelCase            =    downstream_dict[f"""model.framelevel_feature_extractor.module.{i}.kernel.bias"""]

    __lowerCamelCase            =    downstream_dict['''model.utterancelevel_feature_extractor.linear1.weight''']
    __lowerCamelCase            =    downstream_dict['''model.utterancelevel_feature_extractor.linear1.bias''']
    __lowerCamelCase            =    downstream_dict['''model.utterancelevel_feature_extractor.linear2.weight''']
    __lowerCamelCase            =    downstream_dict['''model.utterancelevel_feature_extractor.linear2.bias''']
    __lowerCamelCase            =    downstream_dict['''objective.W''']
    return model






@torch.no_grad()
def __lowerCAmelCase					(				UpperCamelCase__							,    UpperCamelCase__							,    UpperCamelCase__							,    UpperCamelCase__							)      ->       Optional[Any]:
    __lowerCamelCase            =    torch.load(UpperCamelCase__							,    map_location='''cpu'''							)

    __lowerCamelCase            =    checkpoint['''Downstream''']

    __lowerCamelCase            =    UniSpeechSatConfig.from_pretrained(UpperCamelCase__							)
    __lowerCamelCase            =    WavaVecaFeatureExtractor.from_pretrained(
        UpperCamelCase__							,    return_attention_mask=UpperCamelCase__							,    do_normalize=UpperCamelCase__							)

    __lowerCamelCase            =    hf_config.architectures[0]
    if arch.endswith('''ForSequenceClassification'''							):
        __lowerCamelCase            =    convert_classification(UpperCamelCase__							,    UpperCamelCase__							,    UpperCamelCase__							)
    elif arch.endswith('''ForAudioFrameClassification'''							):
        __lowerCamelCase            =    convert_diarization(UpperCamelCase__							,    UpperCamelCase__							,    UpperCamelCase__							)
    elif arch.endswith('''ForXVector'''							):
        __lowerCamelCase            =    convert_xvector(UpperCamelCase__							,    UpperCamelCase__							,    UpperCamelCase__							)
    else:
        raise NotImplementedError(f"""S3PRL weights conversion is not supported for {arch}"""							)

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

    hf_feature_extractor.save_pretrained(UpperCamelCase__							)
    hf_model.save_pretrained(UpperCamelCase__							)


if __name__ == "__main__":
  __UpperCAmelCase					=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.")
  __UpperCAmelCase					=parser.parse_args()
  convert_saprl_checkpoint(args.base_model_name, args.config_path, args.checkpoint_path, args.model_dump_path)
