




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_			(  snake_case_  ):
      _A							:  int    =		model.config

      _A							:  List[str]    =		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,)
      _A							:  str    =		MBartConfig(
          is_decoder=snake_case_,is_encoder_decoder=snake_case_,add_cross_attention=snake_case_,decoder_layers=original_config.decoder_layer,max_position_embeddings=original_config.max_position_embeddings,vocab_size=len(
              model.decoder.tokenizer  ),scale_embedding=snake_case_,add_final_layer_norm=snake_case_,)

      return encoder_config, decoder_config







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

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

      return name







def        lowerCAmelCase_			(  snake_case_,snake_case_  ):
      for key in orig_state_dict.copy().keys():
            _A							:  str    =		orig_state_dict.pop(snake_case_  )

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

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

      return orig_state_dict







def        lowerCAmelCase_			(  snake_case_,snake_case_=None,snake_case_=False  ):
      # load original model
      _A							:  Any    =		DonutModel.from_pretrained(snake_case_  ).eval()

      # load HuggingFace model
      _A    ,     _A							:  Dict    =		get_configs(snake_case_  )
      _A							:  Dict    =		DonutSwinModel(snake_case_  )
      _A							:  str    =		MBartForCausalLM(snake_case_  )
      _A							:  Optional[int]    =		VisionEncoderDecoderModel(encoder=snake_case_,decoder=snake_case_  )
      model.eval()

      _A							:  Dict    =		original_model.state_dict()
      _A							:  Optional[Any]    =		convert_state_dict(snake_case_,snake_case_  )
      model.load_state_dict(snake_case_  )

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

      _A							:  str    =		XLMRobertaTokenizerFast.from_pretrained(snake_case_,from_slow=snake_case_  )
      _A							:  Tuple    =		DonutImageProcessor(
          do_align_long_axis=original_model.config.align_long_axis,size=original_model.config.input_size[::-1]  )
      _A							:  Optional[Any]    =		DonutProcessor(snake_case_,snake_case_  )
      _A							:  List[str]    =		processor(snake_case_,return_tensors="""pt"""  ).pixel_values

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

      _A							:  List[Any]    =		original_model.encoder.model.patch_embed(snake_case_  )
      _A    ,     _A							:  Any    =		model.encoder.embeddings(snake_case_  )
      assert torch.allclose(snake_case_,snake_case_,atol=1e-3  )

      # verify encoder hidden states
      _A							:  Optional[Any]    =		original_model.encoder(snake_case_  )
      _A							:  Union[str, Any]    =		model.encoder(snake_case_  ).last_hidden_state
      assert torch.allclose(snake_case_,snake_case_,atol=1e-2  )

      # verify decoder hidden states
      _A							:  List[str]    =		original_model(snake_case_,snake_case_,snake_case_  ).logits
      _A							:  Union[str, Any]    =		model(snake_case_,decoder_input_ids=snake_case_  ).logits
      assert torch.allclose(snake_case_,snake_case_,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(snake_case_  )
            processor.save_pretrained(snake_case_  )

      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__":
  _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.",
  )

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







import itertools
import os
import random
import tempfile
import unittest

import numpy as np
from datasets import load_dataset

from transformers import is_speech_available
from transformers.testing_utils import check_json_file_has_correct_format, require_torch, require_torchaudio
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_speech_available():
  from transformers import WhisperFeatureExtractor

if is_torch_available():
  import torch

_snake_case            = random.Random()







def        lowerCAmelCase_			(  snake_case_,snake_case_=1.0,snake_case_=None,snake_case_=None  ):
      if rng is None:
            _A							:  str    =		global_rng

      _A							:  List[Any]    =		[]
      for batch_idx in range(shape[0]  ):
            values.append([]  )
            for _ in range(shape[1]  ):
                  values[-1].append(rng.random() * scale  )

      return values







