
from sklearn.metrics import matthews_corrcoef

import datasets


__lowerCamelCase				=			"""
Compute the Matthews correlation coefficient (MCC)

The Matthews correlation coefficient is used in machine learning as a
measure of the quality of binary and multiclass classifications. It takes
into account true and false positives and negatives and is generally
regarded as a balanced measure which can be used even if the classes are of
very different sizes. The MCC is in essence a correlation coefficient value
between -1 and +1. A coefficient of +1 represents a perfect prediction, 0
an average random prediction and -1 an inverse prediction.  The statistic
is also known as the phi coefficient. [source: Wikipedia]
"""

__lowerCamelCase				=			"""
Args:
    predictions (list of int): Predicted labels, as returned by a model.
    references (list of int): Ground truth labels.
    sample_weight (list of int, float, or bool): Sample weights. Defaults to `None`.
Returns:
    matthews_correlation (dict containing float): Matthews correlation.
Examples:
    Example 1, a basic example with only predictions and references as inputs:
        >>> matthews_metric = datasets.load_metric(\"matthews_correlation\")
        >>> results = matthews_metric.compute(references=[1, 3, 2, 0, 3, 2],
        ...                                     predictions=[1, 2, 2, 0, 3, 3])
        >>> print(round(results['matthews_correlation'], 2))
        0.54

    Example 2, the same example as above, but also including sample weights:
        >>> matthews_metric = datasets.load_metric(\"matthews_correlation\")
        >>> results = matthews_metric.compute(references=[1, 3, 2, 0, 3, 2],
        ...                                     predictions=[1, 2, 2, 0, 3, 3],
        ...                                     sample_weight=[0.5, 3, 1, 1, 1, 2])
        >>> print(round(results['matthews_correlation'], 2))
        0.1

    Example 3, the same example as above, but with sample weights that cause a negative correlation:
        >>> matthews_metric = datasets.load_metric(\"matthews_correlation\")
        >>> results = matthews_metric.compute(references=[1, 3, 2, 0, 3, 2],
        ...                                     predictions=[1, 2, 2, 0, 3, 3],
        ...                                     sample_weight=[0.5, 1, 0, 0, 0, 1])
        >>> print(round(results['matthews_correlation'], 2))
        -0.25
"""

__lowerCamelCase				=			"""\
@article{scikit-learn,
  title={Scikit-learn: Machine Learning in {P}ython},
  author={Pedregosa, F. and Varoquaux, G. and Gramfort, A. and Michel, V.
         and Thirion, B. and Grisel, O. and Blondel, M. and Prettenhofer, P.
         and Weiss, R. and Dubourg, V. and Vanderplas, J. and Passos, A. and
         Cournapeau, D. and Brucher, M. and Perrot, M. and Duchesnay, E.},
  journal={Journal of Machine Learning Research},
  volume={12},
  pages={2825--2830},
  year={2011}
}
"""






