



'''simple docstring'''

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_     (    __UpperCAmelCase					)				->					List[Any]:





			"""simple docstring"""







			snake_case:    List[str]      							=model.config

			snake_case:    Optional[Any]      							=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				,							)
			snake_case:    Union[str, Any]      							=MBartConfig(
			    is_decoder=_UpperCAmelCase				,							is_encoder_decoder=_UpperCAmelCase				,							add_cross_attention=_UpperCAmelCase				,							decoder_layers=original_config.decoder_layer				,							max_position_embeddings=original_config.max_position_embeddings				,							vocab_size=len(
			        model.decoder.tokenizer					)				,							scale_embedding=_UpperCAmelCase				,							add_final_layer_norm=_UpperCAmelCase				,							)

			return encoder_config, decoder_config







def a_     (    __UpperCAmelCase					)				->					List[str]:





			"""simple docstring"""







			if "encoder.model" in name:
						snake_case:    int      							=name.replace('encoder.model'				,							'encoder'					)
			if "decoder.model" in name:
						snake_case:    Optional[Any]      							=name.replace('decoder.model'				,							'decoder'					)
			if "patch_embed.proj" in name:
						snake_case:    int      							=name.replace('patch_embed.proj'				,							'embeddings.patch_embeddings.projection'					)
			if "patch_embed.norm" in name:
						snake_case:    Tuple      							=name.replace('patch_embed.norm'				,							'embeddings.norm'					)
			if name.startswith('encoder'					):
						if "layers" in name:
									snake_case:    Tuple      							='encoder.' + name
						if "attn.proj" in name:
									snake_case:    List[str]      							=name.replace('attn.proj'				,							'attention.output.dense'					)
						if "attn" in name and "mask" not in name:
									snake_case:    Any      							=name.replace('attn'				,							'attention.self'					)
						if "norm1" in name:
									snake_case:    Optional[int]      							=name.replace('norm1'				,							'layernorm_before'					)
						if "norm2" in name:
									snake_case:    List[Any]      							=name.replace('norm2'				,							'layernorm_after'					)
						if "mlp.fc1" in name:
									snake_case:    Dict      							=name.replace('mlp.fc1'				,							'intermediate.dense'					)
						if "mlp.fc2" in name:
									snake_case:    List[Any]      							=name.replace('mlp.fc2'				,							'output.dense'					)

						if name == "encoder.norm.weight":
									snake_case:    Union[str, Any]      							='encoder.layernorm.weight'
						if name == "encoder.norm.bias":
									snake_case:    Union[str, Any]      							='encoder.layernorm.bias'

			return name







def a_     (    __UpperCAmelCase				,							__UpperCAmelCase					)				->					Union[str, Any]:





			"""simple docstring"""







			for key in orig_state_dict.copy().keys():
						snake_case:    Any      							=orig_state_dict.pop(_UpperCAmelCase					)

						if "qkv" in key:
									snake_case:    str      							=key.split('.'					)
									snake_case:    Optional[Any]      							=int(key_split[3]					)
									snake_case:    List[str]      							=int(key_split[5]					)
									snake_case:    List[str]      							=model.encoder.encoder.layers[layer_num].blocks[block_num].attention.self.all_head_size

									if "weight" in key:
												snake_case:    int      							=val[:dim, :]
												snake_case:    str      							=val[dim : dim * 2, :]
												snake_case:    Optional[Any]      							=val[-dim:, :]
									else:
												snake_case:    Union[str, Any]      							=val[:dim]
												snake_case:    List[str]      							=val[dim : dim * 2]
												snake_case:    Optional[Any]      							=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:
									snake_case:    Tuple      							=val

			return orig_state_dict







