'''simple docstring'''

import math
import unittest

from transformers import BioGptConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import (
      BioGptForCausalLM,
      BioGptForSequenceClassification,
      BioGptForTokenClassification,
      BioGptModel,
      BioGptTokenizer,
  )
  from transformers.models.biogpt.modeling_biogpt import BIOGPT_PRETRAINED_MODEL_ARCHIVE_LIST



class 			UpperCAmelCase  :





       '''simple docstring'''



       def __init__(						self	,   __lowerCAmelCase	,   __lowerCAmelCase=13	,   __lowerCAmelCase=7	,   __lowerCAmelCase=True	,   __lowerCAmelCase=True	,   __lowerCAmelCase=False	,   __lowerCAmelCase=True	,   __lowerCAmelCase=99	,   __lowerCAmelCase=32	,   __lowerCAmelCase=5	,   __lowerCAmelCase=4	,   __lowerCAmelCase=37	,   __lowerCAmelCase="gelu"	,   __lowerCAmelCase=0.1	,   __lowerCAmelCase=0.1	,   __lowerCAmelCase=512	,   __lowerCAmelCase=16	,   __lowerCAmelCase=2	,   __lowerCAmelCase=0.0_2	,   __lowerCAmelCase=3	,   __lowerCAmelCase=4	,   __lowerCAmelCase=None	,   )  ->  str:
        lowercase__				:    str  =					parent
        lowercase__				:    int  =					batch_size
        lowercase__				:    Dict  =					seq_length
        lowercase__				:    Any  =					is_training
        lowercase__				:    List[str]  =					use_input_mask
        lowercase__				:    Any  =					use_token_type_ids
        lowercase__				:    int  =					use_labels
        lowercase__				:    str  =					vocab_size
        lowercase__				:    Any  =					hidden_size
        lowercase__				:    Optional[Any]  =					num_hidden_layers
        lowercase__				:    Dict  =					num_attention_heads
        lowercase__				:    Dict  =					intermediate_size
        lowercase__				:    Optional[int]  =					hidden_act
        lowercase__				:    int  =					hidden_dropout_prob
        lowercase__				:    Optional[int]  =					attention_probs_dropout_prob
        lowercase__				:    Any  =					max_position_embeddings
        lowercase__				:    List[Any]  =					type_vocab_size
        lowercase__				:    Any  =					type_sequence_label_size
        lowercase__				:    Tuple  =					initializer_range
        lowercase__				:    int  =					num_labels
        lowercase__				:    Optional[int]  =					num_choices
        lowercase__				:    Optional[int]  =					scope

       def _lowerCAmelCase(						self )  ->  str:
        lowercase__				:    Tuple  =					ids_tensor([self.batch_size, self.seq_length]	,   self.vocab_size )

        lowercase__				:    List[str]  =					None
        if self.use_input_mask:
         lowercase__				:    Any  =					random_attention_mask([self.batch_size, self.seq_length] )

        lowercase__				:    int  =					None
        if self.use_token_type_ids:
         lowercase__				:    int  =					ids_tensor([self.batch_size, self.seq_length]	,   self.type_vocab_size )

        lowercase__				:    Tuple  =					None
        lowercase__				:    List[str]  =					None
        lowercase__				:    List[str]  =					None
        if self.use_labels:
         lowercase__				:    Any  =					ids_tensor([self.batch_size]	,   self.type_sequence_label_size )
         lowercase__				:    Dict  =					ids_tensor([self.batch_size, self.seq_length]	,   self.num_labels )
         lowercase__				:    str  =					ids_tensor([self.batch_size]	,   self.num_choices )

        lowercase__				:    Optional[Any]  =					self.get_config()

        return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels

       def _lowerCAmelCase(						self )  ->  Optional[Any]:
        return BioGptConfig(
            vocab_size=self.vocab_size	,   hidden_size=self.hidden_size	,   num_hidden_layers=self.num_hidden_layers	,   num_attention_heads=self.num_attention_heads	,   intermediate_size=self.intermediate_size	,   hidden_act=self.hidden_act	,   hidden_dropout_prob=self.hidden_dropout_prob	,   attention_probs_dropout_prob=self.attention_probs_dropout_prob	,   max_position_embeddings=self.max_position_embeddings	,   type_vocab_size=self.type_vocab_size	,   is_decoder=__lowerCAmelCase	,   initializer_range=self.initializer_range	,   )

       def _lowerCAmelCase(						self	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase )  ->  Any:
        lowercase__				:    List[Any]  =					BioGptModel(config=__lowerCAmelCase )
        model.to(__lowerCAmelCase )
        model.eval()
        lowercase__				:    List[Any]  =					model(__lowerCAmelCase	,   attention_mask=__lowerCAmelCase )
        lowercase__				:    str  =					model(__lowerCAmelCase )
        self.parent.assertEqual(result.last_hidden_state.shape	,   (self.batch_size, self.seq_length, self.hidden_size) )

       def _lowerCAmelCase(						self	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   )  ->  Dict:
        lowercase__				:    List[Any]  =					BioGptForCausalLM(config=__lowerCAmelCase )
        model.to(__lowerCAmelCase )
        model.eval()
        lowercase__				:    Dict  =					model(__lowerCAmelCase	,   attention_mask=__lowerCAmelCase	,   token_type_ids=__lowerCAmelCase	,   labels=__lowerCAmelCase )
        self.parent.assertEqual(result.logits.shape	,   (self.batch_size, self.seq_length, self.vocab_size) )

       def _lowerCAmelCase(						self	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   *__lowerCAmelCase )  ->  Dict:
        lowercase__				:    Optional[Any]  =					BioGptModel(config=__lowerCAmelCase )
        model.to(__lowerCAmelCase )
        model.eval()

        # create attention mask
        lowercase__				:    str  =					torch.ones(input_ids.shape	,   dtype=torch.long	,   device=__lowerCAmelCase )
        lowercase__				:    Optional[Any]  =					self.seq_length // 2
        lowercase__				:    List[Any]  =					0

        # first forward pass
        lowercase__				:    List[str]  =					model(__lowerCAmelCase	,   attention_mask=__lowerCAmelCase ).to_tuple()

        # create hypothetical next token and extent to next_input_ids
        lowercase__				:    Optional[Any]  =					ids_tensor((self.batch_size, 1)	,   config.vocab_size )

        # change a random masked slice from input_ids
        lowercase__				:    Union[str, Any]  =					ids_tensor((1,)	,   __lowerCAmelCase ).item() + 1
        lowercase__				:    List[str]  =					ids_tensor((self.batch_size, 1)	,   config.vocab_size ).squeeze(-1 )
        lowercase__				:    Optional[int]  =					random_other_next_tokens

        # append to next input_ids and attn_mask
        lowercase__				:    Union[str, Any]  =					torch.cat([input_ids, next_tokens]	,   dim=-1 )
        lowercase__				:    Any  =					torch.cat(
            [attn_mask, torch.ones((attn_mask.shape[0], 1)	,   dtype=torch.long	,   device=__lowerCAmelCase )]	,   dim=1	,   )

        # get two different outputs
        lowercase__				:    List[Any]  =					model(__lowerCAmelCase	,   attention_mask=__lowerCAmelCase )["""last_hidden_state"""]
        lowercase__				:    Optional[int]  =					model(__lowerCAmelCase	,   past_key_values=__lowerCAmelCase	,   attention_mask=__lowerCAmelCase )["""last_hidden_state"""]

        # select random slice
        lowercase__				:    List[str]  =					ids_tensor((1,)	,   output_from_past.shape[-1] ).item()
        lowercase__				:    int  =					output_from_no_past[:, -1, random_slice_idx].detach()
        lowercase__				:    List[Any]  =					output_from_past[:, 0, random_slice_idx].detach()

        # test that outputs are equal for slice
        self.parent.assertTrue(torch.allclose(__lowerCAmelCase	,   __lowerCAmelCase	,   atol=1E-3 ) )

       def _lowerCAmelCase(						self	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   *__lowerCAmelCase )  ->  str:
        lowercase__				:    Optional[int]  =					BioGptModel(config=__lowerCAmelCase ).to(__lowerCAmelCase ).eval()

        lowercase__				:    Union[str, Any]  =					torch.ones(input_ids.shape	,   dtype=torch.long	,   device=__lowerCAmelCase )

        # first forward pass
        lowercase__				:    int  =					model(__lowerCAmelCase	,   attention_mask=__lowerCAmelCase	,   use_cache=__lowerCAmelCase )

        lowercase__				:    int  =					outputs.to_tuple()

        # create hypothetical multiple next token and extent to next_input_ids
        lowercase__				:    Optional[int]  =					ids_tensor((self.batch_size, 3)	,   config.vocab_size )
        lowercase__				:    str  =					ids_tensor((self.batch_size, 3)	,   2 )

        # append to next input_ids and
        lowercase__				:    Union[str, Any]  =					torch.cat([input_ids, next_tokens]	,   dim=-1 )
        lowercase__				:    int  =					torch.cat([attention_mask, next_attn_mask]	,   dim=-1 )

        lowercase__				:    List[str]  =					model(__lowerCAmelCase	,   attention_mask=__lowerCAmelCase )["""last_hidden_state"""]
        lowercase__				:    Tuple  =					model(__lowerCAmelCase	,   attention_mask=__lowerCAmelCase	,   past_key_values=__lowerCAmelCase )[
            """last_hidden_state"""
        ]

        # select random slice
        lowercase__				:    List[str]  =					ids_tensor((1,)	,   output_from_past.shape[-1] ).item()
        lowercase__				:    List[str]  =					output_from_no_past[:, -3:, random_slice_idx].detach()
        lowercase__				:    Any  =					output_from_past[:, :, random_slice_idx].detach()

        self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1] )

        # test that outputs are equal for slice
        self.parent.assertTrue(torch.allclose(__lowerCAmelCase	,   __lowerCAmelCase	,   atol=1E-3 ) )

       def _lowerCAmelCase(						self	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   *__lowerCAmelCase	,   __lowerCAmelCase=False )  ->  Dict:
        lowercase__				:    Union[str, Any]  =					BioGptForCausalLM(__lowerCAmelCase )
        model.to(__lowerCAmelCase )
        if gradient_checkpointing:
         model.gradient_checkpointing_enable()

        lowercase__				:    Dict  =					model(__lowerCAmelCase	,   labels=__lowerCAmelCase )
        self.parent.assertEqual(result.loss.shape	,   () )
        self.parent.assertEqual(result.logits.shape	,   (self.batch_size, self.seq_length, self.vocab_size) )
        result.loss.backward()

       def _lowerCAmelCase(						self	,   __lowerCAmelCase	,   *__lowerCAmelCase )  ->  Optional[Any]:
        lowercase__				:    Union[str, Any]  =					BioGptModel(__lowerCAmelCase )
        lowercase__				:    str  =					model.config.initializer_range / math.sqrt(2 * model.config.num_hidden_layers )
        for key in model.state_dict().keys():
         if "c_proj" in key and "weight" in key:
          self.parent.assertLessEqual(abs(torch.std(model.state_dict()[key] ) - model_std )	,   0.0_0_1 )
          self.parent.assertLessEqual(abs(torch.mean(model.state_dict()[key] ) - 0.0 )	,   0.0_1 )

       def _lowerCAmelCase(						self	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   *__lowerCAmelCase )  ->  int:
        lowercase__				:    Union[str, Any]  =					self.num_labels
        lowercase__				:    Optional[int]  =					BioGptForTokenClassification(__lowerCAmelCase )
        model.to(__lowerCAmelCase )
        model.eval()
        lowercase__				:    Dict  =					model(__lowerCAmelCase	,   attention_mask=__lowerCAmelCase	,   token_type_ids=__lowerCAmelCase )
        self.parent.assertEqual(result.logits.shape	,   (self.batch_size, self.seq_length, self.num_labels) )


       def _lowerCAmelCase(						self )  ->  str:
        lowercase__				:    List[Any]  =					self.prepare_config_and_inputs()
        (
            lowercase__    
        )				:    Dict  =					config_and_inputs
        lowercase__				:    Tuple  =					{"""input_ids""": input_ids, """attention_mask""": input_mask}
        return config, inputs_dict



