




'''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  _UpperCamelCase   (					SCREAMING_SNAKE_CASE__      :   Any							):
 '''simple docstring'''





 UpperCAmelCase__      				=     VideoMAEConfig()

 set_architecture_configs(__lowercase							,    __lowercase							)

 if "finetuned" not in model_name:
  UpperCAmelCase__      				=     False

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

 return config

def  _UpperCamelCase   (					SCREAMING_SNAKE_CASE__      :   List[Any]							,    SCREAMING_SNAKE_CASE__      :   Tuple							):
 '''simple docstring'''





 if "small" in model_name:
  UpperCAmelCase__      				=     384
  UpperCAmelCase__      				=     1536
  UpperCAmelCase__      				=     12
  UpperCAmelCase__      				=     16
  UpperCAmelCase__      				=     12
  UpperCAmelCase__      				=     3
  UpperCAmelCase__      				=     192
  UpperCAmelCase__      				=     768
 elif "large" in model_name:
  UpperCAmelCase__      				=     1024
  UpperCAmelCase__      				=     4096
  UpperCAmelCase__      				=     24
  UpperCAmelCase__      				=     16
  UpperCAmelCase__      				=     12
  UpperCAmelCase__      				=     8
  UpperCAmelCase__      				=     512
  UpperCAmelCase__      				=     2048
 elif "huge" in model_name:
  UpperCAmelCase__      				=     1280
  UpperCAmelCase__      				=     5120
  UpperCAmelCase__      				=     32
  UpperCAmelCase__      				=     16
  UpperCAmelCase__      				=     12
  UpperCAmelCase__      				=     8
  UpperCAmelCase__      				=     640
  UpperCAmelCase__      				=     2560
 elif "base" not in model_name:
  raise ValueError("""Model name should include either \"small\", \"base\", \"large\", or \"huge\""""							)

def  _UpperCamelCase   (					SCREAMING_SNAKE_CASE__      :   Optional[int]							):
 '''simple docstring'''





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

 return name

def  _UpperCamelCase   (					SCREAMING_SNAKE_CASE__      :   Dict							,    SCREAMING_SNAKE_CASE__      :   Union[str, Any]							):
 '''simple docstring'''





 for key in orig_state_dict.copy().keys():
  UpperCAmelCase__      				=     orig_state_dict.pop(__lowercase							)

  if key.startswith("""encoder."""							):
   UpperCAmelCase__      				=     key.replace("""encoder."""							,    """"""							)

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

 return orig_state_dict

def  _UpperCamelCase   (					):
 '''simple docstring'''





 UpperCAmelCase__      				=     hf_hub_download(
     repo_id="""hf-internal-testing/spaghetti-video"""							,    filename="""eating_spaghetti.npy"""							,    repo_type="""dataset"""							)
 UpperCAmelCase__      				=     np.load(__lowercase							)
 return list(__lowercase							)

def  _UpperCamelCase   (					SCREAMING_SNAKE_CASE__      :   List[Any]							,    SCREAMING_SNAKE_CASE__      :   Tuple							,    SCREAMING_SNAKE_CASE__      :   str							,    SCREAMING_SNAKE_CASE__      :   int							):
 '''simple docstring'''





 UpperCAmelCase__      				=     get_videomae_config(__lowercase							)

 if "finetuned" in model_name:
  UpperCAmelCase__      				=     VideoMAEForVideoClassification(__lowercase							)
 else:
  UpperCAmelCase__      				=     VideoMAEForPreTraining(__lowercase							)

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

 model.load_state_dict(__lowercase							)
 model.eval()

 # verify model on basic input
 UpperCAmelCase__      				=     VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]							,    image_std=[0.5, 0.5, 0.5]							)
 UpperCAmelCase__      				=     prepare_video()
 UpperCAmelCase__      				=     image_processor(__lowercase							,    return_tensors="""pt"""							)

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

 UpperCAmelCase__      				=     model(**__lowercase							)
 UpperCAmelCase__      				=     outputs.logits

 UpperCAmelCase__      				=     [
     """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":
  UpperCAmelCase__      				=     torch.Size([1, 400]							)
  UpperCAmelCase__      				=     torch.tensor([-0.92_91, -0.40_61, -0.93_07]							)
 elif model_name == "videomae-small-finetuned-ssv2":
  UpperCAmelCase__      				=     torch.Size([1, 174]							)
  UpperCAmelCase__      				=     torch.tensor([0.26_71, -0.46_89, -0.82_35]							)
 elif model_name == "videomae-base":
  UpperCAmelCase__      				=     torch.Size([1, 1408, 1536]							)
  UpperCAmelCase__      				=     torch.tensor([[0.77_39, 0.79_68, 0.70_89], [0.67_01, 0.74_87, 0.62_09], [0.42_87, 0.51_58, 0.47_73]]							)
 elif model_name == "videomae-base-short":
  UpperCAmelCase__      				=     torch.Size([1, 1408, 1536]							)
  UpperCAmelCase__      				=     torch.tensor([[0.79_94, 0.96_12, 0.85_08], [0.74_01, 0.89_58, 0.83_02], [0.58_62, 0.74_68, 0.73_25]]							)
  # we verified the loss both for normalized and unnormalized targets for this one
  UpperCAmelCase__      				=     torch.tensor([0.51_42]							) if config.norm_pix_loss else torch.tensor([0.64_69]							)
 elif model_name == "videomae-large":
  UpperCAmelCase__      				=     torch.Size([1, 1408, 1536]							)
  UpperCAmelCase__      				=     torch.tensor([[0.71_49, 0.79_97, 0.69_66], [0.67_68, 0.78_69, 0.69_48], [0.51_39, 0.62_21, 0.56_05]]							)
 elif model_name == "videomae-large-finetuned-kinetics":
  UpperCAmelCase__      				=     torch.Size([1, 400]							)
  UpperCAmelCase__      				=     torch.tensor([0.07_71, 0.00_11, -0.36_25]							)
 elif model_name == "videomae-huge-finetuned-kinetics":
  UpperCAmelCase__      				=     torch.Size([1, 400]							)
  UpperCAmelCase__      				=     torch.tensor([0.24_33, 0.16_32, -0.48_94]							)
 elif model_name == "videomae-base-short-finetuned-kinetics":
  UpperCAmelCase__      				=     torch.Size([1, 400]							)
  UpperCAmelCase__      				=     torch.tensor([0.65_88, 0.09_90, -0.24_93]							)
 elif model_name == "videomae-base-finetuned-kinetics":
  UpperCAmelCase__      				=     torch.Size([1, 400]							)
  UpperCAmelCase__      				=     torch.tensor([0.36_69, -0.06_88, -0.24_21]							)
 elif model_name == "videomae-base-short-ssv2":
  UpperCAmelCase__      				=     torch.Size([1, 1408, 1536]							)
  UpperCAmelCase__      				=     torch.tensor([[0.47_12, 0.52_96, 0.57_86], [0.22_78, 0.27_29, 0.40_26], [0.03_52, 0.07_30, 0.25_06]]							)
 elif model_name == "videomae-base-short-finetuned-ssv2":
  UpperCAmelCase__      				=     torch.Size([1, 174]							)
  UpperCAmelCase__      				=     torch.tensor([-0.05_37, -0.15_39, -0.32_66]							)
 elif model_name == "videomae-base-ssv2":
  UpperCAmelCase__      				=     torch.Size([1, 1408, 1536]							)
  UpperCAmelCase__      				=     torch.tensor([[0.81_31, 0.87_27, 0.85_46], [0.73_66, 0.93_77, 0.88_70], [0.59_35, 0.88_74, 0.85_64]]							)
 elif model_name == "videomae-base-finetuned-ssv2":
  UpperCAmelCase__      				=     torch.Size([1, 174]							)
  UpperCAmelCase__      				=     torch.tensor([0.19_61, -0.83_37, -0.63_89]							)
 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]							,    __lowercase							,    atol=1e-4							)
 else:
  print("""Logits:"""							,    logits[0, :3, :3]							)
  assert torch.allclose(logits[0, :3, :3]							,    __lowercase							,    atol=1e-4							)
 print("""Logits ok!"""							)

 # verify loss, if applicable
 if model_name == "videomae-base-short":
  UpperCAmelCase__      				=     outputs.loss
  assert torch.allclose(__lowercase							,    __lowercase							,    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(__lowercase							)
  model.save_pretrained(__lowercase							)

 if push_to_hub:
  print("""Pushing to the hub..."""							)
  model.push_to_hub(__lowercase							,    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)








import os
import shutil
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
from datasets import Dataset

from transformers.models.realm.configuration_realm import RealmConfig
from transformers.models.realm.retrieval_realm import _REALM_BLOCK_RECORDS_FILENAME, RealmRetriever
from transformers.models.realm.tokenization_realm import VOCAB_FILES_NAMES, RealmTokenizer


class    __lowercase   (						lowercase_	):







     '''simple docstring'''




     def 		lowerCAmelCase_  (					self							:      int					):


         """simple docstring"""



         __A       =						tempfile.mkdtemp()
         __A       =						5

         # Realm tok
         __A       =						[
             """[UNK]""",
             """[CLS]""",
             """[SEP]""",
             """[PAD]""",
             """[MASK]""",
             """test""",
             """question""",
             """this""",
             """is""",
             """the""",
             """first""",
             """second""",
             """third""",
             """fourth""",
             """fifth""",
             """record""",
             """want""",
             """##want""",
             """##ed""",
             """wa""",
             """un""",
             """runn""",
             """##ing""",
             """,""",
             """low""",
             """lowest""",
         ]
         __A       =						os.path.join(self.tmpdirname			,						"""realm_tokenizer"""					)
         os.makedirs(UpperCamelCase_			,						exist_ok=UpperCamelCase_					)
         __A       =						os.path.join(UpperCamelCase_			,						VOCAB_FILES_NAMES["""vocab_file"""]					)
         with open(self.vocab_file			,						"""w"""			,						encoding="""utf-8"""					) as vocab_writer:
             vocab_writer.write("""""".join([x + """\n""" for x in vocab_tokens]					)					)

         __A       =						os.path.join(self.tmpdirname			,						"""realm_block_records"""					)
         os.makedirs(UpperCamelCase_			,						exist_ok=UpperCamelCase_					)




     def 		lowerCAmelCase_  (					self							:      List[Any]					):


         """simple docstring"""



         return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname			,						"""realm_tokenizer"""					)					)




     def 		lowerCAmelCase_  (					self							:      Union[str, Any]					):


         """simple docstring"""



         shutil.rmtree(self.tmpdirname					)




     def 		lowerCAmelCase_  (					self							:      str					):


         """simple docstring"""



         __A       =						RealmConfig(num_block_records=self.num_block_records					)
         return config




     def 		lowerCAmelCase_  (					self							:      Dict					):


         """simple docstring"""



         __A       =						Dataset.from_dict(
             {
                 """id""": ["""0""", """1"""],
                 """question""": ["""foo""", """bar"""],
                 """answers""": [["""Foo""", """Bar"""], ["""Bar"""]],
             }					)
         return dataset




     def 		lowerCAmelCase_  (					self							:      str					):


         """simple docstring"""



         __A       =						np.array(
             [
                 B"""This is the first record""",
                 B"""This is the second record""",
                 B"""This is the third record""",
                 B"""This is the fourth record""",
                 B"""This is the fifth record""",
                 B"""This is a longer longer longer record""",
             ]			,						dtype=UpperCamelCase_			,						)
         return block_records




     def 		lowerCAmelCase_  (					self							:      List[Any]					):


         """simple docstring"""



         __A       =						RealmRetriever(
             block_records=self.get_dummy_block_records()			,						tokenizer=self.get_tokenizer()			,						)
         return retriever




     def 		lowerCAmelCase_  (					self							:      Dict					):


         """simple docstring"""



         __A       =						self.get_config()
         __A       =						self.get_dummy_retriever()
         __A       =						retriever.tokenizer

         __A       =						np.array([0, 3]			,						dtype="""long"""					)
         __A       =						tokenizer(["""Test question"""]					).input_ids
         __A       =						tokenizer(
             ["""the fourth"""]			,						add_special_tokens=UpperCamelCase_			,						return_token_type_ids=UpperCamelCase_			,						return_attention_mask=UpperCamelCase_			,						).input_ids
         __A       =						config.reader_seq_len

         __A     ,						__A     ,						__A     ,						__A       =						retriever(
             UpperCamelCase_			,						UpperCamelCase_			,						answer_ids=UpperCamelCase_			,						max_length=UpperCamelCase_			,						return_tensors="""np"""					)

         self.assertEqual(len(UpperCamelCase_					)			,						2					)
         self.assertEqual(len(UpperCamelCase_					)			,						2					)
         self.assertEqual(len(UpperCamelCase_					)			,						2					)
         self.assertEqual(concat_inputs.input_ids.shape			,						(2, 10)					)
         self.assertEqual(concat_inputs.attention_mask.shape			,						(2, 10)					)
         self.assertEqual(concat_inputs.token_type_ids.shape			,						(2, 10)					)
         self.assertEqual(concat_inputs.special_tokens_mask.shape			,						(2, 10)					)
         self.assertEqual(
             tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[0]					)			,						["""[CLS]""", """test""", """question""", """[SEP]""", """this""", """is""", """the""", """first""", """record""", """[SEP]"""]			,						)
         self.assertEqual(
             tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[1]					)			,						["""[CLS]""", """test""", """question""", """[SEP]""", """this""", """is""", """the""", """fourth""", """record""", """[SEP]"""]			,						)




     def 		lowerCAmelCase_  (					self							:      str					):


         """simple docstring"""



         __A       =						self.get_config()
         __A       =						self.get_dummy_retriever()
         __A       =						retriever.tokenizer

         __A       =						np.array([0, 3, 5]			,						dtype="""long"""					)
         __A       =						tokenizer(["""Test question"""]					).input_ids
         __A       =						tokenizer(
             ["""the fourth""", """longer longer"""]			,						add_special_tokens=UpperCamelCase_			,						return_token_type_ids=UpperCamelCase_			,						return_attention_mask=UpperCamelCase_			,						).input_ids
         __A       =						config.reader_seq_len

         __A     ,						__A     ,						__A     ,						__A       =						retriever(
             UpperCamelCase_			,						UpperCamelCase_			,						answer_ids=UpperCamelCase_			,						max_length=UpperCamelCase_			,						return_tensors="""np"""					)

         self.assertEqual([False, True, True]			,						UpperCamelCase_					)
         self.assertEqual([[-1, -1, -1], [6, -1, -1], [6, 7, 8]]			,						UpperCamelCase_					)
         self.assertEqual([[-1, -1, -1], [7, -1, -1], [7, 8, 9]]			,						UpperCamelCase_					)





     def 		lowerCAmelCase_  (					self							:      str					):


         """simple docstring"""



         __A       =						self.get_dummy_retriever()
         retriever.save_pretrained(os.path.join(self.tmpdirname			,						"""realm_block_records"""					)					)

         # Test local path
         __A       =						retriever.from_pretrained(os.path.join(self.tmpdirname			,						"""realm_block_records"""					)					)
         self.assertEqual(retriever.block_records[0]			,						B"""This is the first record"""					)

         # Test mocked remote path
         with patch("""transformers.models.realm.retrieval_realm.hf_hub_download"""					) as mock_hf_hub_download:
             __A       =						os.path.join(
                 os.path.join(self.tmpdirname			,						"""realm_block_records"""					)			,						_REALM_BLOCK_RECORDS_FILENAME					)
             __A       =						RealmRetriever.from_pretrained("""google/realm-cc-news-pretrained-openqa"""					)

         self.assertEqual(retriever.block_records[0]			,						B"""This is the first record"""					)