def a_     (    __UpperCAmelCase				,							__UpperCAmelCase=None				,							__UpperCAmelCase=False					)				->					int:





			"""simple docstring"""







			snake_case:    str      							=DonutModel.from_pretrained(_UpperCAmelCase					).eval()

			# load HuggingFace model
			snake_case:    Dict      							=get_configs(_UpperCAmelCase					)
			snake_case:    Optional[int]      							=DonutSwinModel(_UpperCAmelCase					)
			snake_case:    Dict      							=MBartForCausalLM(_UpperCAmelCase					)
			snake_case:    List[Any]      							=VisionEncoderDecoderModel(encoder=_UpperCAmelCase				,							decoder=_UpperCAmelCase					)
			model.eval()

			snake_case:    str      							=original_model.state_dict()
			snake_case:    List[Any]      							=convert_state_dict(_UpperCAmelCase				,							_UpperCAmelCase					)
			model.load_state_dict(_UpperCAmelCase					)

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

			snake_case:    Optional[Any]      							=XLMRobertaTokenizerFast.from_pretrained(_UpperCAmelCase				,							from_slow=_UpperCAmelCase					)
			snake_case:    Any      							=DonutImageProcessor(
			    do_align_long_axis=original_model.config.align_long_axis				,							size=original_model.config.input_size[::-1]					)
			snake_case:    List[Any]      							=DonutProcessor(_UpperCAmelCase				,							_UpperCAmelCase					)
			snake_case:    Dict      							=processor(_UpperCAmelCase				,							return_tensors='pt'					).pixel_values

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

			snake_case:    Union[str, Any]      							=original_model.encoder.model.patch_embed(_UpperCAmelCase					)
			snake_case:    Union[str, Any]      							=model.encoder.embeddings(_UpperCAmelCase					)
			assert torch.allclose(_UpperCAmelCase				,							_UpperCAmelCase				,							atol=1e-3					)

			# verify encoder hidden states
			snake_case:    Tuple      							=original_model.encoder(_UpperCAmelCase					)
			snake_case:    str      							=model.encoder(_UpperCAmelCase					).last_hidden_state
			assert torch.allclose(_UpperCAmelCase				,							_UpperCAmelCase				,							atol=1e-2					)

			# verify decoder hidden states
			snake_case:    Any      							=original_model(_UpperCAmelCase				,							_UpperCAmelCase				,							_UpperCAmelCase					).logits
			snake_case:    List[str]      							=model(_UpperCAmelCase				,							decoder_input_ids=_UpperCAmelCase					).logits
			assert torch.allclose(_UpperCAmelCase				,							_UpperCAmelCase				,							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(_UpperCAmelCase					)
						processor.save_pretrained(_UpperCAmelCase					)

			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)









'''simple docstring'''





import re
import warnings
from contextlib import contextmanager

from ...processing_utils import ProcessorMixin



