





import argparse
import collections
import json
from pathlib import Path

import requests
import torch
import yaml
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    MobileViTImageProcessor,
    MobileViTVaConfig,
    MobileViTVaForImageClassification,
    MobileViTVaForSemanticSegmentation,
)
from transformers.utils import logging


logging.set_verbosity_info()
UpperCAmelCase_  							=				logging.get_logger(__name__)




def 	lowerCAmelCase_					(						__UpperCAmelCase:    Dict   )   ->			Union[str, Any]:
    print('''Loading config file...'''   )

    def flatten_yaml_as_dict(__UpperCAmelCase:    Optional[Any]	,     __UpperCAmelCase:    Optional[Any]=""	,     __UpperCAmelCase:    Union[str, Any]="."   ):
        UpperCamelCase__							:       List[Any]      =   []
        for k, v in d.items():
            UpperCamelCase__							:       List[str]      =   parent_key + sep + k if parent_key else k
            if isinstance(__UpperCAmelCase	,     collections.abc.MutableMapping   ):
                items.extend(flatten_yaml_as_dict(__UpperCAmelCase	,     __UpperCAmelCase	,     sep=__UpperCAmelCase   ).items()   )
            else:
                items.append((new_key, v)   )
        return dict(__UpperCAmelCase   )

    UpperCamelCase__							:       List[Any]      =   argparse.Namespace()
    with open(__UpperCAmelCase	,     '''r'''   ) as yaml_file:
        try:
            UpperCamelCase__							:       int      =   yaml.load(__UpperCAmelCase	,     Loader=yaml.FullLoader   )

            UpperCamelCase__							:       int      =   flatten_yaml_as_dict(__UpperCAmelCase   )
            for k, v in flat_cfg.items():
                setattr(__UpperCAmelCase	,     __UpperCAmelCase	,     __UpperCAmelCase   )
        except yaml.YAMLError as exc:
            logger.error('''Error while loading config file: {}. Error message: {}'''.format(__UpperCAmelCase	,     str(__UpperCAmelCase   )   )   )
    return config




def 	lowerCAmelCase_					(						__UpperCAmelCase:    Tuple	,     __UpperCAmelCase:    Any   )   ->			int:
    UpperCamelCase__							:       Optional[int]      =   MobileViTVaConfig()

    UpperCamelCase__							:       Optional[int]      =   False

    # dataset
    if task_name.startswith('''imagenet1k_'''   ):
        UpperCamelCase__							:       Union[str, Any]      =   1000
        if int(task_name.strip().split('''_'''   )[-1]   ) == 384:
            UpperCamelCase__							:       Tuple      =   384
        else:
            UpperCamelCase__							:       Tuple      =   256
        UpperCamelCase__							:       Union[str, Any]      =   '''imagenet-1k-id2label.json'''
    elif task_name.startswith('''imagenet21k_to_1k_'''   ):
        UpperCamelCase__							:       List[Any]      =   2_1000
        if int(task_name.strip().split('''_'''   )[-1]   ) == 384:
            UpperCamelCase__							:       int      =   384
        else:
            UpperCamelCase__							:       Optional[Any]      =   256
        UpperCamelCase__							:       Optional[Any]      =   '''imagenet-22k-id2label.json'''
    elif task_name.startswith('''ade20k_'''   ):
        UpperCamelCase__							:       Optional[Any]      =   151
        UpperCamelCase__							:       List[Any]      =   512
        UpperCamelCase__							:       int      =   '''ade20k-id2label.json'''
        UpperCamelCase__							:       str      =   True
    elif task_name.startswith('''voc_'''   ):
        UpperCamelCase__							:       Union[str, Any]      =   21
        UpperCamelCase__							:       List[str]      =   512
        UpperCamelCase__							:       Dict      =   '''pascal-voc-id2label.json'''
        UpperCamelCase__							:       Union[str, Any]      =   True

    # orig_config
    UpperCamelCase__							:       Dict      =   load_orig_config_file(__UpperCAmelCase   )
    assert getattr(__UpperCAmelCase	,     '''model.classification.name'''	,     -1   ) == "mobilevit_v2", "Invalid model"
    UpperCamelCase__							:       List[str]      =   getattr(__UpperCAmelCase	,     '''model.classification.mitv2.width_multiplier'''	,     1.0   )
    assert (
        getattr(__UpperCAmelCase	,     '''model.classification.mitv2.attn_norm_layer'''	,     -1   ) == "layer_norm_2d"
    ), "Norm layers other than layer_norm_2d is not supported"
    UpperCamelCase__							:       int      =   getattr(__UpperCAmelCase	,     '''model.classification.activation.name'''	,     '''swish'''   )
    # config.image_size == getattr(orig_config,  'sampler.bs.crop_size_width', 256)

    if is_segmentation_model:
        UpperCamelCase__							:       str      =   getattr(__UpperCAmelCase	,     '''model.segmentation.output_stride'''	,     16   )
        if "_deeplabv3" in task_name:
            UpperCamelCase__							:       List[Any]      =   getattr(__UpperCAmelCase	,     '''model.segmentation.deeplabv3.aspp_rates'''	,     [12, 24, 36]   )
            UpperCamelCase__							:       Dict      =   getattr(__UpperCAmelCase	,     '''model.segmentation.deeplabv3.aspp_out_channels'''	,     512   )
            UpperCamelCase__							:       Union[str, Any]      =   getattr(__UpperCAmelCase	,     '''model.segmentation.deeplabv3.aspp_dropout'''	,     0.1   )

    # id2label
    UpperCamelCase__							:       List[str]      =   '''huggingface/label-files'''
    UpperCamelCase__							:       Tuple      =   json.load(open(hf_hub_download(__UpperCAmelCase	,     __UpperCAmelCase	,     repo_type='''dataset'''   )	,     '''r'''   )   )
    UpperCamelCase__							:       Optional[int]      =   {int(__UpperCAmelCase   ): v for k, v in idalabel.items()}
    UpperCamelCase__							:       Optional[Any]      =   idalabel
    UpperCamelCase__							:       Any      =   {v: k for k, v in idalabel.items()}

    return config




