




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)










from __future__ import annotations

import unittest

from transformers import DebertaVaConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
					import tensorflow as tf

					from transformers import (
					    TFDebertaVaForMaskedLM,
					    TFDebertaVaForQuestionAnswering,
					    TFDebertaVaForSequenceClassification,
					    TFDebertaVaForTokenClassification,
					    TFDebertaVaModel,
					)
class     a__					:






						def __init__(     self	,    _A	,    _A=1_3	,    _A=7	,    _A=True	,    _A=True	,    _A=True	,    _A=True	,    _A=9_9	,    _A=3_2	,    _A=2	,    _A=4	,    _A=3_7	,    _A="gelu"	,    _A=0.1	,    _A=0.1	,    _A=5_1_2	,    _A=1_6	,    _A=2	,    _A=0.02	,    _A=False	,    _A=True	,    _A="None"	,    _A=3	,    _A=4	,    _A=None	,    ):




											"""simple docstring"""



											__lowerCAmelCase            = parent
											__lowerCAmelCase            = batch_size
											__lowerCAmelCase            = seq_length
											__lowerCAmelCase            = is_training
											__lowerCAmelCase            = use_input_mask
											__lowerCAmelCase            = use_token_type_ids
											__lowerCAmelCase            = use_labels
											__lowerCAmelCase            = vocab_size
											__lowerCAmelCase            = hidden_size
											__lowerCAmelCase            = num_hidden_layers
											__lowerCAmelCase            = num_attention_heads
											__lowerCAmelCase            = intermediate_size
											__lowerCAmelCase            = hidden_act
											__lowerCAmelCase            = hidden_dropout_prob
											__lowerCAmelCase            = attention_probs_dropout_prob
											__lowerCAmelCase            = max_position_embeddings
											__lowerCAmelCase            = type_vocab_size
											__lowerCAmelCase            = type_sequence_label_size
											__lowerCAmelCase            = initializer_range
											__lowerCAmelCase            = num_labels
											__lowerCAmelCase            = num_choices
											__lowerCAmelCase            = relative_attention
											__lowerCAmelCase            = position_biased_input
											__lowerCAmelCase            = pos_att_type
											__lowerCAmelCase            = scope






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = ids_tensor([self.batch_size, self.seq_length]	,    self.vocab_size    )

											__lowerCAmelCase            = None
											if self.use_input_mask:
																__lowerCAmelCase            = random_attention_mask([self.batch_size, self.seq_length]    )

											__lowerCAmelCase            = None
											if self.use_token_type_ids:
																__lowerCAmelCase            = ids_tensor([self.batch_size, self.seq_length]	,    self.type_vocab_size    )

											__lowerCAmelCase            = None
											__lowerCAmelCase            = None
											__lowerCAmelCase            = None
											if self.use_labels:
																__lowerCAmelCase            = ids_tensor([self.batch_size]	,    self.type_sequence_label_size    )
																__lowerCAmelCase            = ids_tensor([self.batch_size, self.seq_length]	,    self.num_labels    )

											__lowerCAmelCase            = DebertaVaConfig(
											    vocab_size=self.vocab_size	,    hidden_size=self.hidden_size	,    num_hidden_layers=self.num_hidden_layers	,    num_attention_heads=self.num_attention_heads	,    intermediate_size=self.intermediate_size	,    hidden_act=self.hidden_act	,    hidden_dropout_prob=self.hidden_dropout_prob	,    attention_probs_dropout_prob=self.attention_probs_dropout_prob	,    max_position_embeddings=self.max_position_embeddings	,    type_vocab_size=self.type_vocab_size	,    relative_attention=self.relative_attention	,    position_biased_input=self.position_biased_input	,    initializer_range=self.initializer_range	,    return_dict=_A	,    )

											return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels






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




											"""simple docstring"""



											__lowerCAmelCase            = TFDebertaVaModel(config=_A    )
											__lowerCAmelCase            = {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}

											__lowerCAmelCase            = [input_ids, input_mask]
											__lowerCAmelCase            = model(_A    )

											__lowerCAmelCase            = model(_A    )

											self.parent.assertEqual(result.last_hidden_state.shape	,    (self.batch_size, self.seq_length, self.hidden_size)    )






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




											"""simple docstring"""



											__lowerCAmelCase            = TFDebertaVaForMaskedLM(config=_A    )
											__lowerCAmelCase            = {
											    "input_ids": input_ids,
											    "attention_mask": input_mask,
											    "token_type_ids": token_type_ids,
											}
											__lowerCAmelCase            = model(_A    )
											self.parent.assertEqual(result.logits.shape	,    (self.batch_size, self.seq_length, self.vocab_size)    )






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




											"""simple docstring"""



											__lowerCAmelCase            = self.num_labels
											__lowerCAmelCase            = TFDebertaVaForSequenceClassification(config=_A    )
											__lowerCAmelCase            = {
											    "input_ids": input_ids,
											    "attention_mask": input_mask,
											    "token_type_ids": token_type_ids,
											}

											__lowerCAmelCase            = model(_A    )
											self.parent.assertEqual(result.logits.shape	,    (self.batch_size, self.num_labels)    )






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




											"""simple docstring"""



											__lowerCAmelCase            = self.num_labels
											__lowerCAmelCase            = TFDebertaVaForTokenClassification(config=_A    )
											__lowerCAmelCase            = {
											    "input_ids": input_ids,
											    "attention_mask": input_mask,
											    "token_type_ids": token_type_ids,
											}
											__lowerCAmelCase            = model(_A    )
											self.parent.assertEqual(result.logits.shape	,    (self.batch_size, self.seq_length, self.num_labels)    )






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




											"""simple docstring"""



											__lowerCAmelCase            = TFDebertaVaForQuestionAnswering(config=_A    )
											__lowerCAmelCase            = {
											    "input_ids": input_ids,
											    "attention_mask": input_mask,
											    "token_type_ids": token_type_ids,
											}

											__lowerCAmelCase            = model(_A    )
											self.parent.assertEqual(result.start_logits.shape	,    (self.batch_size, self.seq_length)    )
											self.parent.assertEqual(result.end_logits.shape	,    (self.batch_size, self.seq_length)    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.prepare_config_and_inputs()
											(
											    (
											    __lowerCAmelCase
											)     ,							(
											    __lowerCAmelCase
											)     ,							(
											    __lowerCAmelCase
											)     ,							(
											    __lowerCAmelCase
											)     ,							(
											    __lowerCAmelCase
											)     ,							(
											    __lowerCAmelCase
											)     ,							(
											    __lowerCAmelCase
											)     ,							
											)            = config_and_inputs
											__lowerCAmelCase            = {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
											return config, inputs_dict






@require_tf
class     a__					(      snake_case__				,       snake_case__				,       unittest.TestCase   ):
						_a							:					Optional[Any]      =					(
						    (
						        TFDebertaVaModel,
						        TFDebertaVaForMaskedLM,
						        TFDebertaVaForQuestionAnswering,
						        TFDebertaVaForSequenceClassification,
						        TFDebertaVaForTokenClassification,
						    )
						    if is_tf_available()
						    else ()
						)
						_a							:					Union[str, Any]      =					(
						    {
						        """feature-extraction""": TFDebertaVaModel,
						        """fill-mask""": TFDebertaVaForMaskedLM,
						        """question-answering""": TFDebertaVaForQuestionAnswering,
						        """text-classification""": TFDebertaVaForSequenceClassification,
						        """token-classification""": TFDebertaVaForTokenClassification,
						        """zero-shot""": TFDebertaVaForSequenceClassification,
						    }
						    if is_tf_available()
						    else {}
						)

						_a							:					str      =					False
						_a							:					List[str]      =					False






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = TFDebertaVaModelTester(self    )
											__lowerCAmelCase            = ConfigTester(self	,    config_class=_A	,    hidden_size=3_7    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											self.config_tester.run_common_tests()






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*_A    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_lm(*_A    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_question_answering(*_A    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_sequence_classification(*_A    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_token_classification(*_A    )






						@slow
						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = TFDebertaVaModel.from_pretrained("kamalkraj/deberta-v2-xlarge"    )
											self.assertIsNotNone(_A    )






@require_tf
class     a__					(      unittest.TestCase   ):






						@unittest.skip(reason="Model not available yet"    )
						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											pass






						@slow
						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = TFDebertaVaModel.from_pretrained("kamalkraj/deberta-v2-xlarge"    )
											__lowerCAmelCase            = tf.constant([[0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9, 4_6_0_7_8, 1_5_8_8, 2]]    )
											__lowerCAmelCase            = tf.constant([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]    )
											__lowerCAmelCase            = model(_A	,    attention_mask=_A    )[0]

											__lowerCAmelCase            = tf.constant(
											    [[[0.23_56, 0.19_48, 0.03_69], [-0.10_63, 0.35_86, -0.51_52], [-0.63_99, -0.02_59, -0.25_25]]]    )
											tf.debugging.assert_near(output[:, 1:4, 1:4]	,    _A	,    atol=1E-4    )




