


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  _snake_case       (					A	)						->   List[str]:
							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=128			,   )
							lowerCAmelCase__	       =     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  _snake_case       (					A	)						->   List[Any]:
							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  _snake_case       (					A			,   A	)						->   List[str]:
							for key in orig_state_dict.copy().keys():
														lowerCAmelCase__	       =     orig_state_dict.pop(A	)

														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  _snake_case       (					A			,   A=None			,   A=False	)						->   Union[str, Any]:
							# load original model
							lowerCAmelCase__	       =     DonutModel.from_pretrained(A	).eval()

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

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

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

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

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

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

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



'''simple docstring'''






import unittest

from transformers import (
    MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
    TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
    TextClassificationPipeline,
    pipeline,
)
from transformers.testing_utils import is_pipeline_test, nested_simplify, require_tf, require_torch, slow

from .test_pipelines_common import ANY


# These 2 model types require different inputs than those of the usual text models.
__UpperCAmelCase       =							{'''LayoutLMv2Config''', '''LayoutLMv3Config'''}




@is_pipeline_test
class 		a__    (							unittest.TestCase  ):





							'''simple docstring'''






							lowercase__      :       Optional[Any]     						=				MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING
							lowercase__      :       List[Any]     						=				TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING

							if model_mapping is not None:
														lowercase__      :       Optional[Any]     						=				{config: model for config, model in model_mapping.items() if config.__name__ not in _TO_SKIP}
							if tf_model_mapping is not None:
														lowercase__      :       Tuple     						=				{
														    config: model for config, model in tf_model_mapping.items() if config.__name__ not in _TO_SKIP
														}







							@require_torch
							def   __SCREAMING_SNAKE_CASE							(		self       )	->							Optional[int]:
														lowerCAmelCase__	       =     pipeline(
														    task='''text-classification''' ,   model='''hf-internal-testing/tiny-random-distilbert''' ,   framework='''pt'''       )

														lowerCAmelCase__	       =     text_classifier('''This is great !'''       )
														self.assertEqual(nested_simplify(lowerCamelCase_       ) ,   [{'''label''': '''LABEL_0''', '''score''': 0.504}]       )

														lowerCAmelCase__	       =     text_classifier('''This is great !''' ,   top_k=2       )
														self.assertEqual(
														    nested_simplify(lowerCamelCase_       ) ,   [{'''label''': '''LABEL_0''', '''score''': 0.504}, {'''label''': '''LABEL_1''', '''score''': 0.496}]       )

														lowerCAmelCase__	       =     text_classifier(['''This is great !''', '''This is bad'''] ,   top_k=2       )
														self.assertEqual(
														    nested_simplify(lowerCamelCase_       ) ,   [
														        [{'''label''': '''LABEL_0''', '''score''': 0.504}, {'''label''': '''LABEL_1''', '''score''': 0.496}],
														        [{'''label''': '''LABEL_0''', '''score''': 0.504}, {'''label''': '''LABEL_1''', '''score''': 0.496}],
														    ] ,   )

														lowerCAmelCase__	       =     text_classifier('''This is great !''' ,   top_k=1       )
														self.assertEqual(nested_simplify(lowerCamelCase_       ) ,   [{'''label''': '''LABEL_0''', '''score''': 0.504}]       )

														# Legacy behavior
														lowerCAmelCase__	       =     text_classifier('''This is great !''' ,   return_all_scores=lowerCamelCase_       )
														self.assertEqual(nested_simplify(lowerCamelCase_       ) ,   [{'''label''': '''LABEL_0''', '''score''': 0.504}]       )

														lowerCAmelCase__	       =     text_classifier('''This is great !''' ,   return_all_scores=lowerCamelCase_       )
														self.assertEqual(
														    nested_simplify(lowerCamelCase_       ) ,   [[{'''label''': '''LABEL_0''', '''score''': 0.504}, {'''label''': '''LABEL_1''', '''score''': 0.496}]]       )

														lowerCAmelCase__	       =     text_classifier(['''This is great !''', '''Something else'''] ,   return_all_scores=lowerCamelCase_       )
														self.assertEqual(
														    nested_simplify(lowerCamelCase_       ) ,   [
														        [{'''label''': '''LABEL_0''', '''score''': 0.504}, {'''label''': '''LABEL_1''', '''score''': 0.496}],
														        [{'''label''': '''LABEL_0''', '''score''': 0.504}, {'''label''': '''LABEL_1''', '''score''': 0.496}],
														    ] ,   )

														lowerCAmelCase__	       =     text_classifier(['''This is great !''', '''Something else'''] ,   return_all_scores=lowerCamelCase_       )
														self.assertEqual(
														    nested_simplify(lowerCamelCase_       ) ,   [
														        {'''label''': '''LABEL_0''', '''score''': 0.504},
														        {'''label''': '''LABEL_0''', '''score''': 0.504},
														    ] ,   )







							@require_torch
							def   __SCREAMING_SNAKE_CASE							(		self       )	->							int:
														import torch

														lowerCAmelCase__	       =     pipeline(
														    task='''text-classification''' ,   model='''hf-internal-testing/tiny-random-distilbert''' ,   framework='''pt''' ,   device=torch.device('''cpu'''       ) ,   )

														lowerCAmelCase__	       =     text_classifier('''This is great !'''       )
														self.assertEqual(nested_simplify(lowerCamelCase_       ) ,   [{'''label''': '''LABEL_0''', '''score''': 0.504}]       )







							@require_tf
							def   __SCREAMING_SNAKE_CASE							(		self       )	->							str:
														lowerCAmelCase__	       =     pipeline(
														    task='''text-classification''' ,   model='''hf-internal-testing/tiny-random-distilbert''' ,   framework='''tf'''       )

														lowerCAmelCase__	       =     text_classifier('''This is great !'''       )
														self.assertEqual(nested_simplify(lowerCamelCase_       ) ,   [{'''label''': '''LABEL_0''', '''score''': 0.504}]       )







							@slow
							@require_torch
							def   __SCREAMING_SNAKE_CASE							(		self       )	->							Dict:
														lowerCAmelCase__	       =     pipeline('''text-classification'''       )

														lowerCAmelCase__	       =     text_classifier('''This is great !'''       )
														self.assertEqual(nested_simplify(lowerCamelCase_       ) ,   [{'''label''': '''POSITIVE''', '''score''': 1.0}]       )
														lowerCAmelCase__	       =     text_classifier('''This is bad !'''       )
														self.assertEqual(nested_simplify(lowerCamelCase_       ) ,   [{'''label''': '''NEGATIVE''', '''score''': 1.0}]       )
														lowerCAmelCase__	       =     text_classifier('''Birds are a type of animal'''       )
														self.assertEqual(nested_simplify(lowerCamelCase_       ) ,   [{'''label''': '''POSITIVE''', '''score''': 0.988}]       )







							@slow
							@require_tf
							def   __SCREAMING_SNAKE_CASE							(		self       )	->							Optional[Any]:
														lowerCAmelCase__	       =     pipeline('''text-classification''' ,   framework='''tf'''       )

														lowerCAmelCase__	       =     text_classifier('''This is great !'''       )
														self.assertEqual(nested_simplify(lowerCamelCase_       ) ,   [{'''label''': '''POSITIVE''', '''score''': 1.0}]       )
														lowerCAmelCase__	       =     text_classifier('''This is bad !'''       )
														self.assertEqual(nested_simplify(lowerCamelCase_       ) ,   [{'''label''': '''NEGATIVE''', '''score''': 1.0}]       )
														lowerCAmelCase__	       =     text_classifier('''Birds are a type of animal'''       )
														self.assertEqual(nested_simplify(lowerCamelCase_       ) ,   [{'''label''': '''POSITIVE''', '''score''': 0.988}]       )







							def   __SCREAMING_SNAKE_CASE							(		self ,   lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_       )	->							Tuple:
														lowerCAmelCase__	       =     TextClassificationPipeline(model=lowerCamelCase_ ,   tokenizer=lowerCamelCase_       )
														return text_classifier, ["HuggingFace is in", "This is another test"]







							def   __SCREAMING_SNAKE_CASE							(		self ,   lowerCamelCase_ ,   lowerCamelCase_       )	->							Union[str, Any]:
														lowerCAmelCase__	       =     text_classifier.model
														# Small inputs because BartTokenizer tiny has maximum position embeddings = 22
														lowerCAmelCase__	       =     '''HuggingFace is in'''
														lowerCAmelCase__	       =     text_classifier(lowerCamelCase_       )

														self.assertEqual(nested_simplify(lowerCamelCase_       ) ,   [{'''label''': ANY(lowerCamelCase_       ), '''score''': ANY(lowerCamelCase_       )}]       )
														self.assertTrue(outputs[0]['''label'''] in model.config.idalabel.values()       )

														lowerCAmelCase__	       =     ['''HuggingFace is in ''', '''Paris is in France''']
														lowerCAmelCase__	       =     text_classifier(lowerCamelCase_       )
														self.assertEqual(
														    nested_simplify(lowerCamelCase_       ) ,   [{'''label''': ANY(lowerCamelCase_       ), '''score''': ANY(lowerCamelCase_       )}, {'''label''': ANY(lowerCamelCase_       ), '''score''': ANY(lowerCamelCase_       )}] ,   )
														self.assertTrue(outputs[0]['''label'''] in model.config.idalabel.values()       )
														self.assertTrue(outputs[1]['''label'''] in model.config.idalabel.values()       )

														# Forcing to get all results with `top_k=None`
														# This is NOT the legacy format
														lowerCAmelCase__	       =     text_classifier(lowerCamelCase_ ,   top_k=lowerCamelCase_       )
														lowerCAmelCase__	       =     len(model.config.idalabel.values()       )
														self.assertEqual(
														    nested_simplify(lowerCamelCase_       ) ,   [[{'''label''': ANY(lowerCamelCase_       ), '''score''': ANY(lowerCamelCase_       )}] * N, [{'''label''': ANY(lowerCamelCase_       ), '''score''': ANY(lowerCamelCase_       )}] * N] ,   )

														lowerCAmelCase__	       =     {'''text''': '''HuggingFace is in ''', '''text_pair''': '''Paris is in France'''}
														lowerCAmelCase__	       =     text_classifier(lowerCamelCase_       )
														self.assertEqual(
														    nested_simplify(lowerCamelCase_       ) ,   {'''label''': ANY(lowerCamelCase_       ), '''score''': ANY(lowerCamelCase_       )} ,   )
														self.assertTrue(outputs['''label'''] in model.config.idalabel.values()       )

														# This might be used a text pair, but tokenizer + pipe interaction
														# makes it hard to understand that it's not using the pair properly
														# https://github.com/huggingface/transformers/issues/17305
														# We disabled this usage instead as it was outputting wrong outputs.
														lowerCAmelCase__	       =     [['''HuggingFace is in ''', '''Paris is in France''']]
														with self.assertRaises(lowerCamelCase_       ):
																					text_classifier(lowerCamelCase_       )

														# This used to be valid for doing text pairs
														# We're keeping it working because of backward compatibility
														lowerCAmelCase__	       =     text_classifier([[['''HuggingFace is in ''', '''Paris is in France''']]]       )
														self.assertEqual(
														    nested_simplify(lowerCamelCase_       ) ,   [{'''label''': ANY(lowerCamelCase_       ), '''score''': ANY(lowerCamelCase_       )}] ,   )
														self.assertTrue(outputs[0]['''label'''] in model.config.idalabel.values()       )