def 	lowerCAmelCase_					(						__UpperCAmelCase:    int	,     __UpperCAmelCase:    Tuple	,     __UpperCAmelCase:    Optional[Any]   )   ->			int:
    UpperCamelCase__							:       List[str]      =   dct.pop(__UpperCAmelCase   )
    UpperCamelCase__							:       List[Any]      =   val




def 	lowerCAmelCase_					(						__UpperCAmelCase:    Dict	,     __UpperCAmelCase:    List[Any]=False   )   ->			str:
    if base_model:
        UpperCamelCase__							:       Union[str, Any]      =   ''''''
    else:
        UpperCamelCase__							:       str      =   '''mobilevitv2.'''

    UpperCamelCase__							:       Any      =   []
    for k in state_dict.keys():
        if k[:8] == "encoder.":
            UpperCamelCase__							:       Optional[Any]      =   k[8:]
        else:
            UpperCamelCase__							:       int      =   k

        if ".block." in k:
            UpperCamelCase__							:       Optional[Any]      =   k_new.replace('''.block.'''	,     '''.'''   )
        if ".conv." in k:
            UpperCamelCase__							:       str      =   k_new.replace('''.conv.'''	,     '''.convolution.'''   )
        if ".norm." in k:
            UpperCamelCase__							:       List[str]      =   k_new.replace('''.norm.'''	,     '''.normalization.'''   )

        if "conv_1." in k:
            UpperCamelCase__							:       List[Any]      =   k_new.replace('''conv_1.'''	,     f"{model_prefix}conv_stem."   )
        for i in [1, 2]:
            if f"layer_{i}." in k:
                UpperCamelCase__							:       Dict      =   k_new.replace(f"layer_{i}."	,     f"{model_prefix}encoder.layer.{i-1}.layer."   )
        if ".exp_1x1." in k:
            UpperCamelCase__							:       Optional[int]      =   k_new.replace('''.exp_1x1.'''	,     '''.expand_1x1.'''   )
        if ".red_1x1." in k:
            UpperCamelCase__							:       List[str]      =   k_new.replace('''.red_1x1.'''	,     '''.reduce_1x1.'''   )

        for i in [3, 4, 5]:
            if f"layer_{i}.0." in k:
                UpperCamelCase__							:       Union[str, Any]      =   k_new.replace(f"layer_{i}.0."	,     f"{model_prefix}encoder.layer.{i-1}.downsampling_layer."   )
            if f"layer_{i}.1.local_rep.0." in k:
                UpperCamelCase__							:       Optional[int]      =   k_new.replace(f"layer_{i}.1.local_rep.0."	,     f"{model_prefix}encoder.layer.{i-1}.conv_kxk."   )
            if f"layer_{i}.1.local_rep.1." in k:
                UpperCamelCase__							:       Union[str, Any]      =   k_new.replace(f"layer_{i}.1.local_rep.1."	,     f"{model_prefix}encoder.layer.{i-1}.conv_1x1."   )

        for i in [3, 4, 5]:
            if i == 3:
                UpperCamelCase__							:       List[Any]      =   [0, 1]
            elif i == 4:
                UpperCamelCase__							:       List[Any]      =   [0, 1, 2, 3]
            elif i == 5:
                UpperCamelCase__							:       Any      =   [0, 1, 2]

            for j in j_in:
                if f"layer_{i}.1.global_rep.{j}." in k:
                    UpperCamelCase__							:       Any      =   k_new.replace(
                        f"layer_{i}.1.global_rep.{j}."	,     f"{model_prefix}encoder.layer.{i-1}.transformer.layer.{j}."   )
            if f"layer_{i}.1.global_rep.{j+1}." in k:
                UpperCamelCase__							:       Any      =   k_new.replace(
                    f"layer_{i}.1.global_rep.{j+1}."	,     f"{model_prefix}encoder.layer.{i-1}.layernorm."   )

            if f"layer_{i}.1.conv_proj." in k:
                UpperCamelCase__							:       Optional[int]      =   k_new.replace(f"layer_{i}.1.conv_proj."	,     f"{model_prefix}encoder.layer.{i-1}.conv_projection."   )

        if "pre_norm_attn.0." in k:
            UpperCamelCase__							:       str      =   k_new.replace('''pre_norm_attn.0.'''	,     '''layernorm_before.'''   )
        if "pre_norm_attn.1." in k:
            UpperCamelCase__							:       Any      =   k_new.replace('''pre_norm_attn.1.'''	,     '''attention.'''   )
        if "pre_norm_ffn.0." in k:
            UpperCamelCase__							:       Any      =   k_new.replace('''pre_norm_ffn.0.'''	,     '''layernorm_after.'''   )
        if "pre_norm_ffn.1." in k:
            UpperCamelCase__							:       Any      =   k_new.replace('''pre_norm_ffn.1.'''	,     '''ffn.conv1.'''   )
        if "pre_norm_ffn.3." in k:
            UpperCamelCase__							:       Optional[int]      =   k_new.replace('''pre_norm_ffn.3.'''	,     '''ffn.conv2.'''   )

        if "classifier.1." in k:
            UpperCamelCase__							:       List[str]      =   k_new.replace('''classifier.1.'''	,     '''classifier.'''   )

        if "seg_head." in k:
            UpperCamelCase__							:       Tuple      =   k_new.replace('''seg_head.'''	,     '''segmentation_head.'''   )
        if ".aspp_layer." in k:
            UpperCamelCase__							:       Optional[int]      =   k_new.replace('''.aspp_layer.'''	,     '''.'''   )
        if ".aspp_pool." in k:
            UpperCamelCase__							:       int      =   k_new.replace('''.aspp_pool.'''	,     '''.'''   )

        rename_keys.append((k, k_new)   )
    return rename_keys




