






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 		SCREAMING_SNAKE_CASE							(					__lowerCAmelCase			) ->	Union[str, Any]:
     UpperCamelCase__			:   str       				=  model.config

     UpperCamelCase__			:   str       				=  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					,    )
     UpperCamelCase__			:   Optional[Any]       				=  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 		SCREAMING_SNAKE_CASE							(					__lowerCAmelCase			) ->	Any:

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

          if name == "encoder.norm.weight":
               UpperCamelCase__			:   str       				=  '''encoder.layernorm.weight'''
          if name == "encoder.norm.bias":
               UpperCamelCase__			:   str       				=  '''encoder.layernorm.bias'''

     return name







def 		SCREAMING_SNAKE_CASE							(					__lowerCAmelCase					,    __lowerCAmelCase			) ->	Any:

     for key in orig_state_dict.copy().keys():
          UpperCamelCase__			:   Optional[int]       				=  orig_state_dict.pop(__lowerCAmelCase			)

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

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

     return orig_state_dict







def 		SCREAMING_SNAKE_CASE							(					__lowerCAmelCase					,    __lowerCAmelCase=None					,    __lowerCAmelCase=False			) ->	Optional[int]:
     UpperCamelCase__			:   Any       				=  DonutModel.from_pretrained(__lowerCAmelCase			).eval()

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

     UpperCamelCase__			:   List[str]       				=  original_model.state_dict()
     UpperCamelCase__			:   Optional[Any]       				=  convert_state_dict(__lowerCAmelCase					,    __lowerCAmelCase			)
     model.load_state_dict(__lowerCAmelCase			)

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

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

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

     UpperCamelCase__			:   Any       				=  original_model.encoder.model.patch_embed(__lowerCAmelCase			)
     UpperCamelCase__			:   str       				=  model.encoder.embeddings(__lowerCAmelCase			)
     assert torch.allclose(__lowerCAmelCase					,    __lowerCAmelCase					,    atol=1E-3			)

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

     # verify decoder hidden states
     UpperCamelCase__			:   Optional[int]       				=  original_model(__lowerCAmelCase					,    __lowerCAmelCase					,    __lowerCAmelCase			).logits
     UpperCamelCase__			:   Optional[int]       				=  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__":
 lowerCamelCase	:							List[Any]				=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.''',
 )

 lowerCamelCase	:							Any				=parser.parse_args()
 convert_donut_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)




import json
from typing import List, Optional, Tuple

from tokenizers import normalizers

from ...tokenization_utils_base import BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import PaddingStrategy, logging
from .tokenization_realm import RealmTokenizer


A__								=						logging.get_logger(__name__)

A__								=						{'''vocab_file''': '''vocab.txt''', '''tokenizer_file''': '''tokenizer.json'''}

A__								=						{
    '''vocab_file''': {
        '''google/realm-cc-news-pretrained-embedder''': (
            '''https://huggingface.co/google/realm-cc-news-pretrained-embedder/resolve/main/vocab.txt'''
        ),
        '''google/realm-cc-news-pretrained-encoder''': (
            '''https://huggingface.co/google/realm-cc-news-pretrained-encoder/resolve/main/vocab.txt'''
        ),
        '''google/realm-cc-news-pretrained-scorer''': (
            '''https://huggingface.co/google/realm-cc-news-pretrained-scorer/resolve/main/vocab.txt'''
        ),
        '''google/realm-cc-news-pretrained-openqa''': (
            '''https://huggingface.co/google/realm-cc-news-pretrained-openqa/aresolve/main/vocab.txt'''
        ),
        '''google/realm-orqa-nq-openqa''': '''https://huggingface.co/google/realm-orqa-nq-openqa/resolve/main/vocab.txt''',
        '''google/realm-orqa-nq-reader''': '''https://huggingface.co/google/realm-orqa-nq-reader/resolve/main/vocab.txt''',
        '''google/realm-orqa-wq-openqa''': '''https://huggingface.co/google/realm-orqa-wq-openqa/resolve/main/vocab.txt''',
        '''google/realm-orqa-wq-reader''': '''https://huggingface.co/google/realm-orqa-wq-reader/resolve/main/vocab.txt''',
    },
    '''tokenizer_file''': {
        '''google/realm-cc-news-pretrained-embedder''': (
            '''https://huggingface.co/google/realm-cc-news-pretrained-embedder/resolve/main/tokenizer.jsont'''
        ),
        '''google/realm-cc-news-pretrained-encoder''': (
            '''https://huggingface.co/google/realm-cc-news-pretrained-encoder/resolve/main/tokenizer.json'''
        ),
        '''google/realm-cc-news-pretrained-scorer''': (
            '''https://huggingface.co/google/realm-cc-news-pretrained-scorer/resolve/main/tokenizer.json'''
        ),
        '''google/realm-cc-news-pretrained-openqa''': (
            '''https://huggingface.co/google/realm-cc-news-pretrained-openqa/aresolve/main/tokenizer.json'''
        ),
        '''google/realm-orqa-nq-openqa''': (
            '''https://huggingface.co/google/realm-orqa-nq-openqa/resolve/main/tokenizer.json'''
        ),
        '''google/realm-orqa-nq-reader''': (
            '''https://huggingface.co/google/realm-orqa-nq-reader/resolve/main/tokenizer.json'''
        ),
        '''google/realm-orqa-wq-openqa''': (
            '''https://huggingface.co/google/realm-orqa-wq-openqa/resolve/main/tokenizer.json'''
        ),
        '''google/realm-orqa-wq-reader''': (
            '''https://huggingface.co/google/realm-orqa-wq-reader/resolve/main/tokenizer.json'''
        ),
    },
}