@require_torch
class 			UpperCAmelCase  (						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						unittest.TestCase					):





       '''simple docstring'''



       SCREAMING_SNAKE_CASE							=    (
           (BioGptModel, BioGptForCausalLM, BioGptForSequenceClassification, BioGptForTokenClassification)
           if is_torch_available()
           else ()
       )
       SCREAMING_SNAKE_CASE							=    (BioGptForCausalLM,) if is_torch_available() else ()
       SCREAMING_SNAKE_CASE							=    (
           {
               'feature-extraction': BioGptModel,
               'text-classification': BioGptForSequenceClassification,
               'text-generation': BioGptForCausalLM,
               'token-classification': BioGptForTokenClassification,
               'zero-shot': BioGptForSequenceClassification,
           }
           if is_torch_available()
           else {}
       )
       SCREAMING_SNAKE_CASE							=    False

       def _lowerCAmelCase(						self )  ->  List[Any]:
        lowercase__				:    Tuple  =					BioGptModelTester(self )
        lowercase__				:    List[Any]  =					ConfigTester(self	,   config_class=__lowerCAmelCase	,   hidden_size=37 )

       def _lowerCAmelCase(						self )  ->  Union[str, Any]:
        self.config_tester.run_common_tests()

       def _lowerCAmelCase(						self )  ->  Union[str, Any]:
        lowercase__				:    Optional[Any]  =					self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__lowerCAmelCase )

       def _lowerCAmelCase(						self )  ->  List[Any]:
        lowercase__				:    Tuple  =					self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
         lowercase__				:    int  =					type
         self.model_tester.create_and_check_model(*__lowerCAmelCase )

       def _lowerCAmelCase(						self )  ->  Union[str, Any]:
        lowercase__				:    Union[str, Any]  =					self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_biogpt_model_attention_mask_past(*__lowerCAmelCase )

       def _lowerCAmelCase(						self )  ->  str:
        lowercase__				:    Dict  =					self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_forward_and_backwards(*__lowerCAmelCase	,   gradient_checkpointing=__lowerCAmelCase )

       def _lowerCAmelCase(						self )  ->  List[str]:
        lowercase__				:    Optional[int]  =					self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_biogpt_model_past_large_inputs(*__lowerCAmelCase )

       def _lowerCAmelCase(						self )  ->  Optional[int]:
        lowercase__				:    int  =					self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_biogpt_weight_initialization(*__lowerCAmelCase )

       def _lowerCAmelCase(						self )  ->  List[Any]:
        lowercase__				:    List[str]  =					self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_biogpt_for_token_classification(*__lowerCAmelCase )

       @slow
       def _lowerCAmelCase(						self )  ->  List[str]:
        lowercase__				:    str  =					BioGptForCausalLM.from_pretrained('''microsoft/biogpt''' )
        model.to(__lowerCAmelCase )
        lowercase__				:    List[str]  =					BioGptTokenizer.from_pretrained('''microsoft/biogpt''' )

        lowercase__				:    List[str]  =					"""left"""

        # Define PAD Token = EOS Token = 50256
        lowercase__				:    Any  =					tokenizer.eos_token
        lowercase__				:    Dict  =					model.config.eos_token_id

        # use different length sentences to test batching
        lowercase__				:    List[Any]  =					[
            """Hello, my dog is a little""",
            """Today, I""",
        ]

        lowercase__				:    List[str]  =					tokenizer(__lowerCAmelCase	,   return_tensors='''pt'''	,   padding=__lowerCAmelCase )
        lowercase__				:    List[str]  =					inputs["""input_ids"""].to(__lowerCAmelCase )

        lowercase__				:    List[Any]  =					model.generate(
            input_ids=__lowerCAmelCase	,   attention_mask=inputs['''attention_mask'''].to(__lowerCAmelCase )	,   )

        lowercase__				:    List[str]  =					tokenizer(sentences[0]	,   return_tensors='''pt''' ).input_ids.to(__lowerCAmelCase )
        lowercase__				:    List[Any]  =					model.generate(input_ids=__lowerCAmelCase )

        lowercase__				:    List[str]  =					inputs_non_padded.shape[-1] - inputs["""attention_mask"""][-1].long().sum().cpu().item()
        lowercase__				:    Union[str, Any]  =					tokenizer(sentences[1]	,   return_tensors='''pt''' ).input_ids.to(__lowerCAmelCase )
        lowercase__				:    Any  =					model.generate(input_ids=__lowerCAmelCase	,   max_length=model.config.max_length - num_paddings )

        lowercase__				:    List[str]  =					tokenizer.batch_decode(__lowerCAmelCase	,   skip_special_tokens=__lowerCAmelCase )
        lowercase__				:    str  =					tokenizer.decode(output_non_padded[0]	,   skip_special_tokens=__lowerCAmelCase )
        lowercase__				:    List[str]  =					tokenizer.decode(output_padded[0]	,   skip_special_tokens=__lowerCAmelCase )

        lowercase__				:    Union[str, Any]  =					[
            """Hello, my dog is a little bit bigger than a little bit.""",
            """Today, I have a good idea of how to use the information""",
        ]
        self.assertListEqual(__lowerCAmelCase	,   __lowerCAmelCase )
        self.assertListEqual(__lowerCAmelCase	,   [non_padded_sentence, padded_sentence] )

       @slow
       def _lowerCAmelCase(						self )  ->  str:
        for model_name in BIOGPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         lowercase__				:    int  =					BioGptModel.from_pretrained(__lowerCAmelCase )
         self.assertIsNotNone(__lowerCAmelCase )

       def _lowerCAmelCase(						self )  ->  Any:
        lowercase__				:    Tuple  =					self.model_tester.prepare_config_and_inputs_for_common()
        lowercase__				:    List[Any]  =					3
        lowercase__				:    Union[str, Any]  =					input_dict["""input_ids"""]
        lowercase__				:    List[Any]  =					input_ids.ne(1 ).to(__lowerCAmelCase )
        lowercase__				:    str  =					ids_tensor([self.model_tester.batch_size]	,   self.model_tester.type_sequence_label_size )
        lowercase__				:    Union[str, Any]  =					BioGptForSequenceClassification(__lowerCAmelCase )
        model.to(__lowerCAmelCase )
        model.eval()
        lowercase__				:    Tuple  =					model(__lowerCAmelCase	,   attention_mask=__lowerCAmelCase	,   labels=__lowerCAmelCase )
        self.assertEqual(result.logits.shape	,   (self.model_tester.batch_size, self.model_tester.num_labels) )


       def _lowerCAmelCase(						self )  ->  Dict:
        lowercase__				:    Dict  =					self.model_tester.prepare_config_and_inputs_for_common()
        lowercase__				:    Tuple  =					3
        lowercase__				:    Dict  =					"""multi_label_classification"""
        lowercase__				:    List[Any]  =					input_dict["""input_ids"""]
        lowercase__				:    Tuple  =					input_ids.ne(1 ).to(__lowerCAmelCase )
        lowercase__				:    List[str]  =					ids_tensor(
            [self.model_tester.batch_size, config.num_labels]	,   self.model_tester.type_sequence_label_size ).to(torch.float )
        lowercase__				:    Dict  =					BioGptForSequenceClassification(__lowerCAmelCase )
        model.to(__lowerCAmelCase )
        model.eval()
        lowercase__				:    int  =					model(__lowerCAmelCase	,   attention_mask=__lowerCAmelCase	,   labels=__lowerCAmelCase )
        self.assertEqual(result.logits.shape	,   (self.model_tester.batch_size, self.model_tester.num_labels) )



