
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__    ):
				lowerCamelCase_						=			model.config

				lowerCamelCase_						=			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	, )
				lowerCamelCase_						=			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__    ):
				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 			lowerCamelCase_ (				lowerCamelCase__	, lowerCamelCase__    ):
				for key in orig_state_dict.copy().keys():
								lowerCamelCase_						=			orig_state_dict.pop(lowerCamelCase__    )

								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 			lowerCamelCase_ (				lowerCamelCase__	, lowerCamelCase__=None	, lowerCamelCase__=False    ):
				# load original model
				lowerCamelCase_						=			DonutModel.from_pretrained(lowerCamelCase__    ).eval()

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

				lowerCamelCase_						=			original_model.state_dict()
				lowerCamelCase_						=			convert_state_dict(lowerCamelCase__	, lowerCamelCase__    )
				model.load_state_dict(lowerCamelCase__    )

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

				lowerCamelCase_						=			XLMRobertaTokenizerFast.from_pretrained(lowerCamelCase__	, from_slow=lowerCamelCase__    )
				lowerCamelCase_						=			DonutImageProcessor(
				    do_align_long_axis=original_model.config.align_long_axis	, size=original_model.config.input_size[::-1]    )
				lowerCamelCase_						=			DonutProcessor(lowerCamelCase__	, lowerCamelCase__    )
				lowerCamelCase_						=			processor(lowerCamelCase__	, 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}"	, lowerCamelCase__    )
				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(lowerCamelCase__	, add_special_tokens=lowerCamelCase__	, return_tensors="pt"    )[
				    "input_ids"
				]

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

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

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









import unittest
from pathlib import Path
from tempfile import NamedTemporaryFile, TemporaryDirectory

from transformers import BertConfig, BertTokenizerFast, FeatureExtractionPipeline
from transformers.convert_graph_to_onnx import (
    convert,
    ensure_valid_input,
    generate_identified_filename,
    infer_shapes,
    quantize,
)
from transformers.testing_utils import require_tf, require_tokenizers, require_torch, slow



class 					__lowerCAmelCase       :






			def 			lowerCamelCase  (self					,							__magic_name__					,							__magic_name__					,							__magic_name__    )     ->   Dict:


							'''simple docstring'''



							return None




class 					__lowerCAmelCase       :






			def 			lowerCamelCase  (self					,							__magic_name__					,							__magic_name__					,							__magic_name__					,							__magic_name__    )     ->   List[str]:


							'''simple docstring'''



							return None




