



import unittest

from transformers import AlbertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

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 (
      MODEL_FOR_PRETRAINING_MAPPING,
      AlbertForMaskedLM,
      AlbertForMultipleChoice,
      AlbertForPreTraining,
      AlbertForQuestionAnswering,
      AlbertForSequenceClassification,
      AlbertForTokenClassification,
      AlbertModel,
  )
  from transformers.models.albert.modeling_albert import ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST





class 				__lowerCAmelCase  :




  def __init__(       self							,     lowerCAmelCase							,     lowerCAmelCase=13							,     lowerCAmelCase=7							,     lowerCAmelCase=True							,     lowerCAmelCase=True							,     lowerCAmelCase=True							,     lowerCAmelCase=True							,     lowerCAmelCase=99							,     lowerCAmelCase=16							,     lowerCAmelCase=36							,     lowerCAmelCase=6							,     lowerCAmelCase=6							,     lowerCAmelCase=6							,     lowerCAmelCase=37							,     lowerCAmelCase="gelu"							,     lowerCAmelCase=0.1							,     lowerCAmelCase=0.1							,     lowerCAmelCase=512							,     lowerCAmelCase=16							,     lowerCAmelCase=2							,     lowerCAmelCase=0.02							,     lowerCAmelCase=3							,     lowerCAmelCase=4							,     lowerCAmelCase=None							,     )						->				Tuple:
        '''simple docstring'''






        _lowercase			 =parent
        _lowercase			 =batch_size
        _lowercase			 =seq_length
        _lowercase			 =is_training
        _lowercase			 =use_input_mask
        _lowercase			 =use_token_type_ids
        _lowercase			 =use_labels
        _lowercase			 =vocab_size
        _lowercase			 =embedding_size
        _lowercase			 =hidden_size
        _lowercase			 =num_hidden_layers
        _lowercase			 =num_hidden_groups
        _lowercase			 =num_attention_heads
        _lowercase			 =intermediate_size
        _lowercase			 =hidden_act
        _lowercase			 =hidden_dropout_prob
        _lowercase			 =attention_probs_dropout_prob
        _lowercase			 =max_position_embeddings
        _lowercase			 =type_vocab_size
        _lowercase			 =type_sequence_label_size
        _lowercase			 =initializer_range
        _lowercase			 =num_labels
        _lowercase			 =num_choices
        _lowercase			 =scope




  def 						A__  (       self      )						->				Optional[Any]:
        '''simple docstring'''






        _lowercase			 =ids_tensor([self.batch_size, self.seq_length]							,     self.vocab_size      )

        _lowercase			 =None
        if self.use_input_mask:
              _lowercase			 =random_attention_mask([self.batch_size, self.seq_length]      )

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

        _lowercase			 =None
        _lowercase			 =None
        _lowercase			 =None
        if self.use_labels:
              _lowercase			 =ids_tensor([self.batch_size]							,     self.type_sequence_label_size      )
              _lowercase			 =ids_tensor([self.batch_size, self.seq_length]							,     self.num_labels      )
              _lowercase			 =ids_tensor([self.batch_size]							,     self.num_choices      )

        _lowercase			 =self.get_config()

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




  def 						A__  (       self      )						->				int:
        '''simple docstring'''






        return AlbertConfig(
            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							,     initializer_range=self.initializer_range							,     num_hidden_groups=self.num_hidden_groups							,     )




  def 						A__  (       self							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase      )						->				str:
        '''simple docstring'''






        _lowercase			 =AlbertModel(config=__lowerCamelCase      )
        model.to(__lowerCamelCase      )
        model.eval()
        _lowercase			 =model(__lowerCamelCase							,     attention_mask=__lowerCamelCase							,     token_type_ids=__lowerCamelCase      )
        _lowercase			 =model(__lowerCamelCase							,     token_type_ids=__lowerCamelCase      )
        _lowercase			 =model(__lowerCamelCase      )
        self.parent.assertEqual(result.last_hidden_state.shape							,     (self.batch_size, self.seq_length, self.hidden_size)      )
        self.parent.assertEqual(result.pooler_output.shape							,     (self.batch_size, self.hidden_size)      )




  def 						A__  (       self							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase      )						->				Union[str, Any]:
        '''simple docstring'''






        _lowercase			 =AlbertForPreTraining(config=__lowerCamelCase      )
        model.to(__lowerCamelCase      )
        model.eval()
        _lowercase			 =model(
            __lowerCamelCase							,     attention_mask=__lowerCamelCase							,     token_type_ids=__lowerCamelCase							,     labels=__lowerCamelCase							,     sentence_order_label=__lowerCamelCase							,     )
        self.parent.assertEqual(result.prediction_logits.shape							,     (self.batch_size, self.seq_length, self.vocab_size)      )
        self.parent.assertEqual(result.sop_logits.shape							,     (self.batch_size, config.num_labels)      )




  def 						A__  (       self							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase      )						->				Any:
        '''simple docstring'''






        _lowercase			 =AlbertForMaskedLM(config=__lowerCamelCase      )
        model.to(__lowerCamelCase      )
        model.eval()
        _lowercase			 =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 						A__  (       self							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase      )						->				List[Any]:
        '''simple docstring'''






        _lowercase			 =AlbertForQuestionAnswering(config=__lowerCamelCase      )
        model.to(__lowerCamelCase      )
        model.eval()
        _lowercase			 =model(
            __lowerCamelCase							,     attention_mask=__lowerCamelCase							,     token_type_ids=__lowerCamelCase							,     start_positions=__lowerCamelCase							,     end_positions=__lowerCamelCase							,     )
        self.parent.assertEqual(result.start_logits.shape							,     (self.batch_size, self.seq_length)      )
        self.parent.assertEqual(result.end_logits.shape							,     (self.batch_size, self.seq_length)      )




  def 						A__  (       self							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase      )						->				Any:
        '''simple docstring'''






        _lowercase			 =self.num_labels
        _lowercase			 =AlbertForSequenceClassification(__lowerCamelCase      )
        model.to(__lowerCamelCase      )
        model.eval()
        _lowercase			 =model(__lowerCamelCase							,     attention_mask=__lowerCamelCase							,     token_type_ids=__lowerCamelCase							,     labels=__lowerCamelCase      )
        self.parent.assertEqual(result.logits.shape							,     (self.batch_size, self.num_labels)      )




  def 						A__  (       self							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase      )						->				Tuple:
        '''simple docstring'''






        _lowercase			 =self.num_labels
        _lowercase			 =AlbertForTokenClassification(config=__lowerCamelCase      )
        model.to(__lowerCamelCase      )
        model.eval()
        _lowercase			 =model(__lowerCamelCase							,     attention_mask=__lowerCamelCase							,     token_type_ids=__lowerCamelCase							,     labels=__lowerCamelCase      )
        self.parent.assertEqual(result.logits.shape							,     (self.batch_size, self.seq_length, self.num_labels)      )




  def 						A__  (       self							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase      )						->				Optional[Any]:
        '''simple docstring'''






        _lowercase			 =self.num_choices
        _lowercase			 =AlbertForMultipleChoice(config=__lowerCamelCase      )
        model.to(__lowerCamelCase      )
        model.eval()
        _lowercase			 =input_ids.unsqueeze(1      ).expand(-1							,     self.num_choices							,     -1      ).contiguous()
        _lowercase			 =token_type_ids.unsqueeze(1      ).expand(-1							,     self.num_choices							,     -1      ).contiguous()
        _lowercase			 =input_mask.unsqueeze(1      ).expand(-1							,     self.num_choices							,     -1      ).contiguous()
        _lowercase			 =model(
            __lowerCamelCase							,     attention_mask=__lowerCamelCase							,     token_type_ids=__lowerCamelCase							,     labels=__lowerCamelCase							,     )
        self.parent.assertEqual(result.logits.shape							,     (self.batch_size, self.num_choices)      )




  def 						A__  (       self      )						->				Optional[int]:
        '''simple docstring'''






        _lowercase			 =self.prepare_config_and_inputs()
        (
            _lowercase 
        )			 =config_and_inputs
        _lowercase			 ={"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
        return config, inputs_dict





@require_torch
class 				__lowerCAmelCase  (      _A    ,    _A    ,    unittest.TestCase							):
  _a		    =						(
      (
          AlbertModel,
          AlbertForPreTraining,
          AlbertForMaskedLM,
          AlbertForMultipleChoice,
          AlbertForSequenceClassification,
          AlbertForTokenClassification,
          AlbertForQuestionAnswering,
      )
      if is_torch_available()
      else ()
  )
  _a		    =						(
      {
          """feature-extraction""": AlbertModel,
          """fill-mask""": AlbertForMaskedLM,
          """question-answering""": AlbertForQuestionAnswering,
          """text-classification""": AlbertForSequenceClassification,
          """token-classification""": AlbertForTokenClassification,
          """zero-shot""": AlbertForSequenceClassification,
      }
      if is_torch_available()
      else {}
  )
  _a		    =						True




  def 						A__  (       self							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase=False      )						->				List[Any]:
        '''simple docstring'''






        _lowercase			 =super()._prepare_for_class(__lowerCamelCase							,     __lowerCamelCase							,     return_labels=__lowerCamelCase      )

        if return_labels:
              if model_class in get_values(__lowerCamelCase      ):
                    _lowercase			 =torch.zeros(
                        (self.model_tester.batch_size, self.model_tester.seq_length)							,     dtype=torch.long							,     device=__lowerCamelCase      )
                    _lowercase			 =torch.zeros(
                        self.model_tester.batch_size							,     dtype=torch.long							,     device=__lowerCamelCase      )
        return inputs_dict




  def 						A__  (       self      )						->				Optional[int]:
        '''simple docstring'''






        _lowercase			 =AlbertModelTester(self      )
        _lowercase			 =ConfigTester(self							,     config_class=__lowerCamelCase							,     hidden_size=37      )




  def 						A__  (       self      )						->				int:
        '''simple docstring'''






        self.config_tester.run_common_tests()




  def 						A__  (       self      )						->				Dict:
        '''simple docstring'''






        _lowercase			 =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__lowerCamelCase      )




  def 						A__  (       self      )						->				Optional[Any]:
        '''simple docstring'''






        _lowercase			 =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_pretraining(*__lowerCamelCase      )




  def 						A__  (       self      )						->				Optional[int]:
        '''simple docstring'''






        _lowercase			 =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*__lowerCamelCase      )




  def 						A__  (       self      )						->				Dict:
        '''simple docstring'''






        _lowercase			 =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_multiple_choice(*__lowerCamelCase      )




  def 						A__  (       self      )						->				Optional[Any]:
        '''simple docstring'''






        _lowercase			 =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_question_answering(*__lowerCamelCase      )




  def 						A__  (       self      )						->				Union[str, Any]:
        '''simple docstring'''






        _lowercase			 =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_sequence_classification(*__lowerCamelCase      )




  def 						A__  (       self      )						->				Optional[int]:
        '''simple docstring'''






        _lowercase			 =self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
              _lowercase			 =type
              self.model_tester.create_and_check_model(*__lowerCamelCase      )




  @slow
  def 						A__  (       self      )						->				List[str]:
        '''simple docstring'''






        for model_name in ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              _lowercase			 =AlbertModel.from_pretrained(__lowerCamelCase      )
              self.assertIsNotNone(__lowerCamelCase      )