class  snake_case__						(    SCREAMING_SNAKE_CASE_				):
 A__				= ['''image_processor''', '''tokenizer''']
 A__				= '''AutoImageProcessor'''
 A__				= '''AutoTokenizer'''
 def __init__(   self						:							List[str]      ,     __a						:							Optional[int]=None      ,     __a						:							Any=None      ,     **__a						:							Optional[int]  )      ->   Union[str, Any]:







       '''simple docstring'''
       __snake_case :   str													=    None
       if "feature_extractor" in kwargs:
             warnings.warn(
                 'The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`'
                 ' instead.'      ,     __a      ,     )
             __snake_case :   Tuple													=    kwargs.pop('feature_extractor'  )

       __snake_case :   Union[str, Any]													=    image_processor if image_processor is not None else feature_extractor
       if image_processor is None:
             raise ValueError('You need to specify an `image_processor`.'  )
       if tokenizer is None:
             raise ValueError('You need to specify a `tokenizer`.'  )

       super().__init__(__a      ,     __a  )
       __snake_case :   Any													=    self.image_processor
       __snake_case :   int													=    False
 def __call__(   self						:							str      ,     *__a						:							List[str]      ,     **__a						:							List[str]  )      ->   Dict:







       '''simple docstring'''
       # For backward compatibility
       if self._in_target_context_manager:
             return self.current_processor(*__a      ,     **__a  )

       __snake_case :   List[Any]													=    kwargs.pop('images'      ,     __a  )
       __snake_case :   Any													=    kwargs.pop('text'      ,     __a  )
       if len(__a  ) > 0:
             __snake_case :   List[Any]													=    args[0]
             __snake_case :   Dict													=    args[1:]

       if images is None and text is None:
             raise ValueError('You need to specify either an `images` or `text` input to process.'  )

       if images is not None:
             __snake_case :   int													=    self.image_processor(__a      ,     *__a      ,     **__a  )
       if text is not None:
             __snake_case :   Tuple													=    self.tokenizer(__a      ,     **__a  )

       if text is None:
             return inputs
       elif images is None:
             return encodings
       else:
             __snake_case :   List[str]													=    encodings['input_ids']
             return inputs
 def A_   (   self						:							str      ,     *__a						:							Union[str, Any]      ,     **__a						:							Optional[Any]  )      ->   List[str]:







       '''simple docstring'''
       return self.tokenizer.batch_decode(*__a      ,     **__a  )
 def A_   (   self						:							Any      ,     *__a						:							List[Any]      ,     **__a						:							Dict  )      ->   List[Any]:







       '''simple docstring'''
       return self.tokenizer.decode(*__a      ,     **__a  )
 @contextmanager
 def A_   (   self						:							str  )      ->   List[Any]:







       '''simple docstring'''
       warnings.warn(
           '`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your '
           'labels by using the argument `text` of the regular `__call__` method (either in the same call as '
           'your images inputs, or in a separate call.'  )
       __snake_case :   Union[str, Any]													=    True
       __snake_case :   List[Any]													=    self.tokenizer
       yield
       __snake_case :   str													=    self.image_processor
       __snake_case :   Tuple													=    False
 def A_   (   self						:							List[str]      ,     __a						:							Tuple      ,     __a						:							Union[str, Any]=False      ,     __a						:							str=None  )      ->   List[str]:







       '''simple docstring'''
       if added_vocab is None:
             __snake_case :   List[Any]													=    self.tokenizer.get_added_vocab()

       __snake_case :   Tuple													=    {}

       while tokens:
             __snake_case :   Optional[Any]													=    re.search(r'<s_(.*?)>'      ,     __a      ,     re.IGNORECASE  )
             if start_token is None:
                   break
             __snake_case :   Any													=    start_token.group(1  )
             __snake_case :   Dict													=    re.search(rf'''</s_{key}>'''      ,     __a      ,     re.IGNORECASE  )
             __snake_case :   List[str]													=    start_token.group()
             if end_token is None:
                   __snake_case :   List[Any]													=    tokens.replace(__a      ,     ''  )
             else:
                   __snake_case :   int													=    end_token.group()
                   __snake_case :   Optional[int]													=    re.escape(__a  )
                   __snake_case :   List[str]													=    re.escape(__a  )
                   __snake_case :   Optional[int]													=    re.search(f'''{start_token_escaped}(.*?){end_token_escaped}'''      ,     __a      ,     re.IGNORECASE  )
                   if content is not None:
                         __snake_case :   Tuple													=    content.group(1  ).strip()
                         if r"<s_" in content and r"</s_" in content:  # non-leaf node
                               __snake_case :   str													=    self.tokenajson(__a      ,     is_inner_value=__a      ,     added_vocab=__a  )
                               if value:
                                     if len(__a  ) == 1:
                                           __snake_case :   int													=    value[0]
                                     __snake_case :   Union[str, Any]													=    value
                         else:  # leaf nodes
                               __snake_case :   Any													=    []
                               for leaf in content.split(r'<sep/>'  ):
                                     __snake_case :   str													=    leaf.strip()
                                     if leaf in added_vocab and leaf[0] == "<" and leaf[-2:] == "/>":
                                           __snake_case :   Any													=    leaf[1:-2]  # for categorical special tokens
                                     output[key].append(__a  )
                               if len(output[key]  ) == 1:
                                     __snake_case :   int													=    output[key][0]

                   __snake_case :   Union[str, Any]													=    tokens[tokens.find(__a  ) + len(__a  ) :].strip()
                   if tokens[:6] == r"<sep/>":  # non-leaf nodes
                         return [output] + self.tokenajson(tokens[6:]      ,     is_inner_value=__a      ,     added_vocab=__a  )

       if len(__a  ):
             return [output] if is_inner_value else output
       else:
             return [] if is_inner_value else {"text_sequence": tokens}
 @property
 def A_   (   self						:							Optional[int]  )      ->   Any:







       '''simple docstring'''
       warnings.warn(
           '`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead.'      ,     __a      ,     )
       return self.image_processor_class





 @property
 def A_   (   self						:							Union[str, Any]  )      ->   Optional[int]:







       '''simple docstring'''
       warnings.warn(
           '`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead.'      ,     __a      ,     )
       return self.image_processor



