


'''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 lowerCamelCase__     (		__lowerCamelCase			:				Union[str, Any]				):







   '''simple docstring'''




   _UpperCAmelCase  :  Dict				       =model.config

   _UpperCAmelCase  :  Dict				       =DonutSwinConfig(
       image_size=original_config.input_size			,       patch_size=4			,       depths=original_config.encoder_layer			,       num_heads=[4, 8, 1_6, 3_2]			,       window_size=original_config.window_size			,       embed_dim=1_2_8			,       )
   _UpperCAmelCase  :  Optional[int]				       =MBartConfig(
       is_decoder=__lowerCamelCase			,       is_encoder_decoder=__lowerCamelCase			,       add_cross_attention=__lowerCamelCase			,       decoder_layers=original_config.decoder_layer			,       max_position_embeddings=original_config.max_position_embeddings			,       vocab_size=len(
           model.decoder.tokenizer				)			,       scale_embedding=__lowerCamelCase			,       add_final_layer_norm=__lowerCamelCase			,       )

   return encoder_config, decoder_config





def lowerCamelCase__     (		__lowerCamelCase			:				List[str]				):







   '''simple docstring'''




   if "encoder.model" in name:
      _UpperCAmelCase  :  Optional[Any]				       =name.replace('encoder.model'			,       'encoder'				)
   if "decoder.model" in name:
      _UpperCAmelCase  :  List[Any]				       =name.replace('decoder.model'			,       'decoder'				)
   if "patch_embed.proj" in name:
      _UpperCAmelCase  :  str				       =name.replace('patch_embed.proj'			,       'embeddings.patch_embeddings.projection'				)
   if "patch_embed.norm" in name:
      _UpperCAmelCase  :  Union[str, Any]				       =name.replace('patch_embed.norm'			,       'embeddings.norm'				)
   if name.startswith('encoder'				):
      if "layers" in name:
         _UpperCAmelCase  :  Tuple				       ='encoder.' + name
      if "attn.proj" in name:
         _UpperCAmelCase  :  int				       =name.replace('attn.proj'			,       'attention.output.dense'				)
      if "attn" in name and "mask" not in name:
         _UpperCAmelCase  :  Tuple				       =name.replace('attn'			,       'attention.self'				)
      if "norm1" in name:
         _UpperCAmelCase  :  str				       =name.replace('norm1'			,       'layernorm_before'				)
      if "norm2" in name:
         _UpperCAmelCase  :  Any				       =name.replace('norm2'			,       'layernorm_after'				)
      if "mlp.fc1" in name:
         _UpperCAmelCase  :  Tuple				       =name.replace('mlp.fc1'			,       'intermediate.dense'				)
      if "mlp.fc2" in name:
         _UpperCAmelCase  :  str				       =name.replace('mlp.fc2'			,       'output.dense'				)

      if name == "encoder.norm.weight":
         _UpperCAmelCase  :  Dict				       ='encoder.layernorm.weight'
      if name == "encoder.norm.bias":
         _UpperCAmelCase  :  Any				       ='encoder.layernorm.bias'

   return name





def lowerCamelCase__     (		__lowerCamelCase			:				List[str]			,       __lowerCamelCase			:				Optional[Any]				):







   '''simple docstring'''




   for key in orig_state_dict.copy().keys():
      _UpperCAmelCase  :  Any				       =orig_state_dict.pop(__lowerCamelCase				)

      if "qkv" in key:
         _UpperCAmelCase  :  Tuple				       =key.split('.'				)
         _UpperCAmelCase  :  Optional[Any]				       =int(key_split[3]				)
         _UpperCAmelCase  :  Dict				       =int(key_split[5]				)
         _UpperCAmelCase  :  str				       =model.encoder.encoder.layers[layer_num].blocks[block_num].attention.self.all_head_size

         if "weight" in key:
            _UpperCAmelCase  :  Tuple				       =val[:dim, :]
            _UpperCAmelCase  :  Any				       =val[dim : dim * 2, :]
            _UpperCAmelCase  :  List[Any]				       =val[-dim:, :]
         else:
            _UpperCAmelCase  :  Any				       =val[:dim]
            _UpperCAmelCase  :  List[str]				       =val[dim : dim * 2]
            _UpperCAmelCase  :  List[Any]				       =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  :  Dict				       =val

   return orig_state_dict





def lowerCamelCase__     (		__lowerCamelCase			:				str			,       __lowerCamelCase			:				int=None			,       __lowerCamelCase			:				str=False				):







   '''simple docstring'''




   _UpperCAmelCase  :  Union[str, Any]				       =DonutModel.from_pretrained(__lowerCamelCase				).eval()

   # load HuggingFace model
   _UpperCAmelCase       ,			_UpperCAmelCase  :  Optional[Any]				       =get_configs(__lowerCamelCase				)
   _UpperCAmelCase  :  Optional[Any]				       =DonutSwinModel(__lowerCamelCase				)
   _UpperCAmelCase  :  Optional[int]				       =MBartForCausalLM(__lowerCamelCase				)
   _UpperCAmelCase  :  str				       =VisionEncoderDecoderModel(encoder=__lowerCamelCase			,       decoder=__lowerCamelCase				)
   model.eval()

   _UpperCAmelCase  :  str				       =original_model.state_dict()
   _UpperCAmelCase  :  Union[str, Any]				       =convert_state_dict(__lowerCamelCase			,       __lowerCamelCase				)
   model.load_state_dict(__lowerCamelCase				)

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

   _UpperCAmelCase  :  Optional[int]				       =XLMRobertaTokenizerFast.from_pretrained(__lowerCamelCase			,       from_slow=__lowerCamelCase				)
   _UpperCAmelCase  :  List[str]				       =DonutImageProcessor(
       do_align_long_axis=original_model.config.align_long_axis			,       size=original_model.config.input_size[::-1]				)
   _UpperCAmelCase  :  str				       =DonutProcessor(__lowerCamelCase			,       __lowerCamelCase				)
   _UpperCAmelCase  :  Tuple				       =processor(__lowerCamelCase			,       return_tensors='pt'				).pixel_values

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

   _UpperCAmelCase  :  List[Any]				       =original_model.encoder.model.patch_embed(__lowerCamelCase				)
   _UpperCAmelCase       ,			_UpperCAmelCase  :  Union[str, Any]				       =model.encoder.embeddings(__lowerCamelCase				)
   assert torch.allclose(__lowerCamelCase			,       __lowerCamelCase			,       atol=1e-3				)

   # verify encoder hidden states
   _UpperCAmelCase  :  Optional[int]				       =original_model.encoder(__lowerCamelCase				)
   _UpperCAmelCase  :  List[str]				       =model.encoder(__lowerCamelCase				).last_hidden_state
   assert torch.allclose(__lowerCamelCase			,       __lowerCamelCase			,       atol=1e-2				)

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

   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)





'''simple docstring'''






import math
import time
from typing import Dict, List, Optional

from torch.utils.data import Dataset

from transformers import SeqaSeqTrainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput, speed_metrics


if is_torch_tpu_available(check_device=False):
 import torch_xla.core.xla_model as xm
 import torch_xla.debug.metrics as met






class    __magic_name__				(						lowerCAmelCase							):






       def __init__(     self    ,				*snake_case    ,				snake_case=None    ,				snake_case=None    ,				**snake_case)		->    Tuple:




          '''simple docstring'''
          super().__init__(*snake_case    ,				**snake_case)
          _UpperCAmelCase  :  Any				       =eval_examples
          _UpperCAmelCase  :  Optional[int]				       =post_process_function






       def      lowerCAmelCase						(     self    ,				snake_case = None    ,				snake_case=None    ,				snake_case = None    ,				snake_case = "eval"    ,				**snake_case    ,				)		->    Dict[str, float]:




          '''simple docstring'''
          _UpperCAmelCase  :  Tuple				       =gen_kwargs.copy()
          _UpperCAmelCase  :  int				       =(
              gen_kwargs['max_length'] if gen_kwargs.get('max_length') is not None else self.args.generation_max_length
          )
          _UpperCAmelCase  :  Dict				       =(
              gen_kwargs['num_beams'] if gen_kwargs.get('num_beams') is not None else self.args.generation_num_beams
          )
          _UpperCAmelCase  :  Dict				       =gen_kwargs

          _UpperCAmelCase  :  Tuple				       =self.eval_dataset if eval_dataset is None else eval_dataset
          _UpperCAmelCase  :  Union[str, Any]				       =self.get_eval_dataloader(snake_case)
          _UpperCAmelCase  :  Tuple				       =self.eval_examples if eval_examples is None else eval_examples

          # Temporarily disable metric computation, we will do it in the loop here.
          _UpperCAmelCase  :  Dict				       =self.compute_metrics
          _UpperCAmelCase  :  Optional[Any]				       =None
          _UpperCAmelCase  :  str				       =time.time()
          _UpperCAmelCase  :  int				       =self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
          try:
             _UpperCAmelCase  :  str				       =eval_loop(
                 snake_case    ,				description='Evaluation'    ,				prediction_loss_only=True if compute_metrics is None else None    ,				ignore_keys=snake_case    ,				metric_key_prefix=snake_case    ,				)
          finally:
             _UpperCAmelCase  :  str				       =compute_metrics
          _UpperCAmelCase  :  Optional[int]				       =self.args.eval_batch_size * self.args.world_size
          if f"{metric_key_prefix}_jit_compilation_time" in output.metrics:
             start_time += output.metrics[f"{metric_key_prefix}_jit_compilation_time"]
          output.metrics.update(
              speed_metrics(
                  snake_case    ,				snake_case    ,				num_samples=output.num_samples    ,				num_steps=math.ceil(output.num_samples / total_batch_size)    ,				))

          if self.post_process_function is not None and self.compute_metrics is not None and self.args.should_save:
             # Only the main node write the results by default
             _UpperCAmelCase  :  Tuple				       =self.post_process_function(snake_case    ,				snake_case    ,				snake_case)
             _UpperCAmelCase  :  Union[str, Any]				       =self.compute_metrics(snake_case)

             # Prefix all keys with metric_key_prefix + '_'
             for key in list(metrics.keys()):
                if not key.startswith(f"{metric_key_prefix}_"):
                   _UpperCAmelCase  :  Any				       =metrics.pop(snake_case)

             metrics.update(output.metrics)
          else:
             _UpperCAmelCase  :  str				       =output.metrics

          if self.args.should_log:
             # Only the main node log the results by default
             self.log(snake_case)

          if self.args.tpu_metrics_debug or self.args.debug:
             # tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
             xm.master_print(met.metrics_report())

          _UpperCAmelCase  :  Any				       =self.callback_handler.on_evaluate(self.args    ,				self.state    ,				self.control    ,				snake_case)
          return metrics






       def      lowerCAmelCase						(     self    ,				snake_case    ,				snake_case    ,				snake_case=None    ,				snake_case = "test"    ,				**snake_case)		->    Dict:




          '''simple docstring'''
          _UpperCAmelCase  :  Optional[int]				       =gen_kwargs.copy()

          _UpperCAmelCase  :  int				       =self.get_test_dataloader(snake_case)

          # Temporarily disable metric computation, we will do it in the loop here.
          _UpperCAmelCase  :  str				       =self.compute_metrics
          _UpperCAmelCase  :  Union[str, Any]				       =None
          _UpperCAmelCase  :  Union[str, Any]				       =time.time()
          _UpperCAmelCase  :  Tuple				       =self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
          try:
             _UpperCAmelCase  :  Any				       =eval_loop(
                 snake_case    ,				description='Prediction'    ,				prediction_loss_only=True if compute_metrics is None else None    ,				ignore_keys=snake_case    ,				metric_key_prefix=snake_case    ,				)
          finally:
             _UpperCAmelCase  :  Any				       =compute_metrics

          _UpperCAmelCase  :  Dict				       =self.args.eval_batch_size * self.args.world_size
          if f"{metric_key_prefix}_jit_compilation_time" in output.metrics:
             start_time += output.metrics[f"{metric_key_prefix}_jit_compilation_time"]
          output.metrics.update(
              speed_metrics(
                  snake_case    ,				snake_case    ,				num_samples=output.num_samples    ,				num_steps=math.ceil(output.num_samples / total_batch_size)    ,				))
          if self.post_process_function is None or self.compute_metrics is None:
             return output

          _UpperCAmelCase  :  Tuple				       =self.post_process_function(snake_case    ,				snake_case    ,				snake_case    ,				'predict')
          _UpperCAmelCase  :  Union[str, Any]				       =self.compute_metrics(snake_case)

          # Prefix all keys with metric_key_prefix + '_'
          for key in list(metrics.keys()):
             if not key.startswith(f"{metric_key_prefix}_"):
                _UpperCAmelCase  :  Any				       =metrics.pop(snake_case)
          metrics.update(output.metrics)
          return PredictionOutput(predictions=predictions.predictions    ,				label_ids=predictions.label_ids    ,				metrics=snake_case)