@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION			,_KWARGS_DESCRIPTION					)
class     UpperCAmelCase						(		datasets.Metric					):



   def 							_SCREAMING_SNAKE_CASE  (self						:   Tuple					)						->      Optional[int]:


          '''simple docstring'''


          return datasets.MetricInfo(
              description=_DESCRIPTION       ,		citation=_CITATION       ,		inputs_description=_KWARGS_DESCRIPTION       ,		features=datasets.Features(
                  {
                      "predictions": datasets.Value("int32"					),
                      "references": datasets.Value("int32"					),
                  }					)       ,		reference_urls=[
                  "https://scikit-learn.org/stable/modules/generated/sklearn.metrics.matthews_corrcoef.html"
              ]       ,		)



   def 							_SCREAMING_SNAKE_CASE  (self						:   Optional[int]       ,		snake_case__						:   Optional[Any]       ,		snake_case__						:   Union[str, Any]       ,		snake_case__						:   Any=None					)						->      List[Any]:


          '''simple docstring'''


          return {
              "matthews_correlation": float(matthews_corrcoef(snake_case__       ,		snake_case__       ,		sample_weight=snake_case__					)					),
          }







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      UpperCamelCase   (       __lowerCamelCase	:					Dict       ):
       snake_case			:  Any     							= VideoMAEConfig()

       set_architecture_configs(__lowerCamelCase							,   __lowerCamelCase       )

       if "finetuned" not in model_name:
              snake_case			:  int     							= False

       if "finetuned" in model_name:
              snake_case			:  Dict     							= "huggingface/label-files"
              if "kinetics" in model_name:
                     snake_case			:  List[str]     							= 400
                     snake_case			:  List[str]     							= "kinetics400-id2label.json"
              elif "ssv2" in model_name:
                     snake_case			:  List[str]     							= 174
                     snake_case			:  Optional[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			:  Tuple     							= json.load(open(hf_hub_download(__lowerCamelCase							,   __lowerCamelCase							,   repo_type="dataset"       )							,   "r"       )       )
              snake_case			:  Dict     							= {int(__lowerCamelCase       ): v for k, v in idalabel.items()}
              snake_case			:  Optional[Any]     							= idalabel
              snake_case			:  Tuple     							= {v: k for k, v in idalabel.items()}

       return config


def      UpperCamelCase   (       __lowerCamelCase	:					Optional[Any]							,   __lowerCamelCase	:					int       ):
       if "small" in model_name:
              snake_case			:  Dict     							= 384
              snake_case			:  Optional[Any]     							= 1536
              snake_case			:  Union[str, Any]     							= 12
              snake_case			:  Dict     							= 16
              snake_case			:  List[str]     							= 12
              snake_case			:  List[str]     							= 3
              snake_case			:  str     							= 192
              snake_case			:  Union[str, Any]     							= 768
       elif "large" in model_name:
              snake_case			:  Tuple     							= 1024
              snake_case			:  List[Any]     							= 4096
              snake_case			:  Any     							= 24
              snake_case			:  Optional[int]     							= 16
              snake_case			:  List[str]     							= 12
              snake_case			:  int     							= 8
              snake_case			:  Any     							= 512
              snake_case			:  List[Any]     							= 2048
       elif "huge" in model_name:
              snake_case			:  Any     							= 1280
              snake_case			:  List[str]     							= 5120
              snake_case			:  Dict     							= 32
              snake_case			:  int     							= 16
              snake_case			:  int     							= 12
              snake_case			:  Tuple     							= 8
              snake_case			:  str     							= 640
              snake_case			:  List[Any]     							= 2560
       elif "base" not in model_name:
              raise ValueError("Model name should include either \"small\", \"base\", \"large\", or \"huge\""       )


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

       return name


def      UpperCamelCase   (       __lowerCamelCase	:					Optional[Any]							,   __lowerCamelCase	:					Any       ):
       for key in orig_state_dict.copy().keys():
              snake_case			:  Optional[Any]     							= orig_state_dict.pop(__lowerCamelCase       )

              if key.startswith("encoder."       ):
                     snake_case			:  Optional[Any]     							= key.replace("encoder."							,   ""       )

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

       return orig_state_dict


def      UpperCamelCase   (       ):
       snake_case			:  Union[str, Any]     							= hf_hub_download(
           repo_id="hf-internal-testing/spaghetti-video"							,   filename="eating_spaghetti.npy"							,   repo_type="dataset"       )
       snake_case			:  Optional[Any]     							= np.load(__lowerCamelCase       )
       return list(__lowerCamelCase       )


def      UpperCamelCase   (       __lowerCamelCase	:					Optional[Any]							,   __lowerCamelCase	:					Union[str, Any]							,   __lowerCamelCase	:					Optional[Any]							,   __lowerCamelCase	:					Any       ):
       snake_case			:  List[str]     							= get_videomae_config(__lowerCamelCase       )

       if "finetuned" in model_name:
              snake_case			:  List[Any]     							= VideoMAEForVideoClassification(__lowerCamelCase       )
       else:
              snake_case			:  List[str]     							= VideoMAEForPreTraining(__lowerCamelCase       )

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

       model.load_state_dict(__lowerCamelCase       )
       model.eval()

       # verify model on basic input
       snake_case			:  Tuple     							= VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]							,   image_std=[0.5, 0.5, 0.5]       )
       snake_case			:  Dict     							= prepare_video()
       snake_case			:  Dict     							= image_processor(__lowerCamelCase							,   return_tensors="pt"       )

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

       snake_case			:  Any     							= model(**__lowerCamelCase       )
       snake_case			:  List[Any]     							= outputs.logits

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

   __lowerCamelCase				=			parser.parse_args()
   convert_videomae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)