@require_torch
class 			UpperCAmelCase  (						unittest.TestCase					):





       '''simple docstring'''



       @slow
       def _lowerCAmelCase(						self )  ->  List[str]:
        lowercase__				:    Dict  =					BioGptForCausalLM.from_pretrained('''microsoft/biogpt''' )
        lowercase__				:    Optional[Any]  =					torch.tensor([[2, 4805, 9, 656, 21]] )
        lowercase__				:    Dict  =					model(__lowerCAmelCase )[0]

        lowercase__				:    Any  =					42384

        lowercase__				:    Union[str, Any]  =					torch.Size((1, 5, vocab_size) )
        self.assertEqual(output.shape	,   __lowerCAmelCase )

        lowercase__				:    Union[str, Any]  =					torch.tensor(
            [[[-9.5_2_3_6, -9.8_9_1_8, 1_0.4_5_5_7], [-1_1.0_4_6_9, -9.6_4_2_3, 8.1_0_2_2], [-8.8_6_6_4, -7.8_8_2_6, 5.5_3_2_5]]] )

        self.assertTrue(torch.allclose(output[:, :3, :3]	,   __lowerCAmelCase	,   atol=1E-4 ) )


       @slow
       def _lowerCAmelCase(						self )  ->  Dict:
        lowercase__				:    Union[str, Any]  =					BioGptTokenizer.from_pretrained('''microsoft/biogpt''' )
        lowercase__				:    Dict  =					BioGptForCausalLM.from_pretrained('''microsoft/biogpt''' )
        model.to(__lowerCAmelCase )

        torch.manual_seed(0 )
        lowercase__				:    Dict  =					tokenizer('''COVID-19 is'''	,   return_tensors='''pt''' ).to(__lowerCAmelCase )
        lowercase__				:    Optional[int]  =					model.generate(
            **__lowerCAmelCase	,   min_length=100	,   max_length=1024	,   num_beams=5	,   early_stopping=__lowerCAmelCase	,   )
        lowercase__				:    Union[str, Any]  =					tokenizer.decode(output_ids[0]	,   skip_special_tokens=__lowerCAmelCase )

        lowercase__				:    Any  =					(
            """COVID-19 is a global pandemic caused by severe acute respiratory syndrome coronavirus 2 (SARS-CoV-2), the"""
            """ causative agent of coronavirus disease 2019 (COVID-19), which has spread to more than 200 countries and"""
            """ territories, including the United States (US), Canada, Australia, New Zealand, the United Kingdom (UK),"""
            """ and the United States of America (USA), as of March 11, 2020, with more than 800,000 confirmed cases and"""
            """ more than 800,000 deaths."""
        )
        self.assertEqual(__lowerCAmelCase	,   __lowerCAmelCase )






