




import unittest

from transformers import XLMConfig, 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 (
     XLMForMultipleChoice,
     XLMForQuestionAnswering,
     XLMForQuestionAnsweringSimple,
     XLMForSequenceClassification,
     XLMForTokenClassification,
     XLMModel,
     XLMWithLMHeadModel,
 )
 from transformers.models.xlm.modeling_xlm import XLM_PRETRAINED_MODEL_ARCHIVE_LIST

class  _lowerCamelCase						:

      """simple docstring"""

      def __init__(     self					: Optional[Any] ,       snake_case					: Union[str, Any] ,       snake_case					: Tuple=13 ,       snake_case					: Any=7 ,       snake_case					: str=True ,       snake_case					: Tuple=True ,       snake_case					: Dict=True ,       snake_case					: Union[str, Any]=True ,       snake_case					: List[Any]=True ,       snake_case					: Dict=False ,       snake_case					: Tuple=False ,       snake_case					: List[str]=False ,       snake_case					: Optional[Any]=2 ,       snake_case					: int=99 ,       snake_case					: Any=0 ,       snake_case					: int=32 ,       snake_case					: Optional[int]=5 ,       snake_case					: Any=4 ,       snake_case					: str=0.1 ,       snake_case					: Tuple=0.1 ,       snake_case					: Optional[int]=512 ,       snake_case					: Any=2 ,       snake_case					: Union[str, Any]=0.02 ,       snake_case					: int=2 ,       snake_case					: List[Any]=4 ,       snake_case					: List[Any]="last" ,       snake_case					: Union[str, Any]=True ,       snake_case					: Optional[Any]=None ,       snake_case					: Dict=0 ,       ):
          __UpperCamelCase  		=						parent
          __UpperCamelCase  		=						batch_size
          __UpperCamelCase  		=						seq_length
          __UpperCamelCase  		=						is_training
          __UpperCamelCase  		=						use_input_lengths
          __UpperCamelCase  		=						use_token_type_ids
          __UpperCamelCase  		=						use_labels
          __UpperCamelCase  		=						gelu_activation
          __UpperCamelCase  		=						sinusoidal_embeddings
          __UpperCamelCase  		=						causal
          __UpperCamelCase  		=						asm
          __UpperCamelCase  		=						n_langs
          __UpperCamelCase  		=						vocab_size
          __UpperCamelCase  		=						n_special
          __UpperCamelCase  		=						hidden_size
          __UpperCamelCase  		=						num_hidden_layers
          __UpperCamelCase  		=						num_attention_heads
          __UpperCamelCase  		=						hidden_dropout_prob
          __UpperCamelCase  		=						attention_probs_dropout_prob
          __UpperCamelCase  		=						max_position_embeddings
          __UpperCamelCase  		=						type_sequence_label_size
          __UpperCamelCase  		=						initializer_range
          __UpperCamelCase  		=						num_labels
          __UpperCamelCase  		=						num_choices
          __UpperCamelCase  		=						summary_type
          __UpperCamelCase  		=						use_proj
          __UpperCamelCase  		=						scope
          __UpperCamelCase  		=						bos_token_id







      def       snake_case	(     self					: Any				):
          __UpperCamelCase  		=						ids_tensor([self.batch_size, self.seq_length] ,       self.vocab_size				)
          __UpperCamelCase  		=						random_attention_mask([self.batch_size, self.seq_length]				)

          __UpperCamelCase  		=						None
          if self.use_input_lengths:
              __UpperCamelCase  		=						(
                  ids_tensor([self.batch_size] ,       vocab_size=2				) + self.seq_length - 2
              )  # small variation of seq_length

          __UpperCamelCase  		=						None
          if self.use_token_type_ids:
              __UpperCamelCase  		=						ids_tensor([self.batch_size, self.seq_length] ,       self.n_langs				)

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

          __UpperCamelCase  		=						self.get_config()

          return (
              config,
              input_ids,
              token_type_ids,
              input_lengths,
              sequence_labels,
              token_labels,
              is_impossible_labels,
              choice_labels,
              input_mask,
          )







      def       snake_case	(     self					: str				):
          return XLMConfig(
              vocab_size=self.vocab_size ,       n_special=self.n_special ,       emb_dim=self.hidden_size ,       n_layers=self.num_hidden_layers ,       n_heads=self.num_attention_heads ,       dropout=self.hidden_dropout_prob ,       attention_dropout=self.attention_probs_dropout_prob ,       gelu_activation=self.gelu_activation ,       sinusoidal_embeddings=self.sinusoidal_embeddings ,       asm=self.asm ,       causal=self.causal ,       n_langs=self.n_langs ,       max_position_embeddings=self.max_position_embeddings ,       initializer_range=self.initializer_range ,       summary_type=self.summary_type ,       use_proj=self.use_proj ,       num_labels=self.num_labels ,       bos_token_id=self.bos_token_id ,       )







      def       snake_case	(     self					: List[str] ,       snake_case					: List[Any] ,       snake_case					: Dict ,       snake_case					: List[str] ,       snake_case					: Union[str, Any] ,       snake_case					: List[Any] ,       snake_case					: List[str] ,       snake_case					: int ,       snake_case					: List[Any] ,       snake_case					: Any ,       ):
          __UpperCamelCase  		=						XLMModel(config=snake_case				)
          model.to(snake_case				)
          model.eval()
          __UpperCamelCase  		=						model(snake_case ,       lengths=snake_case ,       langs=snake_case				)
          __UpperCamelCase  		=						model(snake_case ,       langs=snake_case				)
          __UpperCamelCase  		=						model(snake_case				)
          self.parent.assertEqual(result.last_hidden_state.shape ,       (self.batch_size, self.seq_length, self.hidden_size)				)







      def       snake_case	(     self					: List[Any] ,       snake_case					: Optional[int] ,       snake_case					: Optional[Any] ,       snake_case					: List[Any] ,       snake_case					: int ,       snake_case					: Union[str, Any] ,       snake_case					: List[str] ,       snake_case					: int ,       snake_case					: Tuple ,       snake_case					: str ,       ):
          __UpperCamelCase  		=						XLMWithLMHeadModel(snake_case				)
          model.to(snake_case				)
          model.eval()

          __UpperCamelCase  		=						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       snake_case	(     self					: Any ,       snake_case					: List[Any] ,       snake_case					: Any ,       snake_case					: Tuple ,       snake_case					: Any ,       snake_case					: Optional[int] ,       snake_case					: Tuple ,       snake_case					: int ,       snake_case					: str ,       snake_case					: Union[str, Any] ,       ):
          __UpperCamelCase  		=						XLMForQuestionAnsweringSimple(snake_case				)
          model.to(snake_case				)
          model.eval()

          __UpperCamelCase  		=						model(snake_case				)

          __UpperCamelCase  		=						model(snake_case ,       start_positions=snake_case ,       end_positions=snake_case				)
          __UpperCamelCase  		=						outputs
          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       snake_case	(     self					: int ,       snake_case					: Tuple ,       snake_case					: Any ,       snake_case					: List[Any] ,       snake_case					: Union[str, Any] ,       snake_case					: Optional[int] ,       snake_case					: Union[str, Any] ,       snake_case					: str ,       snake_case					: Tuple ,       snake_case					: Dict ,       ):
          __UpperCamelCase  		=						XLMForQuestionAnswering(snake_case				)
          model.to(snake_case				)
          model.eval()

          __UpperCamelCase  		=						model(snake_case				)

          __UpperCamelCase  		=						model(
              snake_case ,       start_positions=snake_case ,       end_positions=snake_case ,       cls_index=snake_case ,       is_impossible=snake_case ,       p_mask=snake_case ,       )

          __UpperCamelCase  		=						model(
              snake_case ,       start_positions=snake_case ,       end_positions=snake_case ,       cls_index=snake_case ,       is_impossible=snake_case ,       )

          ((__UpperCamelCase)		,      )  		=						result_with_labels.to_tuple()

          __UpperCamelCase  		=						model(snake_case ,       start_positions=snake_case ,       end_positions=snake_case				)

          ((__UpperCamelCase)		,      )  		=						result_with_labels.to_tuple()

          self.parent.assertEqual(result_with_labels.loss.shape ,       ()				)
          self.parent.assertEqual(result.start_top_log_probs.shape ,       (self.batch_size, model.config.start_n_top)				)
          self.parent.assertEqual(result.start_top_index.shape ,       (self.batch_size, model.config.start_n_top)				)
          self.parent.assertEqual(
              result.end_top_log_probs.shape ,       (self.batch_size, model.config.start_n_top * model.config.end_n_top)				)
          self.parent.assertEqual(
              result.end_top_index.shape ,       (self.batch_size, model.config.start_n_top * model.config.end_n_top)				)
          self.parent.assertEqual(result.cls_logits.shape ,       (self.batch_size,)				)







      def       snake_case	(     self					: Any ,       snake_case					: List[Any] ,       snake_case					: int ,       snake_case					: List[Any] ,       snake_case					: Union[str, Any] ,       snake_case					: List[Any] ,       snake_case					: Optional[Any] ,       snake_case					: Dict ,       snake_case					: Tuple ,       snake_case					: int ,       ):
          __UpperCamelCase  		=						XLMForSequenceClassification(snake_case				)
          model.to(snake_case				)
          model.eval()

          __UpperCamelCase  		=						model(snake_case				)
          __UpperCamelCase  		=						model(snake_case ,       labels=snake_case				)
          self.parent.assertEqual(result.loss.shape ,       ()				)
          self.parent.assertEqual(result.logits.shape ,       (self.batch_size, self.type_sequence_label_size)				)







      def       snake_case	(     self					: Optional[int] ,       snake_case					: Tuple ,       snake_case					: int ,       snake_case					: List[Any] ,       snake_case					: Optional[int] ,       snake_case					: Optional[int] ,       snake_case					: Optional[int] ,       snake_case					: Any ,       snake_case					: Dict ,       snake_case					: int ,       ):
          __UpperCamelCase  		=						self.num_labels
          __UpperCamelCase  		=						XLMForTokenClassification(snake_case				)
          model.to(snake_case				)
          model.eval()

          __UpperCamelCase  		=						model(snake_case ,       attention_mask=snake_case ,       labels=snake_case				)
          self.parent.assertEqual(result.logits.shape ,       (self.batch_size, self.seq_length, self.num_labels)				)







      def       snake_case	(     self					: Optional[int] ,       snake_case					: Tuple ,       snake_case					: Optional[Any] ,       snake_case					: Any ,       snake_case					: Optional[int] ,       snake_case					: str ,       snake_case					: List[str] ,       snake_case					: Any ,       snake_case					: Optional[Any] ,       snake_case					: List[Any] ,       ):
          __UpperCamelCase  		=						self.num_choices
          __UpperCamelCase  		=						XLMForMultipleChoice(config=snake_case				)
          model.to(snake_case				)
          model.eval()
          __UpperCamelCase  		=						input_ids.unsqueeze(1				).expand(-1 ,       self.num_choices ,       -1				).contiguous()
          __UpperCamelCase  		=						token_type_ids.unsqueeze(1				).expand(-1 ,       self.num_choices ,       -1				).contiguous()
          __UpperCamelCase  		=						input_mask.unsqueeze(1				).expand(-1 ,       self.num_choices ,       -1				).contiguous()
          __UpperCamelCase  		=						model(
              snake_case ,       attention_mask=snake_case ,       token_type_ids=snake_case ,       labels=snake_case ,       )
          self.parent.assertEqual(result.logits.shape ,       (self.batch_size, self.num_choices)				)







      def       snake_case	(     self					: Optional[int]				):
          __UpperCamelCase  		=						self.prepare_config_and_inputs()
          (
              (
              __UpperCamelCase
          )		,      (
              __UpperCamelCase
          )		,      (
              __UpperCamelCase
          )		,      (
              __UpperCamelCase
          )		,      (
              __UpperCamelCase
          )		,      (
              __UpperCamelCase
          )		,      (
              __UpperCamelCase
          )		,      (
              __UpperCamelCase
          )		,      (
              __UpperCamelCase
          )		,      
          )  		=						config_and_inputs
          __UpperCamelCase  		=						{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''lengths''': input_lengths}
          return config, inputs_dict







@require_torch
class  _lowerCamelCase						(			_SCREAMING_SNAKE_CASE							,			_SCREAMING_SNAKE_CASE							,			_SCREAMING_SNAKE_CASE							,			unittest.TestCase   ):

      """simple docstring"""

      lowerCAmelCase__			: Optional[Any]    =					(
          (
              XLMModel,
              XLMWithLMHeadModel,
              XLMForQuestionAnswering,
              XLMForSequenceClassification,
              XLMForQuestionAnsweringSimple,
              XLMForTokenClassification,
              XLMForMultipleChoice,
          )
          if is_torch_available()
          else ()
      )
      lowerCAmelCase__			: List[str]    =					(
          (XLMWithLMHeadModel,) if is_torch_available() else ()
      )  # TODO (PVP): Check other models whether language generation is also applicable
      lowerCAmelCase__			: Tuple    =					(
          {
              "feature-extraction": XLMModel,
              "fill-mask": XLMWithLMHeadModel,
              "question-answering": XLMForQuestionAnsweringSimple,
              "text-classification": XLMForSequenceClassification,
              "text-generation": XLMWithLMHeadModel,
              "token-classification": XLMForTokenClassification,
              "zero-shot": XLMForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )







      def       snake_case	(     self					: Dict ,       snake_case					: List[str] ,       snake_case					: str ,       snake_case					: Tuple ,       snake_case					: List[Any] ,       snake_case					: Any				):
          if (
              pipeline_test_casse_name == "QAPipelineTests"
              and tokenizer_name is not None
              and not tokenizer_name.endswith('''Fast'''				)
          ):
              # `QAPipelineTests` fails for a few models when the slower tokenizer are used.
              # (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
              # TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
              return True

          return False







      def       snake_case	(     self					: Dict ,       snake_case					: List[str] ,       snake_case					: Tuple ,       snake_case					: List[str]=False				):
          __UpperCamelCase  		=						super()._prepare_for_class(snake_case ,       snake_case ,       return_labels=snake_case				)

          if return_labels:
              if model_class.__name__ == "XLMForQuestionAnswering":
                  __UpperCamelCase  		=						torch.zeros(
                      self.model_tester.batch_size ,       dtype=torch.long ,       device=snake_case				)
                  __UpperCamelCase  		=						torch.zeros(
                      self.model_tester.batch_size ,       dtype=torch.long ,       device=snake_case				)

          return inputs_dict







      def       snake_case	(     self					: List[str]				):
          __UpperCamelCase  		=						XLMModelTester(self				)
          __UpperCamelCase  		=						ConfigTester(self ,       config_class=snake_case ,       emb_dim=37				)







      def       snake_case	(     self					: Optional[int]				):
          self.config_tester.run_common_tests()







      def       snake_case	(     self					: Any				):
          __UpperCamelCase  		=						self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_xlm_model(*snake_case				)







      def       snake_case	(     self					: List[str]				):
          __UpperCamelCase  		=						self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_xlm_lm_head(*snake_case				)







      def       snake_case	(     self					: Union[str, Any]				):
          __UpperCamelCase  		=						self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_xlm_simple_qa(*snake_case				)







      def       snake_case	(     self					: Union[str, Any]				):
          __UpperCamelCase  		=						self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_xlm_qa(*snake_case				)







      def       snake_case	(     self					: str				):
          __UpperCamelCase  		=						self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_xlm_sequence_classif(*snake_case				)







      def       snake_case	(     self					: List[Any]				):
          __UpperCamelCase  		=						self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_xlm_token_classif(*snake_case				)







      def       snake_case	(     self					: Any				):
          __UpperCamelCase  		=						self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_xlm_for_multiple_choice(*snake_case				)







      def       snake_case	(     self					: str ,       snake_case					: List[str] ,       snake_case					: Dict ,       snake_case					: Union[str, Any] ,       snake_case					: int ,       snake_case					: List[Any] ,       snake_case					: Any=False ,       snake_case					: Optional[Any]=1				):
          self.assertIsInstance(snake_case ,       snake_case				)
          self.assertListEqual(
              [isinstance(snake_case ,       snake_case				) for iter_attentions in attentions] ,       [True] * len(snake_case				)				)
          self.assertEqual(len(snake_case				) ,       (max_length - min_length) * num_beam_groups				)

          for idx, iter_attentions in enumerate(snake_case				):
              # adds PAD dummy token
              __UpperCamelCase  		=						min_length + idx + 1
              __UpperCamelCase  		=						min_length + idx + 1

              __UpperCamelCase  		=						(
                  batch_size * num_beam_groups,
                  config.num_attention_heads,
                  tgt_len,
                  src_len,
              )
              # check attn size
              self.assertListEqual(
                  [layer_attention.shape for layer_attention in iter_attentions] ,       [expected_shape] * len(snake_case				)				)







      def       snake_case	(     self					: List[str] ,       snake_case					: Union[str, Any] ,       snake_case					: Dict ,       snake_case					: Optional[int] ,       snake_case					: int ,       snake_case					: List[Any] ,       snake_case					: Optional[Any]=False ,       snake_case					: str=1				):
          self.assertIsInstance(snake_case ,       snake_case				)
          self.assertListEqual(
              [isinstance(snake_case ,       snake_case				) for iter_hidden_states in hidden_states] ,       [True] * len(snake_case				) ,       )
          self.assertEqual(len(snake_case				) ,       (max_length - min_length) * num_beam_groups				)

          for idx, iter_hidden_states in enumerate(snake_case				):
              # adds PAD dummy token
              __UpperCamelCase  		=						min_length + idx + 1
              __UpperCamelCase  		=						(batch_size * num_beam_groups, seq_len, config.hidden_size)
              # check hidden size
              self.assertListEqual(
                  [layer_hidden_states.shape for layer_hidden_states in iter_hidden_states] ,       [expected_shape] * len(snake_case				) ,       )
          pass







      @slow
      def       snake_case	(     self					: Any				):
          for model_name in XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              __UpperCamelCase  		=						XLMModel.from_pretrained(snake_case				)
              self.assertIsNotNone(snake_case				)







@require_torch
class  _lowerCamelCase						(			unittest.TestCase   ):

      """simple docstring"""

      @slow
      def       snake_case	(     self					: int				):
          __UpperCamelCase  		=						XLMWithLMHeadModel.from_pretrained('''xlm-mlm-en-2048'''				)
          model.to(snake_case				)
          __UpperCamelCase  		=						torch.tensor([[14, 447]] ,       dtype=torch.long ,       device=snake_case				)  # the president
          __UpperCamelCase  		=						[
              14,
              447,
              14,
              447,
              14,
              447,
              14,
              447,
              14,
              447,
              14,
              447,
              14,
              447,
              14,
              447,
              14,
              447,
              14,
              447,
          ]  # the president the president the president the president the president the president the president the president the president the president
          # TODO(PVP): this and other input_ids I tried for generation give pretty bad results. Not sure why. Model might just not be made for auto-regressive inference
          __UpperCamelCase  		=						model.generate(snake_case ,       do_sample=snake_case				)
          self.assertListEqual(output_ids[0].cpu().numpy().tolist() ,       snake_case				)












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_      )       ->	Optional[int]:

    """simple docstring"""


    __UpperCamelCase  		=						model.config

    __UpperCamelCase  		=						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					,     )
    __UpperCamelCase  		=						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_      )       ->	Any:

    """simple docstring"""


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

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

    return name




def 					__SCREAMING_SNAKE_CASE   (  lowercase_					,     lowercase_      )       ->	int:

    """simple docstring"""


    for key in orig_state_dict.copy().keys():
        __UpperCamelCase  		=						orig_state_dict.pop(lowercase_      )

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

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

    return orig_state_dict




def 					__SCREAMING_SNAKE_CASE   (  lowercase_					,     lowercase_=None					,     lowercase_=False      )       ->	Dict:

    """simple docstring"""


    __UpperCamelCase  		=						DonutModel.from_pretrained(lowercase_      ).eval()

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

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

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

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

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

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

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

    # verify decoder hidden states
    __UpperCamelCase  		=						original_model(lowercase_					,     lowercase_					,     lowercase_      ).logits
    __UpperCamelCase  		=						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__":
 a_   		=     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.",
 )

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