A__								=						{
    '''google/realm-cc-news-pretrained-embedder''': 512,
    '''google/realm-cc-news-pretrained-encoder''': 512,
    '''google/realm-cc-news-pretrained-scorer''': 512,
    '''google/realm-cc-news-pretrained-openqa''': 512,
    '''google/realm-orqa-nq-openqa''': 512,
    '''google/realm-orqa-nq-reader''': 512,
    '''google/realm-orqa-wq-openqa''': 512,
    '''google/realm-orqa-wq-reader''': 512,
}

A__								=						{
    '''google/realm-cc-news-pretrained-embedder''': {'''do_lower_case''': True},
    '''google/realm-cc-news-pretrained-encoder''': {'''do_lower_case''': True},
    '''google/realm-cc-news-pretrained-scorer''': {'''do_lower_case''': True},
    '''google/realm-cc-news-pretrained-openqa''': {'''do_lower_case''': True},
    '''google/realm-orqa-nq-openqa''': {'''do_lower_case''': True},
    '''google/realm-orqa-nq-reader''': {'''do_lower_case''': True},
    '''google/realm-orqa-wq-openqa''': {'''do_lower_case''': True},
    '''google/realm-orqa-wq-reader''': {'''do_lower_case''': True},
}


class       a							(			__lowerCamelCase				):
		__lowerCAmelCase	:    List[Any]											=					VOCAB_FILES_NAMES
		__lowerCAmelCase	:    Optional[Any]											=					PRETRAINED_VOCAB_FILES_MAP
		__lowerCAmelCase	:    Optional[Any]											=					PRETRAINED_INIT_CONFIGURATION
		__lowerCAmelCase	:    Optional[Any]											=					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
		__lowerCAmelCase	:    Optional[Any]											=					RealmTokenizer



		def __init__(     self     :Optional[Any]   ,__lowercase     :Dict=None   ,__lowercase     :Optional[int]=None   ,__lowercase     :Optional[Any]=True   ,__lowercase     :Optional[int]="[UNK]"   ,__lowercase     :List[str]="[SEP]"   ,__lowercase     :List[str]="[PAD]"   ,__lowercase     :int="[CLS]"   ,__lowercase     :str="[MASK]"   ,__lowercase     :Dict=True   ,__lowercase     :List[str]=None   ,**__lowercase     :Any   ,):
						super().__init__(
						    __lowercase   ,tokenizer_file=__lowercase   ,do_lower_case=__lowercase   ,unk_token=__lowercase   ,sep_token=__lowercase   ,pad_token=__lowercase   ,cls_token=__lowercase   ,mask_token=__lowercase   ,tokenize_chinese_chars=__lowercase   ,strip_accents=__lowercase   ,**__lowercase   ,)

						snake_case__	:      List[str]     		=	json.loads(self.backend_tokenizer.normalizer.__getstate__()  )
						if (
						    normalizer_state.get('''lowercase'''   ,__lowercase  ) != do_lower_case
						    or normalizer_state.get('''strip_accents'''   ,__lowercase  ) != strip_accents
						    or normalizer_state.get('''handle_chinese_chars'''   ,__lowercase  ) != tokenize_chinese_chars
						):
										snake_case__	:      Optional[int]     		=	getattr(__lowercase   ,normalizer_state.pop('''type'''  )  )
										snake_case__	:      List[Any]     		=	do_lower_case
										snake_case__	:      Optional[Any]     		=	strip_accents
										snake_case__	:      List[str]     		=	tokenize_chinese_chars
										snake_case__	:      Dict     		=	normalizer_class(**__lowercase  )

						snake_case__	:      Tuple     		=	do_lower_case



		def 	__lowerCamelCase			(     self     :Union[str, Any]   ,__lowercase     :Union[str, Any]   ,**__lowercase     :Any  ):
						snake_case__	:      Dict     		=	PaddingStrategy.MAX_LENGTH

						snake_case__	:      List[str]     		=	text
						snake_case__	:      int     		=	kwargs.pop('''text_pair'''   ,__lowercase  )
						snake_case__	:      Optional[int]     		=	kwargs.pop('''return_tensors'''   ,__lowercase  )

						snake_case__	:      str     		=	{
						    '''input_ids''': [],
						    '''attention_mask''': [],
						    '''token_type_ids''': [],
						}

						for idx, candidate_text in enumerate(__lowercase  ):
										if batch_text_pair is not None:
														snake_case__	:      Optional[int]     		=	batch_text_pair[idx]
										else:
														snake_case__	:      Tuple     		=	None

										snake_case__	:      List[str]     		=	super().__call__(__lowercase   ,__lowercase   ,return_tensors=__lowercase   ,**__lowercase  )

										snake_case__	:      Optional[Any]     		=	encoded_candidates.get('''input_ids'''  )
										snake_case__	:      Optional[Any]     		=	encoded_candidates.get('''attention_mask'''  )
										snake_case__	:      List[Any]     		=	encoded_candidates.get('''token_type_ids'''  )

										if encoded_input_ids is not None:
														output_data["input_ids"].append(__lowercase  )
										if encoded_attention_mask is not None:
														output_data["attention_mask"].append(__lowercase  )
										if encoded_token_type_ids is not None:
														output_data["token_type_ids"].append(__lowercase  )

						snake_case__	:      Any     		=	{key: item for key, item in output_data.items() if len(__lowercase  ) != 0}

						return BatchEncoding(__lowercase   ,tensor_type=__lowercase  )



		def 	__lowerCamelCase			(     self     :List[Any]   ,__lowercase     :Tuple   ,__lowercase     :Tuple=None  ):
						snake_case__	:      List[str]     		=	[self.cls_token_id] + token_ids_a + [self.sep_token_id]

						if token_ids_a:
										output += token_ids_a + [self.sep_token_id]

						return output



		def 	__lowerCamelCase			(     self     :List[str]   ,__lowercase     :List[int]   ,__lowercase     :Optional[List[int]] = None  ):
						snake_case__	:      Tuple     		=	[self.sep_token_id]
						snake_case__	:      Union[str, Any]     		=	[self.cls_token_id]
						if token_ids_a is None:
										return len(cls + token_ids_a + sep  ) * [0]
						return len(cls + token_ids_a + sep  ) * [0] + len(token_ids_a + sep  ) * [1]






		def 	__lowerCamelCase			(     self     :Optional[Any]   ,__lowercase     :str   ,__lowercase     :Optional[str] = None  ):
						snake_case__	:      Tuple     		=	self._tokenizer.model.save(__lowercase   ,name=__lowercase  )
						return tuple(__lowercase  )