@require_torch
class 				__lowerCAmelCase  (      unittest.TestCase							):




  @slow
  def 						A__  (       self      )						->				Tuple:
        '''simple docstring'''






        _lowercase			 =AlbertModel.from_pretrained('albert-base-v2'      )
        _lowercase			 =torch.tensor([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]]      )
        _lowercase			 =torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]      )
        with torch.no_grad():
              _lowercase			 =model(__lowerCamelCase							,     attention_mask=__lowerCamelCase      )[0]
        _lowercase			 =torch.Size((1, 11, 768)      )
        self.assertEqual(output.shape							,     __lowerCamelCase      )
        _lowercase			 =torch.tensor(
            [[[-0.6513, 1.5035, -0.2766], [-0.6515, 1.5046, -0.2780], [-0.6512, 1.5049, -0.2784]]]      )

        self.assertTrue(torch.allclose(output[:, 1:4, 1:4]							,     __lowerCamelCase							,     atol=1e-4      )      )





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   a   (    A__			:     Tuple  )					->   Any:





      """simple docstring"""



      _lowercase			 =model.config

      _lowercase			 =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=128   ,     )
      _lowercase			 =MBartConfig(
          is_decoder=A__   ,     is_encoder_decoder=A__   ,     add_cross_attention=A__   ,     decoder_layers=original_config.decoder_layer   ,     max_position_embeddings=original_config.max_position_embeddings   ,     vocab_size=len(
              model.decoder.tokenizer  )   ,     scale_embedding=A__   ,     add_final_layer_norm=A__   ,     )

      return encoder_config, decoder_config