def 	lowerCAmelCase_					(						__UpperCAmelCase:    Optional[Any]   )   ->			Dict:
    UpperCamelCase__							:       Any      =   []
    for k in state_dict.keys():
        if k.startswith('''seg_head.aux_head.'''   ):
            keys_to_ignore.append(__UpperCAmelCase   )
    for k in keys_to_ignore:
        state_dict.pop(__UpperCAmelCase	,     __UpperCAmelCase   )




def 	lowerCAmelCase_					(						)   ->			Union[str, Any]:
    UpperCamelCase__							:       Any      =   '''http://images.cocodataset.org/val2017/000000039769.jpg'''
    # url = "https://cdn.britannica.com/86/141086-050-9D7C75EE/Gulfstream-G450-business-jet-passengers.jpg"
    UpperCamelCase__							:       Optional[Any]      =   Image.open(requests.get(__UpperCAmelCase	,     stream=__UpperCAmelCase   ).raw   )
    return im




@torch.no_grad()
def 	lowerCAmelCase_					(						__UpperCAmelCase:    int	,     __UpperCAmelCase:    Union[str, Any]	,     __UpperCAmelCase:    Union[str, Any]	,     __UpperCAmelCase:    Dict   )   ->			Optional[int]:
    UpperCamelCase__							:       Any      =   get_mobilevitva_config(__UpperCAmelCase	,     __UpperCAmelCase   )

    # load original state_dict
    UpperCamelCase__							:       int      =   torch.load(__UpperCAmelCase	,     map_location='''cpu'''   )

    # load huggingface model
    if task_name.startswith('''ade20k_'''   ) or task_name.startswith('''voc_'''   ):
        UpperCamelCase__							:       Optional[int]      =   MobileViTVaForSemanticSegmentation(__UpperCAmelCase   ).eval()
        UpperCamelCase__							:       Optional[int]      =   False
    else:
        UpperCamelCase__							:       str      =   MobileViTVaForImageClassification(__UpperCAmelCase   ).eval()
        UpperCamelCase__							:       Dict      =   False

    # remove and rename some keys of load the original model
    UpperCamelCase__							:       Optional[int]      =   checkpoint
    remove_unused_keys(__UpperCAmelCase   )
    UpperCamelCase__							:       Optional[int]      =   create_rename_keys(__UpperCAmelCase	,     base_model=__UpperCAmelCase   )
    for rename_key_src, rename_key_dest in rename_keys:
        rename_key(__UpperCAmelCase	,     __UpperCAmelCase	,     __UpperCAmelCase   )

    # load modified state_dict
    model.load_state_dict(__UpperCAmelCase   )

    # Check outputs on an image, prepared by MobileViTImageProcessor
    UpperCamelCase__							:       Optional[Any]      =   MobileViTImageProcessor(crop_size=config.image_size	,     size=config.image_size + 32   )
    UpperCamelCase__							:       Dict      =   image_processor(images=prepare_img()	,     return_tensors='''pt'''   )
    UpperCamelCase__							:       Optional[int]      =   model(**__UpperCAmelCase   )

    # verify classification model
    if task_name.startswith('''imagenet'''   ):
        UpperCamelCase__							:       Optional[Any]      =   outputs.logits
        UpperCamelCase__							:       str      =   logits.argmax(-1   ).item()
        print('''Predicted class:'''	,     model.config.idalabel[predicted_class_idx]   )
        if task_name.startswith('''imagenet1k_256'''   ) and config.width_multiplier == 1.0:
            # expected_logits for base variant
            UpperCamelCase__							:       Optional[Any]      =   torch.tensor([-1.6_336e00, -7.3_204e-02, -5.1_883e-01]   )
            assert torch.allclose(logits[0, :3]	,     __UpperCAmelCase	,     atol=1e-4   )

    Path(__UpperCAmelCase   ).mkdir(exist_ok=__UpperCAmelCase   )
    print(f"Saving model {task_name} to {pytorch_dump_folder_path}"   )
    model.save_pretrained(__UpperCAmelCase   )
    print(f"Saving image processor to {pytorch_dump_folder_path}"   )
    image_processor.save_pretrained(__UpperCAmelCase   )


