





from sklearn.metrics import matthews_corrcoef

import datasets


UpperCAmelCase_ 							=       """
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]
"""

UpperCAmelCase_ 							=       """
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
"""

UpperCAmelCase_ 							=       """\
@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 		lowerCAmelCase     (						self					)						->							Optional[int]:
										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 		lowerCAmelCase     (						self			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_=None					)						->							str:
										return {
										    "matthews_correlation": float(matthews_corrcoef(lowerCAmelCase_			,  lowerCAmelCase_			,  sample_weight=lowerCAmelCase_					)					),
										}










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__     :    int  )     ->	List[str]:



					'''simple docstring'''




					_snake_case													=			VideoMAEConfig()

					set_architecture_configs(UpperCamelCase__    ,      UpperCamelCase__  )

					if "finetuned" not in model_name:
										_snake_case													=			False

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

					return config





def   lowerCamelCase__       (       UpperCamelCase__     :    str    ,      UpperCamelCase__     :    Dict  )     ->	int:



					'''simple docstring'''




					if "small" in model_name:
										_snake_case													=			384
										_snake_case													=			1_536
										_snake_case													=			12
										_snake_case													=			16
										_snake_case													=			12
										_snake_case													=			3
										_snake_case													=			192
										_snake_case													=			768
					elif "large" in model_name:
										_snake_case													=			1_024
										_snake_case													=			4_096
										_snake_case													=			24
										_snake_case													=			16
										_snake_case													=			12
										_snake_case													=			8
										_snake_case													=			512
										_snake_case													=			2_048
					elif "huge" in model_name:
										_snake_case													=			1_280
										_snake_case													=			5_120
										_snake_case													=			32
										_snake_case													=			16
										_snake_case													=			12
										_snake_case													=			8
										_snake_case													=			640
										_snake_case													=			2_560
					elif "base" not in model_name:
										raise ValueError('Model name should include either "small", "base", "large", or "huge"'  )





def   lowerCamelCase__       (       UpperCamelCase__     :    Any  )     ->	Tuple:



					'''simple docstring'''




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

					return name





def   lowerCamelCase__       (       UpperCamelCase__     :    Any    ,      UpperCamelCase__     :    List[Any]  )     ->	Union[str, Any]:



					'''simple docstring'''




					for key in orig_state_dict.copy().keys():
										_snake_case													=			orig_state_dict.pop(UpperCamelCase__  )

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

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

					return orig_state_dict





def   lowerCamelCase__       (       )     ->	Union[str, Any]:



					'''simple docstring'''




					_snake_case													=			hf_hub_download(
					    repo_id='hf-internal-testing/spaghetti-video'    ,      filename='eating_spaghetti.npy'    ,      repo_type='dataset'  )
					_snake_case													=			np.load(UpperCamelCase__  )
					return list(UpperCamelCase__  )





def   lowerCamelCase__       (       UpperCamelCase__     :    List[str]    ,      UpperCamelCase__     :    Union[str, Any]    ,      UpperCamelCase__     :    Optional[int]    ,      UpperCamelCase__     :    Optional[int]  )     ->	List[Any]:



					'''simple docstring'''




					_snake_case													=			get_videomae_config(UpperCamelCase__  )

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

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

					model.load_state_dict(UpperCamelCase__  )
					model.eval()

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

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

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

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



