
import unittest

from transformers import 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
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import (
      OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
      OpenAIGPTConfig,
      OpenAIGPTDoubleHeadsModel,
      OpenAIGPTForSequenceClassification,
      OpenAIGPTLMHeadModel,
      OpenAIGPTModel,
  )

class __snake_case		:
    def __init__(  self		,snake_case		,snake_case=13		,snake_case=7		,snake_case=True		,snake_case=True		,snake_case=True		,snake_case=99		,snake_case=32		,snake_case=5		,snake_case=4		,snake_case=37		,snake_case="gelu"		,snake_case=0.1		,snake_case=0.1		,snake_case=512		,snake_case=16		,snake_case=2		,snake_case=0.02		,snake_case=3		,snake_case=4		,snake_case=None		,):



           '''simple docstring'''







           lowercase							:     str       	=      parent
           lowercase							:     List[str]       	=      batch_size
           lowercase							:     Dict       	=      seq_length
           lowercase							:     Tuple       	=      is_training
           lowercase							:     Any       	=      use_token_type_ids
           lowercase							:     str       	=      use_labels
           lowercase							:     List[str]       	=      vocab_size
           lowercase							:     Union[str, Any]       	=      hidden_size
           lowercase							:     Union[str, Any]       	=      num_hidden_layers
           lowercase							:     Optional[Any]       	=      num_attention_heads
           lowercase							:     Tuple       	=      intermediate_size
           lowercase							:     Any       	=      hidden_act
           lowercase							:     Any       	=      hidden_dropout_prob
           lowercase							:     Optional[int]       	=      attention_probs_dropout_prob
           lowercase							:     Optional[int]       	=      max_position_embeddings
           lowercase							:     str       	=      type_vocab_size
           lowercase							:     Tuple       	=      type_sequence_label_size
           lowercase							:     Union[str, Any]       	=      initializer_range
           lowercase							:     List[str]       	=      num_labels
           lowercase							:     Any       	=      num_choices
           lowercase							:     Optional[int]       	=      scope
           lowercase							:     Dict       	=      self.vocab_size - 1
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           lowercase							:     int       	=      ids_tensor([self.batch_size, self.seq_length]		,self.vocab_size      )

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

           lowercase							:     Any       	=      None
           lowercase							:     Dict       	=      None
           lowercase							:     Any       	=      None
           if self.use_labels:
                  lowercase							:     Optional[int]       	=      ids_tensor([self.batch_size]		,self.type_sequence_label_size      )
                  lowercase							:     List[str]       	=      ids_tensor([self.batch_size, self.seq_length]		,self.num_labels      )
                  lowercase							:     Union[str, Any]       	=      ids_tensor([self.batch_size]		,self.num_choices      )

           lowercase							:     List[str]       	=      OpenAIGPTConfig(
               vocab_size=self.vocab_size		,n_embd=self.hidden_size		,n_layer=self.num_hidden_layers		,n_head=self.num_attention_heads		,n_positions=self.max_position_embeddings		,pad_token_id=self.pad_token_id		,)

           lowercase							:     str       	=      ids_tensor([self.num_hidden_layers, self.num_attention_heads]		,2      )

           return (
               config,
               input_ids,
               head_mask,
               token_type_ids,
               sequence_labels,
               token_labels,
               choice_labels,
           )
    def       _SCREAMING_SNAKE_CASE       (  self		,snake_case		,snake_case		,snake_case		,snake_case		,*snake_case      ):



           '''simple docstring'''







           lowercase							:     int       	=      OpenAIGPTModel(config=snake_case      )
           model.to(snake_case      )
           model.eval()

           lowercase							:     Dict       	=      model(snake_case		,token_type_ids=snake_case		,head_mask=snake_case      )
           lowercase							:     Union[str, Any]       	=      model(snake_case		,token_type_ids=snake_case      )
           lowercase							:     int       	=      model(snake_case      )

           self.parent.assertEqual(result.last_hidden_state.shape		,(self.batch_size, self.seq_length, self.hidden_size)      )
    def       _SCREAMING_SNAKE_CASE       (  self		,snake_case		,snake_case		,snake_case		,snake_case		,*snake_case      ):



           '''simple docstring'''







           lowercase							:     List[Any]       	=      OpenAIGPTLMHeadModel(snake_case      )
           model.to(snake_case      )
           model.eval()

           lowercase							:     Dict       	=      model(snake_case		,token_type_ids=snake_case		,labels=snake_case      )
           self.parent.assertEqual(result.loss.shape		,()      )
           self.parent.assertEqual(result.logits.shape		,(self.batch_size, self.seq_length, self.vocab_size)      )
    def       _SCREAMING_SNAKE_CASE       (  self		,snake_case		,snake_case		,snake_case		,snake_case		,*snake_case      ):



           '''simple docstring'''







           lowercase							:     List[Any]       	=      OpenAIGPTDoubleHeadsModel(snake_case      )
           model.to(snake_case      )
           model.eval()

           lowercase							:     int       	=      model(snake_case		,token_type_ids=snake_case		,labels=snake_case      )
           self.parent.assertEqual(result.loss.shape		,()      )
           self.parent.assertEqual(result.logits.shape		,(self.batch_size, self.seq_length, self.vocab_size)      )
    def       _SCREAMING_SNAKE_CASE       (  self		,snake_case		,snake_case		,snake_case		,snake_case		,*snake_case      ):



           '''simple docstring'''







           lowercase							:     List[Any]       	=      self.num_labels
           lowercase							:     Any       	=      OpenAIGPTForSequenceClassification(snake_case      )
           model.to(snake_case      )
           model.eval()

           lowercase							:     int       	=      ids_tensor([self.batch_size]		,self.type_sequence_label_size      )
           lowercase							:     List[str]       	=      model(snake_case		,token_type_ids=snake_case		,labels=snake_case      )
           self.parent.assertEqual(result.logits.shape		,(self.batch_size, self.num_labels)      )
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           lowercase							:     List[Any]       	=      self.prepare_config_and_inputs()
           (
               (
               lowercase
           )	,			(
               lowercase
           )	,			(
               lowercase
           )	,			(
               lowercase
           )	,			(
               lowercase
           )	,			(
               lowercase
           )	,			(
               lowercase
           )	,			
           )							:     Any       	=      config_and_inputs
           lowercase							:     int       	=      {
               """input_ids""": input_ids,
               """token_type_ids""": token_type_ids,
               """head_mask""": head_mask,
           }

           return config, inputs_dict