if __name__ == "__main__":
  UpperCAmelCase_  							=				argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      '--task',
      default='imagenet1k_256',
      type=str,
      help=(
          'Name of the task for which the MobileViTV2 model you\'d like to convert is trained on . '
          '\n                Classification (ImageNet-1k)\n                    - MobileViTV2 (256x256) : imagenet1k_256\n                    - MobileViTV2 (Trained on 256x256 and Finetuned on 384x384) : imagenet1k_384\n                    - MobileViTV2 (Trained on ImageNet-21k and Finetuned on ImageNet-1k 256x256) :\n                      imagenet21k_to_1k_256\n                    - MobileViTV2 (Trained on ImageNet-21k, Finetuned on ImageNet-1k 256x256, and Finetuned on\n                      ImageNet-1k 384x384) : imagenet21k_to_1k_384\n                Segmentation\n                    - ADE20K Dataset : ade20k_deeplabv3\n                    - Pascal VOC 2012 Dataset: voc_deeplabv3\n            '
      ),
      choices=[
          'imagenet1k_256',
          'imagenet1k_384',
          'imagenet21k_to_1k_256',
          'imagenet21k_to_1k_384',
          'ade20k_deeplabv3',
          'voc_deeplabv3',
      ],
  )

  parser.add_argument(
      '--orig_checkpoint_path', required=True, type=str, help='Path to the original state dict (.pt file).'
  )
  parser.add_argument('--orig_config_path', required=True, type=str, help='Path to the original config file.')
  parser.add_argument(
      '--pytorch_dump_folder_path', required=True, type=str, help='Path to the output PyTorch model directory.'
  )

  UpperCAmelCase_  							=				parser.parse_args()
  convert_mobilevitva_checkpoint(
      args.task, args.orig_checkpoint_path, args.orig_config_path, args.pytorch_dump_folder_path
  )








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

    UpperCamelCase__							:       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=128	,     )
    UpperCamelCase__							:       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 	lowerCAmelCase_					(						__UpperCAmelCase:    str   )   ->			Tuple:
    if "encoder.model" in name:
        UpperCamelCase__							:       str      =   name.replace('''encoder.model'''	,     '''encoder'''   )
    if "decoder.model" in name:
        UpperCamelCase__							:       Optional[int]      =   name.replace('''decoder.model'''	,     '''decoder'''   )
    if "patch_embed.proj" in name:
        UpperCamelCase__							:       Dict      =   name.replace('''patch_embed.proj'''	,     '''embeddings.patch_embeddings.projection'''   )
    if "patch_embed.norm" in name:
        UpperCamelCase__							:       Optional[int]      =   name.replace('''patch_embed.norm'''	,     '''embeddings.norm'''   )
    if name.startswith('''encoder'''   ):
        if "layers" in name:
            UpperCamelCase__							:       Optional[Any]      =   '''encoder.''' + name
        if "attn.proj" in name:
            UpperCamelCase__							:       Optional[int]      =   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__							:       List[str]      =   name.replace('''norm1'''	,     '''layernorm_before'''   )
        if "norm2" in name:
            UpperCamelCase__							:       Union[str, Any]      =   name.replace('''norm2'''	,     '''layernorm_after'''   )
        if "mlp.fc1" in name:
            UpperCamelCase__							:       Any      =   name.replace('''mlp.fc1'''	,     '''intermediate.dense'''   )
        if "mlp.fc2" in name:
            UpperCamelCase__							:       List[str]      =   name.replace('''mlp.fc2'''	,     '''output.dense'''   )

        if name == "encoder.norm.weight":
            UpperCamelCase__							:       Optional[Any]      =   '''encoder.layernorm.weight'''
        if name == "encoder.norm.bias":
            UpperCamelCase__							:       Any      =   '''encoder.layernorm.bias'''

    return name




