





import itertools
import random
import unittest

import numpy as np

from transformers import ASTFeatureExtractor
from transformers.testing_utils import require_torch, require_torchaudio
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


__lowerCamelCase       :     str   =     random.Random()

if is_torch_available():
    import torch



def 		A_				(     _lowerCAmelCase							,   _lowerCAmelCase=1.0							,   _lowerCAmelCase=None							,   _lowerCAmelCase=None     )						-> Optional[Any]:
 if rng is None:
  UpperCamelCase	:					Optional[int]       	=   global_rng

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

 return values



class 					A__   (       unittest.TestCase			):





      def __init__( self       ,							A_       ,							A_=7       ,							A_=400       ,							A_=2000       ,							A_=1       ,							A_=0.0       ,							A_=1_6000       ,							A_=True       ,							A_=True       ,							):
       '''simple docstring'''




       UpperCamelCase	:					Tuple       	=   parent
       UpperCamelCase	:					List[Any]       	=   batch_size
       UpperCamelCase	:					List[Any]       	=   min_seq_length
       UpperCamelCase	:					List[str]       	=   max_seq_length
       UpperCamelCase	:					int       	=   (self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
       UpperCamelCase	:					Union[str, Any]       	=   feature_size
       UpperCamelCase	:					List[str]       	=   padding_value
       UpperCamelCase	:					Optional[Any]       	=   sampling_rate
       UpperCamelCase	:					List[str]       	=   return_attention_mask
       UpperCamelCase	:					List[Any]       	=   do_normalize





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       return {
           "feature_size": self.feature_size,
           "padding_value": self.padding_value,
           "sampling_rate": self.sampling_rate,
           "return_attention_mask": self.return_attention_mask,
           "do_normalize": self.do_normalize,
       }





      def        __UpperCamelCase( self       ,							A_=False       ,							A_=False   ):
       '''simple docstring'''




       def _flatten(A_   ):
        return list(itertools.chain(*A_   )   )

       if equal_length:
        UpperCamelCase	:					List[str]       	=   floats_list((self.batch_size, self.max_seq_length)   )
       else:
        # make sure that inputs increase in size
        UpperCamelCase	:					Dict       	=   [
            _flatten(floats_list((x, self.feature_size)   )   )
            for x in range(self.min_seq_length       ,							self.max_seq_length       ,							self.seq_length_diff   )
        ]

       if numpify:
        UpperCamelCase	:					Union[str, Any]       	=   [np.asarray(A_   ) for x in speech_inputs]

       return speech_inputs



@require_torch
@require_torchaudio
class 					A__   (       __snake_case       ,   unittest.TestCase			):
      _UpperCAmelCase	:Optional[Any]						    =   ASTFeatureExtractor





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Tuple       	=   ASTFeatureExtractionTester(self   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Optional[Any]       	=   self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
       # create three inputs of length 800, 1000, and 1200
       UpperCamelCase	:					Tuple       	=   [floats_list((1, x)   )[0] for x in range(800       ,							1400       ,							200   )]
       UpperCamelCase	:					Dict       	=   [np.asarray(A_   ) for speech_input in speech_inputs]

       # Test not batched input
       UpperCamelCase	:					Dict       	=   feat_extract(speech_inputs[0]       ,							return_tensors="np"   ).input_values
       UpperCamelCase	:					Union[str, Any]       	=   feat_extract(np_speech_inputs[0]       ,							return_tensors="np"   ).input_values
       self.assertTrue(np.allclose(A_       ,							A_       ,							atol=1e-3   )   )

       # Test batched
       UpperCamelCase	:					Any       	=   feat_extract(A_       ,							padding=A_       ,							return_tensors="np"   ).input_values
       UpperCamelCase	:					Any       	=   feat_extract(A_       ,							padding=A_       ,							return_tensors="np"   ).input_values
       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.
       UpperCamelCase	:					Dict       	=   [floats_list((1, x)   )[0] for x in (800, 800, 800)]
       UpperCamelCase	:					int       	=   np.asarray(A_   )
       UpperCamelCase	:					Any       	=   feat_extract(A_       ,							return_tensors="np"   ).input_values
       UpperCamelCase	:					List[str]       	=   feat_extract(A_       ,							return_tensors="np"   ).input_values
       for enc_seq_a, enc_seq_a in zip(A_       ,							A_   ):
        self.assertTrue(np.allclose(A_       ,							A_       ,							atol=1e-3   )   )





      @require_torch
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       import torch

       UpperCamelCase	:					List[Any]       	=   self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
       UpperCamelCase	:					int       	=   np.random.rand(100   ).astype(np.floataa   )
       UpperCamelCase	:					str       	=   np_speech_inputs.tolist()

       for inputs in [py_speech_inputs, np_speech_inputs]:
        UpperCamelCase	:					List[Any]       	=   feature_extractor.pad([{"input_values": inputs}]       ,							return_tensors="np"   )
        self.assertTrue(np_processed.input_values.dtype == np.floataa   )
        UpperCamelCase	:					List[str]       	=   feature_extractor.pad([{"input_values": inputs}]       ,							return_tensors="pt"   )
        self.assertTrue(pt_processed.input_values.dtype == torch.floataa   )





      def        __UpperCamelCase( self       ,							A_   ):
       '''simple docstring'''




       from datasets import load_dataset

       UpperCamelCase	:					Dict       	=   load_dataset("hf-internal-testing/librispeech_asr_dummy"       ,							"clean"       ,							split="validation"   )
       # automatic decoding with librispeech
       UpperCamelCase	:					Any       	=   ds.sort("id"   ).select(range(A_   )   )[:num_samples]["audio"]

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





      @require_torch
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Any       	=   torch.tensor(
           [-0.98_94, -1.27_76, -0.90_66, -1.27_76, -0.93_49, -1.26_09, -1.03_86, -1.27_76,
            -1.15_61, -1.27_76, -1.20_52, -1.27_23, -1.21_90, -1.21_32, -1.27_76, -1.11_33,
            -1.19_53, -1.13_43, -1.15_84, -1.22_03, -1.17_70, -1.24_74, -1.23_81, -1.19_36,
            -0.92_70, -0.83_17, -0.80_49, -0.77_06, -0.75_65, -0.78_69]   )
       # fmt: on

       UpperCamelCase	:					List[Any]       	=   self._load_datasamples(1   )
       UpperCamelCase	:					Tuple       	=   ASTFeatureExtractor()
       UpperCamelCase	:					str       	=   feature_extractor(A_       ,							return_tensors="pt"   ).input_values
       self.assertEquals(input_values.shape       ,							(1, 1024, 128)   )
       self.assertTrue(torch.allclose(input_values[0, 0, :30]       ,							A_       ,							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       )		->       Optional[Any]:
					"""simple docstring"""





					a					     =   model.config

					a					     =   DonutSwinConfig(
					    image_size=original_config.input_size ,     patch_size=4 ,     depths=original_config.encoder_layer ,     num_heads=[4, 8, 16, 32] ,     window_size=original_config.window_size ,     embed_dim=128 ,     )
					a					     =   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       :    List[Any]       )		->       Union[str, Any]:
					"""simple docstring"""





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

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

					return name






def     a(      A       :    Union[str, Any] ,     A       :    Tuple       )		->       List[Any]:
					"""simple docstring"""





					for key in orig_state_dict.copy().keys():
										a					     =   orig_state_dict.pop(A       )

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

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

					return orig_state_dict






def     a(      A       :    List[Any] ,     A       :    Tuple=None ,     A       :    List[Any]=False       )		->       Optional[int]:
					"""simple docstring"""





					a					     =   DonutModel.from_pretrained(A       ).eval()

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

					a					     =   original_model.state_dict()
					a					     =   convert_state_dict(A ,     A       )
					model.load_state_dict(A       )

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

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

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

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

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

					# verify decoder hidden states
					a					     =   original_model(A ,     A ,     A       ).logits
					a					     =   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: Tuple      			=		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: Optional[Any]      			=		parser.parse_args()
				convert_donut_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)