'''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 	a			( lowerCamelCase__    ):
			'''simple docstring'''





			A_	:						Optional[int]	=     VideoMAEConfig()

			set_architecture_configs(lowerCamelCase__						,      lowerCamelCase__    )

			if "finetuned" not in model_name:
						A_	:						Dict	=     False

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

			return config



def 	a			( lowerCamelCase__						,      lowerCamelCase__    ):
			'''simple docstring'''





			if "small" in model_name:
						A_	:						int	=     3_84
						A_	:						Union[str, Any]	=     15_36
						A_	:						List[str]	=     12
						A_	:						Optional[int]	=     16
						A_	:						Any	=     12
						A_	:						int	=     3
						A_	:						Optional[Any]	=     1_92
						A_	:						Union[str, Any]	=     7_68
			elif "large" in model_name:
						A_	:						List[Any]	=     10_24
						A_	:						Optional[Any]	=     40_96
						A_	:						Optional[Any]	=     24
						A_	:						List[str]	=     16
						A_	:						Any	=     12
						A_	:						str	=     8
						A_	:						str	=     5_12
						A_	:						int	=     20_48
			elif "huge" in model_name:
						A_	:						Optional[Any]	=     12_80
						A_	:						str	=     51_20
						A_	:						str	=     32
						A_	:						int	=     16
						A_	:						Any	=     12
						A_	:						Union[str, Any]	=     8
						A_	:						Dict	=     6_40
						A_	:						Optional[Any]	=     25_60
			elif "base" not in model_name:
						raise ValueError("""Model name should include either \"small\", \"base\", \"large\", or \"huge\""""    )



def 	a			( lowerCamelCase__    ):
			'''simple docstring'''





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

			return name



def 	a			( lowerCamelCase__						,      lowerCamelCase__    ):
			'''simple docstring'''





			for key in orig_state_dict.copy().keys():
						A_	:						str	=     orig_state_dict.pop(lowerCamelCase__    )

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

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

			return orig_state_dict



def 	a			( ):
			'''simple docstring'''





			A_	:						List[Any]	=     hf_hub_download(
			    repo_id="""hf-internal-testing/spaghetti-video"""						,      filename="""eating_spaghetti.npy"""						,      repo_type="""dataset"""    )
			A_	:						Optional[Any]	=     np.load(lowerCamelCase__    )
			return list(lowerCamelCase__    )



def 	a			( lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__    ):
			'''simple docstring'''





			A_	:						Any	=     get_videomae_config(lowerCamelCase__    )

			if "finetuned" in model_name:
						A_	:						List[str]	=     VideoMAEForVideoClassification(lowerCamelCase__    )
			else:
						A_	:						Optional[Any]	=     VideoMAEForPreTraining(lowerCamelCase__    )

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

			model.load_state_dict(lowerCamelCase__    )
			model.eval()

			# verify model on basic input
			A_	:						int	=     VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]						,      image_std=[0.5, 0.5, 0.5]    )
			A_	:						Union[str, Any]	=     prepare_video()
			A_	:						str	=     image_processor(lowerCamelCase__						,      return_tensors="""pt"""    )

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

			A_	:						Dict	=     model(**lowerCamelCase__    )
			A_	:						List[Any]	=     outputs.logits

			A_	:						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":
						A_	:						str	=     torch.Size([1, 4_00]    )
						A_	:						Optional[Any]	=     torch.tensor([-0.9_291, -0.4_061, -0.9_307]    )
			elif model_name == "videomae-small-finetuned-ssv2":
						A_	:						str	=     torch.Size([1, 1_74]    )
						A_	:						Union[str, Any]	=     torch.tensor([0.2_671, -0.4_689, -0.8_235]    )
			elif model_name == "videomae-base":
						A_	:						Tuple	=     torch.Size([1, 14_08, 15_36]    )
						A_	:						List[str]	=     torch.tensor([[0.7_739, 0.7_968, 0.7_089], [0.6_701, 0.7_487, 0.6_209], [0.4_287, 0.5_158, 0.4_773]]    )
			elif model_name == "videomae-base-short":
						A_	:						Dict	=     torch.Size([1, 14_08, 15_36]    )
						A_	:						List[str]	=     torch.tensor([[0.7_994, 0.9_612, 0.8_508], [0.7_401, 0.8_958, 0.8_302], [0.5_862, 0.7_468, 0.7_325]]    )
						# we verified the loss both for normalized and unnormalized targets for this one
						A_	:						List[Any]	=     torch.tensor([0.5_142]    ) if config.norm_pix_loss else torch.tensor([0.6_469]    )
			elif model_name == "videomae-large":
						A_	:						str	=     torch.Size([1, 14_08, 15_36]    )
						A_	:						Dict	=     torch.tensor([[0.7_149, 0.7_997, 0.6_966], [0.6_768, 0.7_869, 0.6_948], [0.5_139, 0.6_221, 0.5_605]]    )
			elif model_name == "videomae-large-finetuned-kinetics":
						A_	:						int	=     torch.Size([1, 4_00]    )
						A_	:						Optional[Any]	=     torch.tensor([0.0_771, 0.0_011, -0.3_625]    )
			elif model_name == "videomae-huge-finetuned-kinetics":
						A_	:						Union[str, Any]	=     torch.Size([1, 4_00]    )
						A_	:						Optional[int]	=     torch.tensor([0.2_433, 0.1_632, -0.4_894]    )
			elif model_name == "videomae-base-short-finetuned-kinetics":
						A_	:						List[Any]	=     torch.Size([1, 4_00]    )
						A_	:						Optional[Any]	=     torch.tensor([0.6_588, 0.0_990, -0.2_493]    )
			elif model_name == "videomae-base-finetuned-kinetics":
						A_	:						Union[str, Any]	=     torch.Size([1, 4_00]    )
						A_	:						Tuple	=     torch.tensor([0.3_669, -0.0_688, -0.2_421]    )
			elif model_name == "videomae-base-short-ssv2":
						A_	:						Optional[Any]	=     torch.Size([1, 14_08, 15_36]    )
						A_	:						List[Any]	=     torch.tensor([[0.4_712, 0.5_296, 0.5_786], [0.2_278, 0.2_729, 0.4_026], [0.0_352, 0.0_730, 0.2_506]]    )
			elif model_name == "videomae-base-short-finetuned-ssv2":
						A_	:						Any	=     torch.Size([1, 1_74]    )
						A_	:						Any	=     torch.tensor([-0.0_537, -0.1_539, -0.3_266]    )
			elif model_name == "videomae-base-ssv2":
						A_	:						Dict	=     torch.Size([1, 14_08, 15_36]    )
						A_	:						Dict	=     torch.tensor([[0.8_131, 0.8_727, 0.8_546], [0.7_366, 0.9_377, 0.8_870], [0.5_935, 0.8_874, 0.8_564]]    )
			elif model_name == "videomae-base-finetuned-ssv2":
						A_	:						Any	=     torch.Size([1, 1_74]    )
						A_	:						str	=     torch.tensor([0.1_961, -0.8_337, -0.6_389]    )
			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":
						A_	:						Optional[int]	=     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							:Tuple				    =					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							:Union[str, Any]				    =					parser.parse_args()
						convert_videomae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)