def   a   (    A__			:     Union[str, Any]  )					->   str:





      """simple docstring"""



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

            if name == "encoder.norm.weight":
                  _lowercase			 ='encoder.layernorm.weight'
            if name == "encoder.norm.bias":
                  _lowercase			 ='encoder.layernorm.bias'

      return name
def   a   (    A__			:     Optional[Any]   ,     A__			:     List[Any]  )					->   List[str]:





      """simple docstring"""



      for key in orig_state_dict.copy().keys():
            _lowercase			 =orig_state_dict.pop(A__  )

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

                  if "weight" in key:
                        _lowercase			 =val[:dim, :]
                        _lowercase			 =val[dim : dim * 2, :]
                        _lowercase			 =val[-dim:, :]
                  else:
                        _lowercase			 =val[:dim]
                        _lowercase			 =val[dim : dim * 2]
                        _lowercase			 =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:
                  _lowercase			 =val

      return orig_state_dict
def   a   (    A__			:     str   ,     A__			:     List[str]=None   ,     A__			:     List[Any]=False  )					->   List[str]:





      """simple docstring"""



      _lowercase			 =DonutModel.from_pretrained(A__  ).eval()

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

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

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

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

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

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

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

      # verify decoder hidden states
      _lowercase			 =original_model(A__   ,     A__   ,     A__  ).logits
      _lowercase			 =model(A__   ,     decoder_input_ids=A__  ).logits
      assert torch.allclose(A__   ,     A__   ,     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(A__  )
            processor.save_pretrained(A__  )

      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__":
  lowercase_             =     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.',
  )

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