class 					__lowerCAmelCase       (			unittest.TestCase     ):
			lowerCamelCase_     :     Dict       			=      [
			    # (model_name, model_kwargs)
			    ('''bert-base-cased''', {}),
			    ('''gpt2''', {'''use_cache''': False}),  # We don't support exporting GPT2 past keys anymore
			]






			@require_tf
			@slow
			def 			lowerCamelCase  (self    )     ->   Union[str, Any]:


							'''simple docstring'''



							for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
											self._test_export(__magic_name__					,							'''tf'''					,							12					,							**__magic_name__    )






			@require_torch
			@slow
			def 			lowerCamelCase  (self    )     ->   int:


							'''simple docstring'''



							for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
											self._test_export(__magic_name__					,							'''pt'''					,							12					,							**__magic_name__    )






			@require_torch
			@slow
			def 			lowerCamelCase  (self    )     ->   int:


							'''simple docstring'''



							from transformers import BertModel

							snake_case_						:    str				=  ['''[UNK]''', '''[SEP]''', '''[CLS]''', '''[PAD]''', '''[MASK]''', '''some''', '''other''', '''words''']
							with NamedTemporaryFile(mode='''w+t'''    ) as vocab_file:
											vocab_file.write('''\n'''.join(__magic_name__    )    )
											vocab_file.flush()
											snake_case_						:    Optional[Any]				=  BertTokenizerFast(vocab_file.name    )

							with TemporaryDirectory() as bert_save_dir:
											snake_case_						:    str				=  BertModel(BertConfig(vocab_size=len(__magic_name__    )    )    )
											model.save_pretrained(__magic_name__    )
											self._test_export(__magic_name__					,							'''pt'''					,							12					,							__magic_name__    )






			@require_tf
			@slow
			def 			lowerCamelCase  (self    )     ->   Tuple:


							'''simple docstring'''



							for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
											snake_case_						:    Tuple				=  self._test_export(__magic_name__					,							'''tf'''					,							12					,							**__magic_name__    )
											snake_case_						:    List[str]				=  quantize(Path(__magic_name__    )    )

											# Ensure the actual quantized model is not bigger than the original one
											if quantized_path.stat().st_size >= Path(__magic_name__    ).stat().st_size:
															self.fail('''Quantized model is bigger than initial ONNX model'''    )






			@require_torch
			@slow
			def 			lowerCamelCase  (self    )     ->   Any:


							'''simple docstring'''



							for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
											snake_case_						:    Any				=  self._test_export(__magic_name__					,							'''pt'''					,							12					,							**__magic_name__    )
											snake_case_						:    Any				=  quantize(__magic_name__    )

											# Ensure the actual quantized model is not bigger than the original one
											if quantized_path.stat().st_size >= Path(__magic_name__    ).stat().st_size:
															self.fail('''Quantized model is bigger than initial ONNX model'''    )






			def 			lowerCamelCase  (self					,							__magic_name__					,							__magic_name__					,							__magic_name__					,							__magic_name__=None					,							**__magic_name__    )     ->   Tuple:


							'''simple docstring'''



							try:
											# Compute path
											with TemporaryDirectory() as tempdir:
															snake_case_						:    List[str]				=  Path(__magic_name__    ).joinpath('''model.onnx'''    )

											# Remove folder if exists
											if path.parent.exists():
															path.parent.rmdir()

											# Export
											convert(__magic_name__					,							__magic_name__					,							__magic_name__					,							__magic_name__					,							__magic_name__					,							**__magic_name__    )

											return path
							except Exception as e:
											self.fail(__magic_name__    )






			@require_torch
			@require_tokenizers
			@slow
			def 			lowerCamelCase  (self    )     ->   Optional[Any]:


							'''simple docstring'''



							from transformers import BertModel

							snake_case_						:    Optional[Any]				=  BertModel(BertConfig.from_pretrained('''lysandre/tiny-bert-random'''    )    )
							snake_case_						:    int				=  BertTokenizerFast.from_pretrained('''lysandre/tiny-bert-random'''    )
							self._test_infer_dynamic_axis(__magic_name__					,							__magic_name__					,							'''pt'''    )






			@require_tf
			@require_tokenizers
			@slow
			def 			lowerCamelCase  (self    )     ->   List[str]:


							'''simple docstring'''



							from transformers import TFBertModel

							snake_case_						:    Any				=  TFBertModel(BertConfig.from_pretrained('''lysandre/tiny-bert-random'''    )    )
							snake_case_						:    str				=  BertTokenizerFast.from_pretrained('''lysandre/tiny-bert-random'''    )
							self._test_infer_dynamic_axis(__magic_name__					,							__magic_name__					,							'''tf'''    )






			def 			lowerCamelCase  (self					,							__magic_name__					,							__magic_name__					,							__magic_name__    )     ->   int:


							'''simple docstring'''



							snake_case_						:    Tuple				=  FeatureExtractionPipeline(__magic_name__					,							__magic_name__    )

							snake_case_						:    Optional[int]				=  ['''input_ids''', '''token_type_ids''', '''attention_mask''', '''output_0''', '''output_1''']
							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_						:    Optional[int]				=  infer_shapes(__magic_name__					,							__magic_name__    )

							# Assert all variables are present
							self.assertEqual(len(__magic_name__    )					,							len(__magic_name__    )    )
							self.assertTrue(all(var_name in shapes for var_name in variable_names    )    )
							self.assertSequenceEqual(variable_names[:3]					,							__magic_name__    )
							self.assertSequenceEqual(variable_names[3:]					,							__magic_name__    )

							# Assert inputs are {0: batch, 1: sequence}
							for var_name in ["input_ids", "token_type_ids", "attention_mask"]:
											self.assertDictEqual(shapes[var_name]					,							{0: '''batch''', 1: '''sequence'''}    )

							# Assert outputs are {0: batch, 1: sequence} and {0: batch}
							self.assertDictEqual(shapes['''output_0''']					,							{0: '''batch''', 1: '''sequence'''}    )
							self.assertDictEqual(shapes['''output_1''']					,							{0: '''batch'''}    )






			def 			lowerCamelCase  (self    )     ->   Optional[int]:


							'''simple docstring'''



							snake_case_						:    Tuple				=  ['''input_ids''', '''attention_mask''', '''token_type_ids''']
							snake_case_						:    List[str]				=  {'''input_ids''': [1, 2, 3, 4], '''attention_mask''': [0, 0, 0, 0], '''token_type_ids''': [1, 1, 1, 1]}
							snake_case_			,							snake_case_						:    Tuple				=  ensure_valid_input(FuncContiguousArgs()					,							__magic_name__					,							__magic_name__    )

							# Should have exactly the same number of args (all are valid)
							self.assertEqual(len(__magic_name__    )					,							3    )

							# Should have exactly the same input names
							self.assertEqual(set(__magic_name__    )					,							set(__magic_name__    )    )

							# Parameter should be reordered according to their respective place in the function:
							# (input_ids, token_type_ids, attention_mask)
							self.assertEqual(__magic_name__					,							(tokens['''input_ids'''], tokens['''token_type_ids'''], tokens['''attention_mask'''])    )

							# Generated args are interleaved with another args (for instance parameter "past" in GPT2)
							snake_case_			,							snake_case_						:    Dict				=  ensure_valid_input(FuncNonContiguousArgs()					,							__magic_name__					,							__magic_name__    )

							# Should have exactly the one arg (all before the one not provided "some_other_args")
							self.assertEqual(len(__magic_name__    )					,							1    )
							self.assertEqual(len(__magic_name__    )					,							1    )

							# Should have only "input_ids"
							self.assertEqual(inputs_args[0]					,							tokens['''input_ids''']    )
							self.assertEqual(ordered_input_names[0]					,							'''input_ids'''    )






			def 			lowerCamelCase  (self    )     ->   Any:


							'''simple docstring'''



							snake_case_						:    Optional[int]				=  generate_identified_filename(Path('''/home/something/my_fake_model.onnx'''    )					,							'''-test'''    )
							self.assertEqual('''/home/something/my_fake_model-test.onnx'''					,							generated.as_posix()    )
