




"""simple docstring"""



import inspect
import math
import tempfile
import unittest

import numpy as np

from transformers import ViTMAEConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
      import torch
      from torch import nn

      from transformers import ViTMAEForPreTraining, ViTMAEModel
      from transformers.models.vit.modeling_vit import VIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
      from PIL import Image

      from transformers import ViTImageProcessor




class      UpperCAmelCase_  :




    def __init__(   self						:	Optional[int]     ,			snake_case_						:	str     ,			snake_case_						:	Any=13     ,			snake_case_						:	Tuple=30     ,			snake_case_						:	List[Any]=2     ,			snake_case_						:	str=3     ,			snake_case_						:	List[Any]=True     ,			snake_case_						:	List[Any]=True     ,			snake_case_						:	str=32     ,			snake_case_						:	str=5     ,			snake_case_						:	List[Any]=4     ,			snake_case_						:	str=37     ,			snake_case_						:	int="gelu"     ,			snake_case_						:	Any=0.1     ,			snake_case_						:	Union[str, Any]=0.1     ,			snake_case_						:	Union[str, Any]=10     ,			snake_case_						:	Tuple=0.02     ,			snake_case_						:	List[str]=3     ,			snake_case_						:	Any=0.6     ,			snake_case_						:	Dict=None     ,			)				->							Tuple:




     '''simple docstring'''

     A__	=    parent
     A__	=    batch_size
     A__	=    image_size
     A__	=    patch_size
     A__	=    num_channels
     A__	=    is_training
     A__	=    use_labels
     A__	=    hidden_size
     A__	=    num_hidden_layers
     A__	=    num_attention_heads
     A__	=    intermediate_size
     A__	=    hidden_act
     A__	=    hidden_dropout_prob
     A__	=    attention_probs_dropout_prob
     A__	=    type_sequence_label_size
     A__	=    initializer_range
     A__	=    mask_ratio
     A__	=    scope

     # in ViTMAE, the expected sequence length = (num_patches + 1) * (1 - config.mask_ratio), rounded above
     # (we add 1 for the [CLS] token)
     A__	=    (image_size // patch_size) ** 2
     A__	=    int(math.ceil((1 - mask_ratio) * (num_patches + 1)				)				)




    def 							__magic_name__							(   self						:	Optional[int]				)				->							Union[str, Any]:




     '''simple docstring'''

     A__	=    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]				)

     A__	=    None
     if self.use_labels:
      A__	=    ids_tensor([self.batch_size]     ,			self.type_sequence_label_size				)

     A__	=    self.get_config()

     return config, pixel_values, labels




    def 							__magic_name__							(   self						:	int				)				->							Dict:




     '''simple docstring'''

     return ViTMAEConfig(
         image_size=self.image_size     ,			patch_size=self.patch_size     ,			num_channels=self.num_channels     ,			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     ,			is_decoder=_SCREAMING_SNAKE_CASE     ,			initializer_range=self.initializer_range     ,			mask_ratio=self.mask_ratio     ,			)




    def 							__magic_name__							(   self						:	Any     ,			snake_case_						:	List[str]     ,			snake_case_						:	Any     ,			snake_case_						:	str				)				->							Tuple:




     '''simple docstring'''

     A__	=    ViTMAEModel(config=_SCREAMING_SNAKE_CASE				)
     model.to(_SCREAMING_SNAKE_CASE				)
     model.eval()
     A__	=    model(_SCREAMING_SNAKE_CASE				)
     self.parent.assertEqual(result.last_hidden_state.shape     ,			(self.batch_size, self.seq_length, self.hidden_size)				)




    def 							__magic_name__							(   self						:	Any     ,			snake_case_						:	str     ,			snake_case_						:	int     ,			snake_case_						:	Tuple				)				->							Optional[Any]:




     '''simple docstring'''

     A__	=    ViTMAEForPreTraining(_SCREAMING_SNAKE_CASE				)
     model.to(_SCREAMING_SNAKE_CASE				)
     model.eval()
     A__	=    model(_SCREAMING_SNAKE_CASE				)
     A__	=    (self.image_size // self.patch_size) ** 2
     A__	=    self.patch_size**2 * self.num_channels
     self.parent.assertEqual(result.logits.shape     ,			(self.batch_size, num_patches, expected_num_channels)				)

     # test greyscale images
     A__	=    1
     A__	=    ViTMAEForPreTraining(_SCREAMING_SNAKE_CASE				)
     model.to(_SCREAMING_SNAKE_CASE				)
     model.eval()
     A__	=    floats_tensor([self.batch_size, 1, self.image_size, self.image_size]				)
     A__	=    model(_SCREAMING_SNAKE_CASE				)
     A__	=    self.patch_size**2
     self.parent.assertEqual(result.logits.shape     ,			(self.batch_size, num_patches, expected_num_channels)				)




    def 							__magic_name__							(   self						:	Optional[Any]				)				->							int:




     '''simple docstring'''

     A__	=    self.prepare_config_and_inputs()
     A__,				A__,				A__	=    config_and_inputs
     A__	=    {"pixel_values": pixel_values}
     return config, inputs_dict






@require_torch
class      UpperCAmelCase_  (						_UpperCamelCase,       _UpperCamelCase,       unittest.TestCase							):
    lowercase__						 =				(ViTMAEModel, ViTMAEForPreTraining) if is_torch_available() else ()
    lowercase__						 =				{'feature-extraction': ViTMAEModel} if is_torch_available() else {}

    lowercase__						 =				False
    lowercase__						 =				False
    lowercase__						 =				False
    lowercase__						 =				False




    def 							__magic_name__							(   self						:	Any				)				->							Any:




     '''simple docstring'''

     A__	=    ViTMAEModelTester(self				)
     A__	=    ConfigTester(self     ,			config_class=_SCREAMING_SNAKE_CASE     ,			has_text_modality=_SCREAMING_SNAKE_CASE     ,			hidden_size=37				)




    def 							__magic_name__							(   self						:	List[Any]				)				->							Optional[int]:




     '''simple docstring'''

     self.config_tester.run_common_tests()




    @unittest.skip(reason="ViTMAE does not use inputs_embeds"				)
    def 							__magic_name__							(   self						:	Any				)				->							int:




     '''simple docstring'''

     pass




    def 							__magic_name__							(   self						:	Union[str, Any]				)				->							int:




     '''simple docstring'''

     A__,				A__	=    self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
      A__	=    model_class(_SCREAMING_SNAKE_CASE				)
      self.assertIsInstance(model.get_input_embeddings()     ,			(nn.Module)				)
      A__	=    model.get_output_embeddings()
      self.assertTrue(x is None or isinstance(_SCREAMING_SNAKE_CASE     ,			nn.Linear				)				)




    def 							__magic_name__							(   self						:	Union[str, Any]				)				->							int:




     '''simple docstring'''

     A__,				A__	=    self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
      A__	=    model_class(_SCREAMING_SNAKE_CASE				)
      A__	=    inspect.signature(model.forward				)
      # signature.parameters is an OrderedDict => so arg_names order is deterministic
      A__	=    [*signature.parameters.keys()]

      A__	=    ["pixel_values"]
      self.assertListEqual(arg_names[:1]     ,			_SCREAMING_SNAKE_CASE				)




    def 							__magic_name__							(   self						:	Tuple				)				->							Union[str, Any]:




     '''simple docstring'''

     A__	=    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*_SCREAMING_SNAKE_CASE				)




    def 							__magic_name__							(   self						:	int				)				->							Optional[int]:




     '''simple docstring'''

     A__	=    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_pretraining(*_SCREAMING_SNAKE_CASE				)




    def 							__magic_name__							(   self						:	Optional[int]     ,			snake_case_						:	Optional[Any]     ,			snake_case_						:	Dict     ,			snake_case_						:	Any				)				->							int:




     '''simple docstring'''

     np.random.seed(2				)

     A__	=    int((pt_model.config.image_size // pt_model.config.patch_size) ** 2				)
     A__	=    np.random.uniform(size=(self.model_tester.batch_size, num_patches)				)
     A__	=    torch.from_numpy(_SCREAMING_SNAKE_CASE				)

     # Add `noise` argument.
     # PT inputs will be prepared in `super().check_pt_tf_models()` with this added `noise` argument
     A__	=    pt_noise

     super().check_pt_tf_models(_SCREAMING_SNAKE_CASE     ,			_SCREAMING_SNAKE_CASE     ,			_SCREAMING_SNAKE_CASE				)




    def 							__magic_name__							(   self						:	List[str]				)				->							str:




     '''simple docstring'''

     A__,				A__	=    self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
      A__	=    model_class(_SCREAMING_SNAKE_CASE				)
      model.to(_SCREAMING_SNAKE_CASE				)
      model.eval()
      # make random mask reproducible
      torch.manual_seed(2				)
      with torch.no_grad():
       A__	=    model(**self._prepare_for_class(_SCREAMING_SNAKE_CASE     ,			_SCREAMING_SNAKE_CASE				)				)

      A__	=    outputs[0].cpu().numpy()
      A__	=    0

      with tempfile.TemporaryDirectory() as tmpdirname:
       model.save_pretrained(_SCREAMING_SNAKE_CASE				)
       A__	=    model_class.from_pretrained(_SCREAMING_SNAKE_CASE				)
       model.to(_SCREAMING_SNAKE_CASE				)
       # make random mask reproducible
       torch.manual_seed(2				)
       with torch.no_grad():
        A__	=    model(**self._prepare_for_class(_SCREAMING_SNAKE_CASE     ,			_SCREAMING_SNAKE_CASE				)				)

       # Make sure we don't have nans
       A__	=    after_outputs[0].cpu().numpy()
       A__	=    0
       A__	=    np.amax(np.abs(out_a - out_a				)				)
       self.assertLessEqual(_SCREAMING_SNAKE_CASE     ,			1e-5				)




    @unittest.skip(
        reason="ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results."				)
    def 							__magic_name__							(   self						:	Any				)				->							List[str]:




     '''simple docstring'''

     pass




    @unittest.skip(
        reason="ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results."				)
    def 							__magic_name__							(   self						:	Optional[int]				)				->							Dict:




     '''simple docstring'''

     pass




    @unittest.skip(
        reason="ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results."				)
    def 							__magic_name__							(   self						:	Any				)				->							Tuple:




     '''simple docstring'''

     pass




    @unittest.skip(reason="ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load"				)
    def 							__magic_name__							(   self						:	List[str]				)				->							str:




     '''simple docstring'''

     pass




    @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."				)
    def 							__magic_name__							(   self						:	int				)				->							Union[str, Any]:




     '''simple docstring'''

     pass




    @slow
    def 							__magic_name__							(   self						:	Optional[int]				)				->							Dict:




     '''simple docstring'''

     for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
      A__	=    ViTMAEModel.from_pretrained(_SCREAMING_SNAKE_CASE				)
      self.assertIsNotNone(_SCREAMING_SNAKE_CASE				)







def   _SCREAMING_SNAKE_CASE       (							)						->		Dict:
 A__	=    Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"						)
 return image






@require_torch
@require_vision
class      UpperCAmelCase_  (						unittest.TestCase							):




    @cached_property
    def 							__magic_name__							(   self						:	Union[str, Any]				)				->							List[str]:




     '''simple docstring'''

     return ViTImageProcessor.from_pretrained("facebook/vit-mae-base"				) if is_vision_available() else None




    @slow
    def 							__magic_name__							(   self						:	Union[str, Any]				)				->							List[str]:




     '''simple docstring'''

     np.random.seed(2				)

     A__	=    ViTMAEForPreTraining.from_pretrained("facebook/vit-mae-base"				).to(_SCREAMING_SNAKE_CASE				)

     A__	=    self.default_image_processor
     A__	=    prepare_img()
     A__	=    image_processor(images=_SCREAMING_SNAKE_CASE     ,			return_tensors="pt"				).to(_SCREAMING_SNAKE_CASE				)

     # prepare a noise vector that will be also used for testing the TF model
     # (this way we can ensure that the PT and TF models operate on the same inputs)
     A__	=    ViTMAEConfig()
     A__	=    int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2				)
     A__	=    np.random.uniform(size=(1, num_patches)				)

     # forward pass
     with torch.no_grad():
      A__	=    model(**_SCREAMING_SNAKE_CASE     ,			noise=torch.from_numpy(_SCREAMING_SNAKE_CASE				).to(device=_SCREAMING_SNAKE_CASE				)				)

     # verify the logits
     A__	=    torch.Size((1, 196, 768)				)
     self.assertEqual(outputs.logits.shape     ,			_SCREAMING_SNAKE_CASE				)

     A__	=    torch.tensor(
         [[-0.0548, -1.7023, -0.9325], [0.3721, -0.5670, -0.2233], [0.8235, -1.3878, -0.3524]]				)

     self.assertTrue(torch.allclose(outputs.logits[0, :3, :3]     ,			expected_slice.to(_SCREAMING_SNAKE_CASE				)     ,			atol=1e-4				)				)











"""simple docstring"""



import argparse

import torch
from datasets import load_dataset
from donut import DonutModel

from transformers import (
    DonutImageProcessor,
    DonutProcessor,
    DonutSwinConfig,
    DonutSwinModel,
    MBartConfig,
    MBartForCausalLM,
    VisionEncoderDecoderModel,
    XLMRobertaTokenizerFast,
)







def   _SCREAMING_SNAKE_CASE       (							lowercase_						)						->		List[str]:
 A__	=    model.config

 A__	=    DonutSwinConfig(
     image_size=original_config.input_size      ,			patch_size=4      ,			depths=original_config.encoder_layer      ,			num_heads=[4, 8, 16, 32]      ,			window_size=original_config.window_size      ,			embed_dim=1_28      ,			)
 A__	=    MBartConfig(
     is_decoder=lowercase_      ,			is_encoder_decoder=lowercase_      ,			add_cross_attention=lowercase_      ,			decoder_layers=original_config.decoder_layer      ,			max_position_embeddings=original_config.max_position_embeddings      ,			vocab_size=len(
         model.decoder.tokenizer						)      ,			scale_embedding=lowercase_      ,			add_final_layer_norm=lowercase_      ,			)

 return encoder_config, decoder_config







def   _SCREAMING_SNAKE_CASE       (							lowercase_						)						->		int:
 if "encoder.model" in name:
  A__	=    name.replace("encoder.model"      ,			"encoder"						)
 if "decoder.model" in name:
  A__	=    name.replace("decoder.model"      ,			"decoder"						)
 if "patch_embed.proj" in name:
  A__	=    name.replace("patch_embed.proj"      ,			"embeddings.patch_embeddings.projection"						)
 if "patch_embed.norm" in name:
  A__	=    name.replace("patch_embed.norm"      ,			"embeddings.norm"						)
 if name.startswith("encoder"						):
  if "layers" in name:
   A__	=    "encoder." + name
  if "attn.proj" in name:
   A__	=    name.replace("attn.proj"      ,			"attention.output.dense"						)
  if "attn" in name and "mask" not in name:
   A__	=    name.replace("attn"      ,			"attention.self"						)
  if "norm1" in name:
   A__	=    name.replace("norm1"      ,			"layernorm_before"						)
  if "norm2" in name:
   A__	=    name.replace("norm2"      ,			"layernorm_after"						)
  if "mlp.fc1" in name:
   A__	=    name.replace("mlp.fc1"      ,			"intermediate.dense"						)
  if "mlp.fc2" in name:
   A__	=    name.replace("mlp.fc2"      ,			"output.dense"						)

  if name == "encoder.norm.weight":
   A__	=    "encoder.layernorm.weight"
  if name == "encoder.norm.bias":
   A__	=    "encoder.layernorm.bias"

 return name







def   _SCREAMING_SNAKE_CASE       (							lowercase_      ,			lowercase_						)						->		Any:
 for key in orig_state_dict.copy().keys():
  A__	=    orig_state_dict.pop(lowercase_						)

  if "qkv" in key:
   A__	=    key.split("."						)
   A__	=    int(key_split[3]						)
   A__	=    int(key_split[5]						)
   A__	=    model.encoder.encoder.layers[layer_num].blocks[block_num].attention.self.all_head_size

   if "weight" in key:
    A__	=    val[:dim, :]
    A__	=    val[dim : dim * 2, :]
    A__	=    val[-dim:, :]
   else:
    A__	=    val[:dim]
    A__	=    val[dim : dim * 2]
    A__	=    val[-dim:]
  elif "attn_mask" in key or key in ["encoder.model.norm.weight", "encoder.model.norm.bias"]:
   # HuggingFace implementation doesn't use attn_mask buffer
   # and model doesn't use final LayerNorms for the encoder
   pass
  else:
   A__	=    val

 return orig_state_dict







def   _SCREAMING_SNAKE_CASE       (							lowercase_      ,			lowercase_=None      ,			lowercase_=False						)						->		Dict:
 # load original model
 A__	=    DonutModel.from_pretrained(lowercase_						).eval()

 # load HuggingFace model
 A__,				A__	=    get_configs(lowercase_						)
 A__	=    DonutSwinModel(lowercase_						)
 A__	=    MBartForCausalLM(lowercase_						)
 A__	=    VisionEncoderDecoderModel(encoder=lowercase_      ,			decoder=lowercase_						)
 model.eval()

 A__	=    original_model.state_dict()
 A__	=    convert_state_dict(lowercase_      ,			lowercase_						)
 model.load_state_dict(lowercase_						)

 # verify results on scanned document
 A__	=    load_dataset("hf-internal-testing/example-documents"						)
 A__	=    dataset["test"][0]["image"].convert("RGB"						)

 A__	=    XLMRobertaTokenizerFast.from_pretrained(lowercase_      ,			from_slow=lowercase_						)
 A__	=    DonutImageProcessor(
     do_align_long_axis=original_model.config.align_long_axis      ,			size=original_model.config.input_size[::-1]						)
 A__	=    DonutProcessor(lowercase_      ,			lowercase_						)
 A__	=    processor(lowercase_      ,			return_tensors="pt"						).pixel_values

 if model_name == "naver-clova-ix/donut-base-finetuned-docvqa":
  A__	=    "<s_docvqa><s_question>{user_input}</s_question><s_answer>"
  A__	=    "When is the coffee break?"
  A__	=    task_prompt.replace("{user_input}"      ,			lowercase_						)
 elif model_name == "naver-clova-ix/donut-base-finetuned-rvlcdip":
  A__	=    "<s_rvlcdip>"
 elif model_name in [
     "naver-clova-ix/donut-base-finetuned-cord-v1",
     "naver-clova-ix/donut-base-finetuned-cord-v1-2560",
 ]:
  A__	=    "<s_cord>"
 elif model_name == "naver-clova-ix/donut-base-finetuned-cord-v2":
  A__	=    "s_cord-v2>"
 elif model_name == "naver-clova-ix/donut-base-finetuned-zhtrainticket":
  A__	=    "<s_zhtrainticket>"
 elif model_name in ["naver-clova-ix/donut-proto", "naver-clova-ix/donut-base"]:
  # use a random prompt
  A__	=    "hello world"
 else:
  raise ValueError("Model name not supported"						)
 A__	=    original_model.decoder.tokenizer(lowercase_      ,			add_special_tokens=lowercase_      ,			return_tensors="pt"						)[
     "input_ids"
 ]

 A__	=    original_model.encoder.model.patch_embed(lowercase_						)
 A__,				A__	=    model.encoder.embeddings(lowercase_						)
 assert torch.allclose(lowercase_      ,			lowercase_      ,			atol=1E-3						)

 # verify encoder hidden states
 A__	=    original_model.encoder(lowercase_						)
 A__	=    model.encoder(lowercase_						).last_hidden_state
 assert torch.allclose(lowercase_      ,			lowercase_      ,			atol=1E-2						)

 # verify decoder hidden states
 A__	=    original_model(lowercase_      ,			lowercase_      ,			lowercase_						).logits
 A__	=    model(lowercase_      ,			decoder_input_ids=lowercase_						).logits
 assert torch.allclose(lowercase_      ,			lowercase_      ,			atol=1E-3						)
 print("Looks ok!"						)

 if pytorch_dump_folder_path is not None:
  print(f"""Saving model and processor to {pytorch_dump_folder_path}"""						)
  model.save_pretrained(lowercase_						)
  processor.save_pretrained(lowercase_						)

 if push_to_hub:
  model.push_to_hub("nielsr/" + model_name.split("/"						)[-1]      ,			commit_message="Update model"						)
  processor.push_to_hub("nielsr/" + model_name.split("/"						)[-1]      ,			commit_message="Update model"						)


if __name__ == "__main__":
      SCREAMING_SNAKE_CASE    =	argparse.ArgumentParser()
      # Required parameters
      parser.add_argument(
          "--model_name",
          default="naver-clova-ix/donut-base-finetuned-docvqa",
          required=False,
          type=str,
          help="Name of the original model you'd like to convert.",
      )
      parser.add_argument(
          "--pytorch_dump_folder_path",
          default=None,
          required=False,
          type=str,
          help="Path to the output PyTorch model directory.",
      )
      parser.add_argument(
          "--push_to_hub",
          action="store_true",
          help="Whether or not to push the converted model and processor to the 🤗 hub.",
      )

      SCREAMING_SNAKE_CASE    =	parser.parse_args()
      convert_donut_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)