@require_torch
class __snake_case		(      lowerCAmelCase					,	lowerCAmelCase					,	lowerCAmelCase					,	unittest.TestCase     ):
    _a		:						Optional[Any]=				(
        (OpenAIGPTModel, OpenAIGPTLMHeadModel, OpenAIGPTDoubleHeadsModel, OpenAIGPTForSequenceClassification)
        if is_torch_available()
        else ()
    )
    _a		:						int=				(
        (OpenAIGPTLMHeadModel,) if is_torch_available() else ()
    )  # TODO (PVP): Add Double HeadsModel when generate() function is changed accordingly
    _a		:						str=				(
        {
            "feature-extraction": OpenAIGPTModel,
            "text-classification": OpenAIGPTForSequenceClassification,
            "text-generation": OpenAIGPTLMHeadModel,
            "zero-shot": OpenAIGPTForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    def       _SCREAMING_SNAKE_CASE       (  self		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case      ):



           '''simple docstring'''







           if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
                  # Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
                  # `OpenAIGPTConfig` was never used in pipeline tests, either because of a missing checkpoint or because a
                  # tiny config could not be created.
                  return True

           return False
    def       _SCREAMING_SNAKE_CASE       (  self		,snake_case		,snake_case		,snake_case=False      ):



           '''simple docstring'''







           lowercase							:     Tuple       	=      super()._prepare_for_class(snake_case		,snake_case		,return_labels=snake_case      )

           if return_labels:
                  if model_class.__name__ == "OpenAIGPTDoubleHeadsModel":
                         lowercase							:     List[str]       	=      torch.zeros(
                             (self.model_tester.batch_size, self.model_tester.num_choices, self.model_tester.seq_length)		,dtype=torch.long		,device=snake_case		,)
                         lowercase							:     Tuple       	=      inputs_dict["""labels"""]
                         lowercase							:     Dict       	=      inputs_dict["""labels"""]
                         lowercase							:     List[Any]       	=      torch.zeros(
                             (self.model_tester.batch_size, self.model_tester.num_choices)		,dtype=torch.long		,device=snake_case		,)
                         lowercase							:     List[str]       	=      torch.zeros(
                             self.model_tester.batch_size		,dtype=torch.long		,device=snake_case      )
           return inputs_dict
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           lowercase							:     Dict       	=      OpenAIGPTModelTester(self      )
           lowercase							:     str       	=      ConfigTester(self		,config_class=snake_case		,n_embd=37      )
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           self.config_tester.run_common_tests()
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           lowercase							:     int       	=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_openai_gpt_model(*snake_case      )
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           lowercase							:     Tuple       	=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_lm_head_model(*snake_case      )
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           lowercase							:     str       	=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_double_lm_head_model(*snake_case      )
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           lowercase							:     str       	=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_openai_gpt_for_sequence_classification(*snake_case      )
    @slow
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           for model_name in OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  lowercase							:     Any       	=      OpenAIGPTModel.from_pretrained(snake_case      )
                  self.assertIsNotNone(snake_case      )






@require_torch
class __snake_case		(      unittest.TestCase     ):
    @slow
    def       _SCREAMING_SNAKE_CASE       (  self      ):



           '''simple docstring'''







           lowercase							:     int       	=      OpenAIGPTLMHeadModel.from_pretrained("""openai-gpt"""      )
           model.to(snake_case      )
           lowercase							:     List[Any]       	=      torch.tensor([[481, 4735, 544]]		,dtype=torch.long		,device=snake_case      )  # the president is
           lowercase							:     str       	=      [
               481,
               4735,
               544,
               246,
               963,
               870,
               762,
               239,
               244,
               40477,
               244,
               249,
               719,
               881,
               487,
               544,
               240,
               244,
               603,
               481,
           ]  # the president is a very good man. " \n " i\'m sure he is, " said the

           lowercase							:     int       	=      model.generate(snake_case		,do_sample=snake_case      )
           self.assertListEqual(output_ids[0].tolist()		,snake_case      )


import argparse

import requests
import torch
from PIL import Image

from transformers import ViTMAEConfig, ViTMAEForPreTraining, ViTMAEImageProcessor







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

       return name







def 						_snake_case(  SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__				) ->    List[Any]:
       for key in orig_state_dict.copy().keys():
              lowercase							:     List[Any]       	=      orig_state_dict.pop(SCREAMING_SNAKE_CASE__				)

              if "qkv" in key:
                     lowercase							:     int       	=      key.split("""."""				)
                     lowercase							:     List[str]       	=      int(key_split[1]				)
                     if "decoder_blocks" in key:
                            lowercase							:     Tuple       	=      config.decoder_hidden_size
                            lowercase							:     int       	=      """decoder.decoder_layers."""
                            if "weight" in key:
                                   lowercase							:     List[Any]       	=      val[:dim, :]
                                   lowercase							:     Tuple       	=      val[dim : dim * 2, :]
                                   lowercase							:     List[Any]       	=      val[-dim:, :]
                            elif "bias" in key:
                                   lowercase							:     str       	=      val[:dim]
                                   lowercase							:     Dict       	=      val[dim : dim * 2]
                                   lowercase							:     Union[str, Any]       	=      val[-dim:]
                     else:
                            lowercase							:     Tuple       	=      config.hidden_size
                            lowercase							:     Union[str, Any]       	=      """vit.encoder.layer."""
                            if "weight" in key:
                                   lowercase							:     Tuple       	=      val[:dim, :]
                                   lowercase							:     List[str]       	=      val[dim : dim * 2, :]
                                   lowercase							:     Dict       	=      val[-dim:, :]
                            elif "bias" in key:
                                   lowercase							:     Any       	=      val[:dim]
                                   lowercase							:     str       	=      val[dim : dim * 2]
                                   lowercase							:     Union[str, Any]       	=      val[-dim:]

              else:
                     lowercase							:     Union[str, Any]       	=      val

       return orig_state_dict







def 						_snake_case(  SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__				) ->    int:
       lowercase							:     int       	=      ViTMAEConfig()
       if "large" in checkpoint_url:
              lowercase							:     Dict       	=      1_024
              lowercase							:     str       	=      4_096
              lowercase							:     Optional[Any]       	=      24
              lowercase							:     Optional[Any]       	=      16
       elif "huge" in checkpoint_url:
              lowercase							:     int       	=      14
              lowercase							:     List[Any]       	=      1_280
              lowercase							:     int       	=      5_120
              lowercase							:     List[Any]       	=      32
              lowercase							:     Any       	=      16

       lowercase							:     List[str]       	=      ViTMAEForPreTraining(SCREAMING_SNAKE_CASE__				)

       lowercase							:     Optional[Any]       	=      torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE__					,	map_location="""cpu"""				)["""model"""]

       lowercase							:     Tuple       	=      ViTMAEImageProcessor(size=config.image_size				)

       lowercase							:     Optional[int]       	=      convert_state_dict(SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__				)

       model.load_state_dict(SCREAMING_SNAKE_CASE__				)
       model.eval()

       lowercase							:     Union[str, Any]       	=      """https://user-images.githubusercontent.com/11435359/147738734-196fd92f-9260-48d5-ba7e-bf103d29364d.jpg"""

       lowercase							:     Union[str, Any]       	=      Image.open(requests.get(SCREAMING_SNAKE_CASE__					,	stream=SCREAMING_SNAKE_CASE__				).raw				)
       lowercase							:     Optional[Any]       	=      ViTMAEImageProcessor(size=config.image_size				)
       lowercase							:     List[Any]       	=      image_processor(images=SCREAMING_SNAKE_CASE__					,	return_tensors="""pt"""				)

       # forward pass
       torch.manual_seed(2				)
       lowercase							:     int       	=      model(**SCREAMING_SNAKE_CASE__				)
       lowercase							:     str       	=      outputs.logits

       if "large" in checkpoint_url:
              lowercase							:     List[Any]       	=      torch.tensor(
                  [[-0.7309, -0.7128, -1.0169], [-1.0161, -0.9058, -1.1878], [-1.0478, -0.9411, -1.1911]]				)
       elif "huge" in checkpoint_url:
              lowercase							:     Tuple       	=      torch.tensor(
                  [[-1.1599, -0.9199, -1.2221], [-1.1952, -0.9269, -1.2307], [-1.2143, -0.9337, -1.2262]]				)
       else:
              lowercase							:     List[str]       	=      torch.tensor(
                  [[-0.9192, -0.8481, -1.1259], [-1.1349, -1.0034, -1.2599], [-1.1757, -1.0429, -1.2726]]				)

       # verify logits
       assert torch.allclose(logits[0, :3, :3]					,	SCREAMING_SNAKE_CASE__					,	atol=1e-4				)

       print(f"Saving model to {pytorch_dump_folder_path}"				)
       model.save_pretrained(SCREAMING_SNAKE_CASE__				)

       print(f"Saving image processor to {pytorch_dump_folder_path}"				)
       image_processor.save_pretrained(SCREAMING_SNAKE_CASE__				)


if __name__ == "__main__":
  lowercase  :					Union[str, Any]						    =							argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      """--checkpoint_url""",
      default="""https://dl.fbaipublicfiles.com/mae/visualize/mae_visualize_vit_base.pth""",
      type=str,
      help="""URL of the checkpoint you'd like to convert.""",
  )
  parser.add_argument(
      """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model directory."""
  )

  lowercase  :					List[Any]						    =							parser.parse_args()
  convert_vit_mae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)