@require_torch
@require_torchaudio
class 		lowercase							(     unittest.TestCase	):

     def __init__(						self	,  _a	,  _a=7	,  _a=400	,  _a=2000	,  _a=10	,  _a=160	,  _a=8	,  _a=0.0	,  _a=4000	,  _a=False	,  _a=True	,  )		->							Optional[int]:
           _A							:  Any    =		parent
           _A							:  List[Any]    =		batch_size
           _A							:  List[Any]    =		min_seq_length
           _A							:  Dict    =		max_seq_length
           _A							:  Optional[Any]    =		(self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
           _A							:  Tuple    =		padding_value
           _A							:  Tuple    =		sampling_rate
           _A							:  str    =		return_attention_mask
           _A							:  Any    =		do_normalize
           _A							:  Union[str, Any]    =		feature_size
           _A							:  List[Any]    =		chunk_length
           _A							:  List[Any]    =		hop_length

     def       a__			(						self       )		->							List[str]:
           return {
               "feature_size": self.feature_size,
               "hop_length": self.hop_length,
               "chunk_length": self.chunk_length,
               "padding_value": self.padding_value,
               "sampling_rate": self.sampling_rate,
               "return_attention_mask": self.return_attention_mask,
               "do_normalize": self.do_normalize,
           }



     def       a__			(						self	,  _a=False	,  _a=False       )		->							List[str]:
           def _flatten(_a       ):
                 return list(itertools.chain(*_a       )       )

           if equal_length:
                 _A							:  Union[str, Any]    =		[floats_list((self.max_seq_length, self.feature_size)       ) for _ in range(self.batch_size       )]
           else:
                 # make sure that inputs increase in size
                 _A							:  int    =		[
                     floats_list((x, self.feature_size)       )
                     for x in range(self.min_seq_length	,  self.max_seq_length	,  self.seq_length_diff       )
                 ]
           if numpify:
                 _A							:  Any    =		[np.asarray(_a       ) for x in speech_inputs]
           return speech_inputs







@require_torch
@require_torchaudio
class 		lowercase							(     UpperCamelCase__,unittest.TestCase	):
     _a											=		WhisperFeatureExtractor if is_speech_available() else None

     def       a__			(						self       )		->							Tuple:
           _A							:  Optional[int]    =		WhisperFeatureExtractionTester(self       )

     def       a__			(						self       )		->							Optional[Any]:
           _A							:  List[Any]    =		self.feature_extraction_class(**self.feat_extract_dict       )

           with tempfile.TemporaryDirectory() as tmpdirname:
                 _A							:  List[str]    =		feat_extract_first.save_pretrained(_a       )[0]
                 check_json_file_has_correct_format(_a       )
                 _A							:  Optional[int]    =		self.feature_extraction_class.from_pretrained(_a       )

           _A							:  Tuple    =		feat_extract_first.to_dict()
           _A							:  List[Any]    =		feat_extract_second.to_dict()
           _A							:  List[Any]    =		feat_extract_first.mel_filters
           _A							:  Union[str, Any]    =		feat_extract_second.mel_filters
           self.assertTrue(np.allclose(_a	,  _a       )       )
           self.assertEqual(_a	,  _a       )

     def       a__			(						self       )		->							Dict:
           _A							:  Any    =		self.feature_extraction_class(**self.feat_extract_dict       )

           with tempfile.TemporaryDirectory() as tmpdirname:
                 _A							:  Dict    =		os.path.join(_a	,  """feat_extract.json"""       )
                 feat_extract_first.to_json_file(_a       )
                 _A							:  Optional[int]    =		self.feature_extraction_class.from_json_file(_a       )

           _A							:  str    =		feat_extract_first.to_dict()
           _A							:  Any    =		feat_extract_second.to_dict()
           _A							:  Union[str, Any]    =		feat_extract_first.mel_filters
           _A							:  Union[str, Any]    =		feat_extract_second.mel_filters
           self.assertTrue(np.allclose(_a	,  _a       )       )
           self.assertEqual(_a	,  _a       )

     def       a__			(						self       )		->							Union[str, Any]:
           # Tests that all call wrap to encode_plus and batch_encode_plus
           _A							:  Optional[int]    =		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()       )
           # create three inputs of length 800, 1000, and 1200
           _A							:  Optional[Any]    =		[floats_list((1, x)       )[0] for x in range(800	,  1400	,  200       )]
           _A							:  Any    =		[np.asarray(_a       ) for speech_input in speech_inputs]

           # Test feature size
           _A							:  Dict    =		feature_extractor(_a	,  padding="""max_length"""	,  return_tensors="""np"""       ).input_features
           self.assertTrue(input_features.ndim == 3       )
           self.assertTrue(input_features.shape[-1] == feature_extractor.nb_max_frames       )
           self.assertTrue(input_features.shape[-2] == feature_extractor.feature_size       )

           # Test not batched input
           _A							:  List[Any]    =		feature_extractor(speech_inputs[0]	,  return_tensors="""np"""       ).input_features
           _A							:  List[str]    =		feature_extractor(np_speech_inputs[0]	,  return_tensors="""np"""       ).input_features
           self.assertTrue(np.allclose(_a	,  _a	,  atol=1e-3       )       )

           # Test batched
           _A							:  Union[str, Any]    =		feature_extractor(_a	,  return_tensors="""np"""       ).input_features
           _A							:  Tuple    =		feature_extractor(_a	,  return_tensors="""np"""       ).input_features
           for enc_seq_a, enc_seq_a in zip(_a	,  _a       ):
                 self.assertTrue(np.allclose(_a	,  _a	,  atol=1e-3       )       )

           # Test 2-D numpy arrays are batched.
           _A							:  List[str]    =		[floats_list((1, x)       )[0] for x in (800, 800, 800)]
           _A							:  Any    =		np.asarray(_a       )
           _A							:  Union[str, Any]    =		feature_extractor(_a	,  return_tensors="""np"""       ).input_features
           _A							:  int    =		feature_extractor(_a	,  return_tensors="""np"""       ).input_features
           for enc_seq_a, enc_seq_a in zip(_a	,  _a       ):
                 self.assertTrue(np.allclose(_a	,  _a	,  atol=1e-3       )       )

           # Test truncation required
           _A							:  List[Any]    =		[floats_list((1, x)       )[0] for x in range(200	,  (feature_extractor.n_samples + 500)	,  200       )]
           _A							:  Union[str, Any]    =		[np.asarray(_a       ) for speech_input in speech_inputs]

           _A							:  Tuple    =		[x[: feature_extractor.n_samples] for x in speech_inputs]
           _A							:  Union[str, Any]    =		[np.asarray(_a       ) for speech_input in speech_inputs_truncated]

           _A							:  Optional[int]    =		feature_extractor(_a	,  return_tensors="""np"""       ).input_features
           _A							:  List[Any]    =		feature_extractor(_a	,  return_tensors="""np"""       ).input_features
           for enc_seq_a, enc_seq_a in zip(_a	,  _a       ):
                 self.assertTrue(np.allclose(_a	,  _a	,  atol=1e-3       )       )

     def       a__			(						self       )		->							Dict:
           import torch

           _A							:  Dict    =		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()       )
           _A							:  Optional[int]    =		np.random.rand(100	,  32       ).astype(np.floataa       )
           _A							:  str    =		np_speech_inputs.tolist()

           for inputs in [py_speech_inputs, np_speech_inputs]:
                 _A							:  Optional[Any]    =		feature_extractor.pad([{"""input_features""": inputs}]	,  return_tensors="""np"""       )
                 self.assertTrue(np_processed.input_features.dtype == np.floataa       )
                 _A							:  Optional[int]    =		feature_extractor.pad([{"""input_features""": inputs}]	,  return_tensors="""pt"""       )
                 self.assertTrue(pt_processed.input_features.dtype == torch.floataa       )

     def       a__			(						self	,  _a       )		->							Dict:
           _A							:  Optional[Any]    =		load_dataset("""hf-internal-testing/librispeech_asr_dummy"""	,  """clean"""	,  split="""validation"""       )
           # automatic decoding with librispeech
           _A							:  Optional[Any]    =		ds.sort("""id"""       ).select(range(_a       )       )[:num_samples]["""audio"""]

           return [x["array"] for x in speech_samples]

     def       a__			(						self       )		->							Tuple:
           # fmt: off
           _A							:  Dict    =		torch.tensor(
               [
                   0.1193, -0.0946, -0.1098, -0.0196, 0.0225, -0.0690, -0.1736, 0.0951,
                   0.0971, -0.0817, -0.0702, 0.0162, 0.0260, 0.0017, -0.0192, -0.1678,
                   0.0709, -0.1867, -0.0655, -0.0274, -0.0234, -0.1884, -0.0516, -0.0554,
                   -0.0274, -0.1425, -0.1423, 0.0837, 0.0377, -0.0854
               ]       )
           # fmt: on

           _A							:  Dict    =		self._load_datasamples(1       )
           _A							:  Optional[Any]    =		WhisperFeatureExtractor()
           _A							:  Optional[Any]    =		feature_extractor(_a	,  return_tensors="""pt"""       ).input_features
           self.assertEqual(input_features.shape	,  (1, 80, 3000)       )
           self.assertTrue(torch.allclose(input_features[0, 0, :30]	,  _a	,  atol=1e-4       )       )



     def       a__			(						self       )		->							str:
           _A							:  Dict    =		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()       )
           _A							:  str    =		self._load_datasamples(1       )[0]
           _A							:  Union[str, Any]    =		((audio - audio.min()) / (audio.max() - audio.min())) * 6_5535  # Rescale to [0, 65535] to show issue
           _A							:  List[Any]    =		feat_extract.zero_mean_unit_var_norm([audio]	,  attention_mask=_a       )[0]

           self.assertTrue(np.all(np.mean(_a       ) < 1e-3       )       )
           self.assertTrue(np.all(np.abs(np.var(_a       ) - 1       ) < 1e-3       )       )