def 	lowerCAmelCase_					(						__UpperCAmelCase:    str	,     __UpperCAmelCase:    str   )   ->			Dict:
    for key in orig_state_dict.copy().keys():
        UpperCamelCase__							:       Dict      =   orig_state_dict.pop(__UpperCAmelCase   )

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

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

    return orig_state_dict




def 	lowerCAmelCase_					(						__UpperCAmelCase:    Optional[Any]	,     __UpperCAmelCase:    List[Any]=None	,     __UpperCAmelCase:    Tuple=False   )   ->			Optional[int]:
    # load original model
    UpperCamelCase__							:       List[Any]      =   DonutModel.from_pretrained(__UpperCAmelCase   ).eval()

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

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

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

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

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

    UpperCamelCase__							:       List[str]      =   original_model.encoder.model.patch_embed(__UpperCAmelCase   )
    UpperCamelCase__     ,UpperCamelCase__							:       List[str]      =   model.encoder.embeddings(__UpperCAmelCase   )
    assert torch.allclose(__UpperCAmelCase	,     __UpperCAmelCase	,     atol=1e-3   )

    # verify encoder hidden states
    UpperCamelCase__							:       Dict      =   original_model.encoder(__UpperCAmelCase   )
    UpperCamelCase__							:       Optional[Any]      =   model.encoder(__UpperCAmelCase   ).last_hidden_state
    assert torch.allclose(__UpperCAmelCase	,     __UpperCAmelCase	,     atol=1e-2   )

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

