



"""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_				:  Any       =			VideoMAEConfig()

    set_architecture_configs(_a	,    _a					)

    if "finetuned" not in model_name:
        snake_case_				:  Optional[Any]       =			False

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

    return config




def __lowercase    (							_a	,    _a					):







    if "small" in model_name:
        snake_case_				:  Union[str, Any]       =			384
        snake_case_				:  Optional[int]       =			1_536
        snake_case_				:  Union[str, Any]       =			12
        snake_case_				:  Optional[int]       =			16
        snake_case_				:  Optional[Any]       =			12
        snake_case_				:  List[Any]       =			3
        snake_case_				:  Tuple       =			192
        snake_case_				:  List[Any]       =			768
    elif "large" in model_name:
        snake_case_				:  List[Any]       =			1_024
        snake_case_				:  Optional[int]       =			4_096
        snake_case_				:  Optional[int]       =			24
        snake_case_				:  Tuple       =			16
        snake_case_				:  List[Any]       =			12
        snake_case_				:  int       =			8
        snake_case_				:  Tuple       =			512
        snake_case_				:  Union[str, Any]       =			2_048
    elif "huge" in model_name:
        snake_case_				:  Tuple       =			1_280
        snake_case_				:  Optional[Any]       =			5_120
        snake_case_				:  Dict       =			32
        snake_case_				:  Dict       =			16
        snake_case_				:  Any       =			12
        snake_case_				:  Tuple       =			8
        snake_case_				:  Optional[int]       =			640
        snake_case_				:  Optional[Any]       =			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_				:  List[Any]       =			name.replace('''encoder.'''	,    ''''''					)
    if "cls_token" in name:
        snake_case_				:  Tuple       =			name.replace('''cls_token'''	,    '''videomae.embeddings.cls_token'''					)
    if "decoder_pos_embed" in name:
        snake_case_				:  List[Any]       =			name.replace('''decoder_pos_embed'''	,    '''decoder.decoder_pos_embed'''					)
    if "pos_embed" in name and "decoder" not in name:
        snake_case_				:  Dict       =			name.replace('''pos_embed'''	,    '''videomae.embeddings.position_embeddings'''					)
    if "patch_embed.proj" in name:
        snake_case_				:  str       =			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_				:  Union[str, Any]       =			name.replace('''decoder.blocks'''	,    '''decoder.decoder_layers'''					)
    if "blocks" in name:
        snake_case_				:  Optional[Any]       =			name.replace('''blocks'''	,    '''videomae.encoder.layer'''					)
    if "attn.proj" in name:
        snake_case_				:  Any       =			name.replace('''attn.proj'''	,    '''attention.output.dense'''					)
    if "attn" in name and "bias" not in name:
        snake_case_				:  Tuple       =			name.replace('''attn'''	,    '''attention.self'''					)
    if "attn" in name:
        snake_case_				:  Dict       =			name.replace('''attn'''	,    '''attention.attention'''					)
    if "norm1" in name:
        snake_case_				:  List[str]       =			name.replace('''norm1'''	,    '''layernorm_before'''					)
    if "norm2" in name:
        snake_case_				:  Optional[Any]       =			name.replace('''norm2'''	,    '''layernorm_after'''					)
    if "mlp.fc1" in name:
        snake_case_				:  Optional[Any]       =			name.replace('''mlp.fc1'''	,    '''intermediate.dense'''					)
    if "mlp.fc2" in name:
        snake_case_				:  Dict       =			name.replace('''mlp.fc2'''	,    '''output.dense'''					)
    if "decoder_embed" in name:
        snake_case_				:  Optional[Any]       =			name.replace('''decoder_embed'''	,    '''decoder.decoder_embed'''					)
    if "decoder_norm" in name:
        snake_case_				:  Union[str, Any]       =			name.replace('''decoder_norm'''	,    '''decoder.decoder_norm'''					)
    if "decoder_pred" in name:
        snake_case_				:  str       =			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_				:  int       =			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_				:  List[str]       =			orig_state_dict.pop(_a					)

        if key.startswith('''encoder.'''					):
            snake_case_				:  Optional[Any]       =			key.replace('''encoder.'''	,    ''''''					)

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

    return orig_state_dict




def __lowercase    (							):
    snake_case_				:  Tuple       =			hf_hub_download(
        repo_id='''hf-internal-testing/spaghetti-video'''	,    filename='''eating_spaghetti.npy'''	,    repo_type='''dataset'''					)
    snake_case_				:  Dict       =			np.load(_a					)
    return list(_a					)




def __lowercase    (							_a	,    _a	,    _a	,    _a					):
    snake_case_				:  Dict       =			get_videomae_config(_a					)

    if "finetuned" in model_name:
        snake_case_				:  Any       =			VideoMAEForVideoClassification(_a					)
    else:
        snake_case_				:  Tuple       =			VideoMAEForPreTraining(_a					)

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

    model.load_state_dict(_a					)
    model.eval()

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

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

    snake_case_				:  str       =			model(**_a					)
    snake_case_				:  int       =			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_				:  Tuple       =			torch.Size([1, 400]					)
        snake_case_				:  List[Any]       =			torch.tensor([-0.9291, -0.4061, -0.9307]					)
    elif model_name == "videomae-small-finetuned-ssv2":
        snake_case_				:  Tuple       =			torch.Size([1, 174]					)
        snake_case_				:  Optional[int]       =			torch.tensor([0.2671, -0.4689, -0.8235]					)
    elif model_name == "videomae-base":
        snake_case_				:  Any       =			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_				:  List[str]       =			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_				:  Optional[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_				:  Any       =			torch.Size([1, 400]					)
        snake_case_				:  Dict       =			torch.tensor([0.0771, 0.0011, -0.3625]					)
    elif model_name == "videomae-huge-finetuned-kinetics":
        snake_case_				:  List[str]       =			torch.Size([1, 400]					)
        snake_case_				:  List[Any]       =			torch.tensor([0.2433, 0.1632, -0.4894]					)
    elif model_name == "videomae-base-short-finetuned-kinetics":
        snake_case_				:  int       =			torch.Size([1, 400]					)
        snake_case_				:  Dict       =			torch.tensor([0.6588, 0.0990, -0.2493]					)
    elif model_name == "videomae-base-finetuned-kinetics":
        snake_case_				:  Union[str, Any]       =			torch.Size([1, 400]					)
        snake_case_				:  List[str]       =			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_				:  Tuple       =			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_				:  str       =			torch.Size([1, 174]					)
        snake_case_				:  Optional[Any]       =			torch.tensor([-0.0537, -0.1539, -0.3266]					)
    elif model_name == "videomae-base-ssv2":
        snake_case_				:  List[str]       =			torch.Size([1, 1_408, 1_536]					)
        snake_case_				:  str       =			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_				:  int       =			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]	,    _a	,    atol=1E-4					)
    else:
        print('''Logits:'''	,    logits[0, :3, :3]					)
        assert torch.allclose(logits[0, :3, :3]	,    _a	,    atol=1E-4					)
    print('''Logits ok!'''					)

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

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


if __name__ == "__main__":
     lowercase__					:						Optional[int]												=			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__					:						Tuple												=			parser.parse_args()
     convert_videomae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)







"""simple docstring"""

import numpy as np
from numpy import ndarray
from scipy.optimize import Bounds, LinearConstraint, minimize




def __lowercase    (							_a					):
    return np.dot(_a	,    _a					)




class  _UpperCAmelCase :

 def __init__(  self	: int						,  *,
     lowercase_	: float = np.inf						,  lowercase_	: str = "linear"						,  lowercase_	: float = 0.0						,  ):
     snake_case_				:  Optional[Any]       =			regularization
     snake_case_				:  Tuple       =			gamma
     if kernel == "linear":
         snake_case_				:  int       =			self.__linear
     elif kernel == "rbf":
         if self.gamma == 0:
             raise ValueError('''rbf kernel requires gamma'''       )
         if not isinstance(self.gamma						,  (float, int)       ):
             raise ValueError('''gamma must be float or int'''       )
         if not self.gamma > 0:
             raise ValueError('''gamma must be > 0'''       )
         snake_case_				:  int       =			self.__rbf
         # in the future, there could be a default value like in sklearn
         # sklear: def_gamma = 1/(n_features * X.var()) (wiki)
         # previously it was 1/(n_features)
     else:
         snake_case_				:  List[Any]       =			f"Unknown kernel: {kernel}"
         raise ValueError(lowercase_       )

 def 		_snake_case      (  self	: int						,  lowercase_	: ndarray						,  lowercase_	: ndarray       ):
     return np.dot(lowercase_						,  lowercase_       )

 def 		_snake_case      (  self	: int						,  lowercase_	: ndarray						,  lowercase_	: ndarray       ):
     return np.exp(-(self.gamma * norm_squared(vectora - vectora       ))       )

 def 		_snake_case      (  self	: Any						,  lowercase_	: list[ndarray]						,  lowercase_	: ndarray       ):
     snake_case_				:  Union[str, Any]       =			observations
     snake_case_				:  int       =			classes

     # using Wolfe's Dual to calculate w.
     # Primal problem: minimize 1/2*norm_squared(w)
     #   constraint: yn(w . xn + b) >= 1
     #
     # With l a vector
     # Dual problem: maximize sum_n(ln) -
     #       1/2 * sum_n(sum_m(ln*lm*yn*ym*xn . xm))
     #   constraint: self.C >= ln >= 0
     #           and sum_n(ln*yn) = 0
     # Then we get w using w = sum_n(ln*yn*xn)
     # At the end we can get b ~= mean(yn - w . xn)
     #
     # Since we use kernels, we only need l_star to calculate b
     # and to classify observations

     ((snake_case_),			)				:  List[str]       =			np.shape(lowercase_       )

     def to_minimize(lowercase_	: ndarray       ) -> float:
         snake_case_				:  Tuple       =			0
         ((snake_case_),			)				:  Optional[Any]       =			np.shape(lowercase_       )
         for i in range(lowercase_       ):
             for j in range(lowercase_       ):
                 s += (
                     candidate[i]
                     * candidate[j]
                     * classes[i]
                     * classes[j]
                     * self.kernel(observations[i]						,  observations[j]       )
                 )
         return 1 / 2 * s - sum(lowercase_       )

     snake_case_				:  Optional[Any]       =			LinearConstraint(lowercase_						,  0						,  0       )
     snake_case_				:  str       =			Bounds(0						,  self.regularization       )

     snake_case_				:  int       =			minimize(
         lowercase_						,  np.ones(lowercase_       )						,  bounds=lowercase_						,  constraints=[ly_contraint]       ).x
     snake_case_				:  Optional[Any]       =			l_star

     # calculating mean offset of separation plane to points
     snake_case_				:  List[Any]       =			0
     for i in range(lowercase_       ):
         for j in range(lowercase_       ):
             s += classes[i] - classes[i] * self.optimum[i] * self.kernel(
                 observations[i]						,  observations[j]       )
     snake_case_				:  Union[str, Any]       =			s / n






 def 		_snake_case      (  self	: List[str]						,  lowercase_	: ndarray       ):
     snake_case_				:  int       =			sum(
         self.optimum[n]
         * self.classes[n]
         * self.kernel(self.observations[n]						,  lowercase_       )
         for n in range(len(self.classes       )       )       )
     return 1 if s + self.offset >= 0 else -1


if __name__ == "__main__":
     import doctest

     doctest.testmod()





