




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   (		SCREAMING_SNAKE_CASE_							:       int  ):
					__lowerCAmelCase            = model.config

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

					return encoder_config, decoder_config







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

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

					return name







def 			_a   (		SCREAMING_SNAKE_CASE_							:       str   ,   SCREAMING_SNAKE_CASE_							:       Optional[int]  ):
					for key in orig_state_dict.copy().keys():
										__lowerCAmelCase            = orig_state_dict.pop(SCREAMING_SNAKE_CASE_  )

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

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

					return orig_state_dict







def 			_a   (		SCREAMING_SNAKE_CASE_							:       int   ,   SCREAMING_SNAKE_CASE_							:       Any=None   ,   SCREAMING_SNAKE_CASE_							:       List[str]=False  ):
					# load original model
					__lowerCAmelCase            = DonutModel.from_pretrained(SCREAMING_SNAKE_CASE_  ).eval()

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

					__lowerCAmelCase            = original_model.state_dict()
					__lowerCAmelCase            = convert_state_dict(SCREAMING_SNAKE_CASE_   ,   SCREAMING_SNAKE_CASE_  )
					model.load_state_dict(SCREAMING_SNAKE_CASE_  )

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

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

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

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

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

					# verify decoder hidden states
					__lowerCAmelCase            = original_model(SCREAMING_SNAKE_CASE_   ,   SCREAMING_SNAKE_CASE_   ,   SCREAMING_SNAKE_CASE_  ).logits
					__lowerCAmelCase            = model(SCREAMING_SNAKE_CASE_   ,   decoder_input_ids=SCREAMING_SNAKE_CASE_  ).logits
					assert torch.allclose(SCREAMING_SNAKE_CASE_   ,   SCREAMING_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(SCREAMING_SNAKE_CASE_  )
										processor.save_pretrained(SCREAMING_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__":
					UpperCamelCase__					    =    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.""",
					)

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










import itertools
import random
import unittest

import numpy as np

from transformers import is_speech_available
from transformers.testing_utils import require_torch, require_torchaudio

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_speech_available():
					from transformers import SpeechaTextFeatureExtractor

UpperCamelCase__					    =    random.Random()







def 			_a   (		SCREAMING_SNAKE_CASE_							:       Union[str, Any]   ,   SCREAMING_SNAKE_CASE_							:       Union[str, Any]=1.0   ,   SCREAMING_SNAKE_CASE_							:       Any=None   ,   SCREAMING_SNAKE_CASE_							:       Any=None  ):
					if rng is None:
										__lowerCAmelCase            = global_rng

					__lowerCAmelCase            = []
					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     a__					(      unittest.TestCase   ):






						def __init__(     self	,    _A	,    _A=7	,    _A=4_0_0	,    _A=2_0_0_0	,    _A=2_4	,    _A=2_4	,    _A=0.0	,    _A=1_6_0_0_0	,    _A=True	,    _A=True	,    ):




											"""simple docstring"""



											__lowerCAmelCase            = parent
											__lowerCAmelCase            = batch_size
											__lowerCAmelCase            = min_seq_length
											__lowerCAmelCase            = max_seq_length
											__lowerCAmelCase            = (self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
											__lowerCAmelCase            = feature_size
											__lowerCAmelCase            = num_mel_bins
											__lowerCAmelCase            = padding_value
											__lowerCAmelCase            = sampling_rate
											__lowerCAmelCase            = return_attention_mask
											__lowerCAmelCase            = do_normalize






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



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






						def        __SCREAMING_SNAKE_CASE(     self	,    _A=False	,    _A=False    ):




											"""simple docstring"""



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

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






@require_torch
@require_torchaudio
class     a__					(      snake_case__				,       unittest.TestCase   ):
						_a							:					Optional[int]      =					SpeechaTextFeatureExtractor if is_speech_available() else None






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = SpeechaTextFeatureExtractionTester(self    )






						def        __SCREAMING_SNAKE_CASE(     self	,    _A    ):




											"""simple docstring"""



											self.assertTrue(np.all(np.mean(_A	,    axis=0    ) < 1E-3    )    )
											self.assertTrue(np.all(np.abs(np.var(_A	,    axis=0    ) - 1    ) < 1E-3    )    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()    )
											# create three inputs of length 800, 1000, and 1200
											__lowerCAmelCase            = [floats_list((1, x)    )[0] for x in range(8_0_0	,    1_4_0_0	,    2_0_0    )]
											__lowerCAmelCase            = [np.asarray(_A    ) for speech_input in speech_inputs]

											# Test feature size
											__lowerCAmelCase            = feature_extractor(_A	,    padding=_A	,    return_tensors="np"    ).input_features
											self.assertTrue(input_features.ndim == 3    )
											self.assertTrue(input_features.shape[-1] == feature_extractor.feature_size    )

											# Test not batched input
											__lowerCAmelCase            = feature_extractor(speech_inputs[0]	,    return_tensors="np"    ).input_features
											__lowerCAmelCase            = feature_extractor(np_speech_inputs[0]	,    return_tensors="np"    ).input_features
											self.assertTrue(np.allclose(_A	,    _A	,    atol=1E-3    )    )

											# Test batched
											__lowerCAmelCase            = feature_extractor(_A	,    return_tensors="np"    ).input_features
											__lowerCAmelCase            = 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.
											__lowerCAmelCase            = [floats_list((1, x)    )[0] for x in (8_0_0, 8_0_0, 8_0_0)]
											__lowerCAmelCase            = np.asarray(_A    )
											__lowerCAmelCase            = feature_extractor(_A	,    return_tensors="np"    ).input_features
											__lowerCAmelCase            = 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        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()    )
											__lowerCAmelCase            = [floats_list((1, x)    )[0] for x in range(8_0_0	,    1_4_0_0	,    2_0_0    )]

											__lowerCAmelCase            = ["longest", "max_length", "do_not_pad"]
											__lowerCAmelCase            = [None, 1_6, None]
											for max_length, padding in zip(_A	,    _A    ):
																__lowerCAmelCase            = feature_extractor(
																    _A	,    padding=_A	,    max_length=_A	,    return_attention_mask=_A    )
																__lowerCAmelCase            = inputs.input_features
																__lowerCAmelCase            = inputs.attention_mask
																__lowerCAmelCase            = [np.sum(_A    ) for x in attention_mask]

																self._check_zero_mean_unit_variance(input_features[0][: fbank_feat_lengths[0]]    )
																self._check_zero_mean_unit_variance(input_features[1][: fbank_feat_lengths[1]]    )
																self._check_zero_mean_unit_variance(input_features[2][: fbank_feat_lengths[2]]    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()    )
											__lowerCAmelCase            = [floats_list((1, x)    )[0] for x in range(8_0_0	,    1_4_0_0	,    2_0_0    )]

											__lowerCAmelCase            = ["longest", "max_length", "do_not_pad"]
											__lowerCAmelCase            = [None, 1_6, None]
											for max_length, padding in zip(_A	,    _A    ):
																__lowerCAmelCase            = feature_extractor(
																    _A	,    max_length=_A	,    padding=_A	,    return_tensors="np"	,    return_attention_mask=_A    )
																__lowerCAmelCase            = inputs.input_features
																__lowerCAmelCase            = inputs.attention_mask
																__lowerCAmelCase            = [np.sum(_A    ) for x in attention_mask]

																self._check_zero_mean_unit_variance(input_features[0][: fbank_feat_lengths[0]]    )
																self.assertTrue(input_features[0][fbank_feat_lengths[0] :].sum() < 1E-6    )
																self._check_zero_mean_unit_variance(input_features[1][: fbank_feat_lengths[1]]    )
																self.assertTrue(input_features[0][fbank_feat_lengths[1] :].sum() < 1E-6    )
																self._check_zero_mean_unit_variance(input_features[2][: fbank_feat_lengths[2]]    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()    )
											__lowerCAmelCase            = [floats_list((1, x)    )[0] for x in range(8_0_0	,    1_4_0_0	,    2_0_0    )]
											__lowerCAmelCase            = feature_extractor(
											    _A	,    padding="max_length"	,    max_length=4	,    truncation=_A	,    return_tensors="np"	,    return_attention_mask=_A	,    )
											__lowerCAmelCase            = inputs.input_features
											__lowerCAmelCase            = inputs.attention_mask
											__lowerCAmelCase            = np.sum(attention_mask == 1	,    axis=1    )

											self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]    )
											self._check_zero_mean_unit_variance(input_features[1]    )
											self._check_zero_mean_unit_variance(input_features[2]    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()    )
											__lowerCAmelCase            = [floats_list((1, x)    )[0] for x in range(8_0_0	,    1_4_0_0	,    2_0_0    )]
											__lowerCAmelCase            = feature_extractor(
											    _A	,    padding="longest"	,    max_length=4	,    truncation=_A	,    return_tensors="np"	,    return_attention_mask=_A	,    )
											__lowerCAmelCase            = inputs.input_features
											__lowerCAmelCase            = inputs.attention_mask
											__lowerCAmelCase            = np.sum(attention_mask == 1	,    axis=1    )

											self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]    )
											self._check_zero_mean_unit_variance(input_features[1, : fbank_feat_lengths[1]]    )
											self._check_zero_mean_unit_variance(input_features[2]    )

											# make sure that if max_length < longest -> then pad to max_length
											self.assertEqual(input_features.shape	,    (3, 4, 2_4)    )

											__lowerCAmelCase            = [floats_list((1, x)    )[0] for x in range(8_0_0	,    1_4_0_0	,    2_0_0    )]
											__lowerCAmelCase            = feature_extractor(
											    _A	,    padding="longest"	,    max_length=1_6	,    truncation=_A	,    return_tensors="np"	,    return_attention_mask=_A	,    )
											__lowerCAmelCase            = inputs.input_features
											__lowerCAmelCase            = inputs.attention_mask
											__lowerCAmelCase            = np.sum(attention_mask == 1	,    axis=1    )

											self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]    )
											self._check_zero_mean_unit_variance(input_features[1, : fbank_feat_lengths[1]]    )
											self._check_zero_mean_unit_variance(input_features[2]    )

											# make sure that if max_length < longest -> then pad to max_length
											self.assertEqual(input_features.shape	,    (3, 6, 2_4)    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											import torch

											__lowerCAmelCase            = self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()    )
											__lowerCAmelCase            = np.random.rand(1_0_0	,    3_2    ).astype(np.floataa    )
											__lowerCAmelCase            = np_speech_inputs.tolist()

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






						def        __SCREAMING_SNAKE_CASE(     self	,    _A    ):




											"""simple docstring"""



											from datasets import load_dataset

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

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






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = np.array([
											    -1.57_45, -1.77_13, -1.70_20, -1.60_69, -1.22_50, -1.11_05, -0.90_72, -0.82_41,
											    -1.23_10, -0.80_98, -0.33_20, -0.41_01, -0.79_85, -0.49_96, -0.82_13, -0.91_28,
											    -1.04_20, -1.12_86, -1.04_40, -0.79_99, -0.84_05, -1.22_75, -1.54_43, -1.46_25,
											]    )
											# fmt: on

											__lowerCAmelCase            = self._load_datasamples(1    )
											__lowerCAmelCase            = self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()    )
											__lowerCAmelCase            = feature_extractor(_A	,    return_tensors="pt"    ).input_features
											self.assertEquals(input_features.shape	,    (1, 5_8_4, 2_4)    )
											self.assertTrue(np.allclose(input_features[0, 0, :3_0]	,    _A	,    atol=1E-4    )    )




