





"""simple docstring"""





import argparse
import json

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

from transformers import ViTImageProcessor, ViTMSNConfig, ViTMSNModel
from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD


torch.set_grad_enabled(False)



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							Tuple    ,						SCREAMING_SNAKE_CASE       :							List[Any]=False   ):


							'''simple docstring'''







							__lowerCamelCase      :      Optional[int]         =[]
							for i in range(config.num_hidden_layers   ):
														# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
														rename_keys.append((F'module.blocks.{i}.norm1.weight', F'vit.encoder.layer.{i}.layernorm_before.weight')   )
														rename_keys.append((F'module.blocks.{i}.norm1.bias', F'vit.encoder.layer.{i}.layernorm_before.bias')   )
														rename_keys.append(
														    (F'module.blocks.{i}.attn.proj.weight', F'vit.encoder.layer.{i}.attention.output.dense.weight')   )
														rename_keys.append((F'module.blocks.{i}.attn.proj.bias', F'vit.encoder.layer.{i}.attention.output.dense.bias')   )
														rename_keys.append((F'module.blocks.{i}.norm2.weight', F'vit.encoder.layer.{i}.layernorm_after.weight')   )
														rename_keys.append((F'module.blocks.{i}.norm2.bias', F'vit.encoder.layer.{i}.layernorm_after.bias')   )
														rename_keys.append((F'module.blocks.{i}.mlp.fc1.weight', F'vit.encoder.layer.{i}.intermediate.dense.weight')   )
														rename_keys.append((F'module.blocks.{i}.mlp.fc1.bias', F'vit.encoder.layer.{i}.intermediate.dense.bias')   )
														rename_keys.append((F'module.blocks.{i}.mlp.fc2.weight', F'vit.encoder.layer.{i}.output.dense.weight')   )
														rename_keys.append((F'module.blocks.{i}.mlp.fc2.bias', F'vit.encoder.layer.{i}.output.dense.bias')   )

							# projection layer + position embeddings
							rename_keys.extend(
							    [
							        ('''module.cls_token''', '''vit.embeddings.cls_token'''),
							        ('''module.patch_embed.proj.weight''', '''vit.embeddings.patch_embeddings.projection.weight'''),
							        ('''module.patch_embed.proj.bias''', '''vit.embeddings.patch_embeddings.projection.bias'''),
							        ('''module.pos_embed''', '''vit.embeddings.position_embeddings'''),
							    ]   )

							if base_model:
														# layernorm + pooler
														rename_keys.extend(
														    [
														        ('''module.norm.weight''', '''layernorm.weight'''),
														        ('''module.norm.bias''', '''layernorm.bias'''),
														    ]   )

														# if just the base model, we should remove "vit" from all keys that start with "vit"
														__lowerCamelCase      :      List[Any]         =[(pair[0], pair[1][4:]) if pair[1].startswith('''vit'''   ) else pair for pair in rename_keys]
							else:
														# layernorm + classification head
														rename_keys.extend(
														    [
														        ('''norm.weight''', '''vit.layernorm.weight'''),
														        ('''norm.bias''', '''vit.layernorm.bias'''),
														        ('''head.weight''', '''classifier.weight'''),
														        ('''head.bias''', '''classifier.bias'''),
														    ]   )

							return rename_keys



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							int    ,						SCREAMING_SNAKE_CASE       :							List[str]    ,						SCREAMING_SNAKE_CASE       :							Any=False   ):


							'''simple docstring'''







							for i in range(config.num_hidden_layers   ):
														if base_model:
																					__lowerCamelCase      :      Optional[int]         =''''''
														else:
																					__lowerCamelCase      :      List[str]         ='''vit.'''
														# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
														__lowerCamelCase      :      int         =state_dict.pop(F'module.blocks.{i}.attn.qkv.weight'   )
														__lowerCamelCase      :      Dict         =state_dict.pop(F'module.blocks.{i}.attn.qkv.bias'   )
														# next, add query, keys and values (in that order) to the state dict
														__lowerCamelCase      :      Optional[Any]         =in_proj_weight[
														    : config.hidden_size, :
														]
														__lowerCamelCase      :      Dict         =in_proj_bias[: config.hidden_size]
														__lowerCamelCase      :      Union[str, Any]         =in_proj_weight[
														    config.hidden_size : config.hidden_size * 2, :
														]
														__lowerCamelCase      :      Optional[int]         =in_proj_bias[
														    config.hidden_size : config.hidden_size * 2
														]
														__lowerCamelCase      :      Optional[Any]         =in_proj_weight[
														    -config.hidden_size :, :
														]
														__lowerCamelCase      :      str         =in_proj_bias[-config.hidden_size :]



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							Optional[int]   ):


							'''simple docstring'''







							__lowerCamelCase      :      Any         =['''head.weight''', '''head.bias''']
							for k in ignore_keys:
														state_dict.pop(SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE   )



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							List[Any]   ):


							'''simple docstring'''







							__lowerCamelCase      :      Optional[int]         =[
							    '''module.fc.fc1.weight''',
							    '''module.fc.fc1.bias''',
							    '''module.fc.bn1.weight''',
							    '''module.fc.bn1.bias''',
							    '''module.fc.bn1.running_mean''',
							    '''module.fc.bn1.running_var''',
							    '''module.fc.bn1.num_batches_tracked''',
							    '''module.fc.fc2.weight''',
							    '''module.fc.fc2.bias''',
							    '''module.fc.bn2.weight''',
							    '''module.fc.bn2.bias''',
							    '''module.fc.bn2.running_mean''',
							    '''module.fc.bn2.running_var''',
							    '''module.fc.bn2.num_batches_tracked''',
							    '''module.fc.fc3.weight''',
							    '''module.fc.fc3.bias''',
							]
							for k in ignore_keys:
														state_dict.pop(SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE   )



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							str    ,						SCREAMING_SNAKE_CASE       :							int    ,						SCREAMING_SNAKE_CASE       :							Any   ):


							'''simple docstring'''







							__lowerCamelCase      :      List[Any]         =dct.pop(SCREAMING_SNAKE_CASE   )
							__lowerCamelCase      :      Optional[Any]         =val



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							List[str]    ,						SCREAMING_SNAKE_CASE       :							Any   ):


							'''simple docstring'''







							__lowerCamelCase      :      Any         =ViTMSNConfig()
							__lowerCamelCase      :      int         =1000

							__lowerCamelCase      :      Dict         ='''datasets/huggingface/label-files'''
							__lowerCamelCase      :      Dict         ='''imagenet-1k-id2label.json'''
							__lowerCamelCase      :      Tuple         =json.load(open(hf_hub_download(SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE   )    ,						'''r'''   )   )
							__lowerCamelCase      :      List[Any]         ={int(SCREAMING_SNAKE_CASE   ): v for k, v in idalabel.items()}
							__lowerCamelCase      :      int         =idalabel
							__lowerCamelCase      :      List[str]         ={v: k for k, v in idalabel.items()}

							if "s16" in checkpoint_url:
														__lowerCamelCase      :      Any         =384
														__lowerCamelCase      :      Union[str, Any]         =1536
														__lowerCamelCase      :      Any         =6
							elif "l16" in checkpoint_url:
														__lowerCamelCase      :      Dict         =1024
														__lowerCamelCase      :      Any         =4096
														__lowerCamelCase      :      List[Any]         =24
														__lowerCamelCase      :      Union[str, Any]         =16
														__lowerCamelCase      :      str         =0.1
							elif "b4" in checkpoint_url:
														__lowerCamelCase      :      Dict         =4
							elif "l7" in checkpoint_url:
														__lowerCamelCase      :      List[str]         =7
														__lowerCamelCase      :      str         =1024
														__lowerCamelCase      :      Tuple         =4096
														__lowerCamelCase      :      Union[str, Any]         =24
														__lowerCamelCase      :      Tuple         =16
														__lowerCamelCase      :      Optional[int]         =0.1

							__lowerCamelCase      :      Union[str, Any]         =ViTMSNModel(SCREAMING_SNAKE_CASE   )

							__lowerCamelCase      :      Tuple         =torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE    ,						map_location='''cpu'''   )['''target_encoder''']

							__lowerCamelCase      :      str         =ViTImageProcessor(size=config.image_size   )

							remove_projection_head(SCREAMING_SNAKE_CASE   )
							__lowerCamelCase      :      Optional[Any]         =create_rename_keys(SCREAMING_SNAKE_CASE    ,						base_model=SCREAMING_SNAKE_CASE   )

							for src, dest in rename_keys:
														rename_key(SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE   )
							read_in_q_k_v(SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE    ,						base_model=SCREAMING_SNAKE_CASE   )

							model.load_state_dict(SCREAMING_SNAKE_CASE   )
							model.eval()

							__lowerCamelCase      :      Dict         ='''http://images.cocodataset.org/val2017/000000039769.jpg'''

							__lowerCamelCase      :      Tuple         =Image.open(requests.get(SCREAMING_SNAKE_CASE    ,						stream=SCREAMING_SNAKE_CASE   ).raw   )
							__lowerCamelCase      :      int         =ViTImageProcessor(
							    size=config.image_size    ,						image_mean=SCREAMING_SNAKE_CASE    ,						image_std=SCREAMING_SNAKE_CASE   )
							__lowerCamelCase      :      Union[str, Any]         =image_processor(images=SCREAMING_SNAKE_CASE    ,						return_tensors='''pt'''   )

							# forward pass
							torch.manual_seed(2   )
							__lowerCamelCase      :      Dict         =model(**SCREAMING_SNAKE_CASE   )
							__lowerCamelCase      :      Any         =outputs.last_hidden_state

							# The following Colab Notebook was used to generate these outputs:
							# https://colab.research.google.com/gist/sayakpaul/3672419a04f5997827503fd84079bdd1/scratchpad.ipynb
							if "s16" in checkpoint_url:
														__lowerCamelCase      :      Union[str, Any]         =torch.tensor([[-1.0_915, -1.4_876, -1.1_809]]   )
							elif "b16" in checkpoint_url:
														__lowerCamelCase      :      List[Any]         =torch.tensor([[14.2_889, -18.9_045, 11.7_281]]   )
							elif "l16" in checkpoint_url:
														__lowerCamelCase      :      int         =torch.tensor([[41.5_028, -22.8_681, 45.6_475]]   )
							elif "b4" in checkpoint_url:
														__lowerCamelCase      :      Union[str, Any]         =torch.tensor([[-4.3_868, 5.2_932, -0.4_137]]   )
							else:
														__lowerCamelCase      :      str         =torch.tensor([[-0.1_792, -0.6_465, 2.4_263]]   )

							# verify logits
							assert torch.allclose(last_hidden_state[:, 0, :3]    ,						SCREAMING_SNAKE_CASE    ,						atol=1E-4   )

							print(F'Saving model to {pytorch_dump_folder_path}'   )
							model.save_pretrained(SCREAMING_SNAKE_CASE   )

							print(F'Saving image processor to {pytorch_dump_folder_path}'   )
							image_processor.save_pretrained(SCREAMING_SNAKE_CASE   )


if __name__ == "__main__":
							_UpperCamelCase							=      argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    '--checkpoint_url',
							    default='https://dl.fbaipublicfiles.com/msn/vits16_800ep.pth.tar',
							    type=str,
							    help='URL of the checkpoint you\'d like to convert.',
							)
							parser.add_argument(
							    '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
							)

							_UpperCamelCase							=      parser.parse_args()
							convert_vit_msn_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)











"""simple docstring"""





import argparse
from pathlib import Path

import requests
import torch
from PIL import Image

from transformers import (
    RobertaTokenizer,
    TrOCRConfig,
    TrOCRForCausalLM,
    TrOCRProcessor,
    VisionEncoderDecoderModel,
    ViTConfig,
    ViTImageProcessor,
    ViTModel,
)
from transformers.utils import logging


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



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							Union[str, Any]    ,						SCREAMING_SNAKE_CASE       :							List[str]   ):


							'''simple docstring'''







							__lowerCamelCase      :      Any         =[]
							for i in range(encoder_config.num_hidden_layers   ):
														# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
														rename_keys.append(
														    (F'encoder.deit.blocks.{i}.norm1.weight', F'encoder.encoder.layer.{i}.layernorm_before.weight')   )
														rename_keys.append((F'encoder.deit.blocks.{i}.norm1.bias', F'encoder.encoder.layer.{i}.layernorm_before.bias')   )
														rename_keys.append(
														    (F'encoder.deit.blocks.{i}.attn.proj.weight', F'encoder.encoder.layer.{i}.attention.output.dense.weight')   )
														rename_keys.append(
														    (F'encoder.deit.blocks.{i}.attn.proj.bias', F'encoder.encoder.layer.{i}.attention.output.dense.bias')   )
														rename_keys.append(
														    (F'encoder.deit.blocks.{i}.norm2.weight', F'encoder.encoder.layer.{i}.layernorm_after.weight')   )
														rename_keys.append((F'encoder.deit.blocks.{i}.norm2.bias', F'encoder.encoder.layer.{i}.layernorm_after.bias')   )
														rename_keys.append(
														    (F'encoder.deit.blocks.{i}.mlp.fc1.weight', F'encoder.encoder.layer.{i}.intermediate.dense.weight')   )
														rename_keys.append(
														    (F'encoder.deit.blocks.{i}.mlp.fc1.bias', F'encoder.encoder.layer.{i}.intermediate.dense.bias')   )
														rename_keys.append(
														    (F'encoder.deit.blocks.{i}.mlp.fc2.weight', F'encoder.encoder.layer.{i}.output.dense.weight')   )
														rename_keys.append((F'encoder.deit.blocks.{i}.mlp.fc2.bias', F'encoder.encoder.layer.{i}.output.dense.bias')   )

							# cls token, position embeddings and patch embeddings of encoder
							rename_keys.extend(
							    [
							        ('''encoder.deit.cls_token''', '''encoder.embeddings.cls_token'''),
							        ('''encoder.deit.pos_embed''', '''encoder.embeddings.position_embeddings'''),
							        ('''encoder.deit.patch_embed.proj.weight''', '''encoder.embeddings.patch_embeddings.projection.weight'''),
							        ('''encoder.deit.patch_embed.proj.bias''', '''encoder.embeddings.patch_embeddings.projection.bias'''),
							        ('''encoder.deit.norm.weight''', '''encoder.layernorm.weight'''),
							        ('''encoder.deit.norm.bias''', '''encoder.layernorm.bias'''),
							    ]   )

							return rename_keys



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							int    ,						SCREAMING_SNAKE_CASE       :							Optional[int]   ):


							'''simple docstring'''







							for i in range(encoder_config.num_hidden_layers   ):
														# queries, keys and values (only weights, no biases)
														__lowerCamelCase      :      str         =state_dict.pop(F'encoder.deit.blocks.{i}.attn.qkv.weight'   )

														__lowerCamelCase      :      Dict         =in_proj_weight[
														    : encoder_config.hidden_size, :
														]
														__lowerCamelCase      :      Any         =in_proj_weight[
														    encoder_config.hidden_size : encoder_config.hidden_size * 2, :
														]
														__lowerCamelCase      :      Dict         =in_proj_weight[
														    -encoder_config.hidden_size :, :
														]



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							List[Any]    ,						SCREAMING_SNAKE_CASE       :							Optional[int]    ,						SCREAMING_SNAKE_CASE       :							Dict   ):


							'''simple docstring'''







							__lowerCamelCase      :      Optional[int]         =dct.pop(SCREAMING_SNAKE_CASE   )
							__lowerCamelCase      :      Optional[Any]         =val



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							List[Any]   ):


							'''simple docstring'''







							if "handwritten" in checkpoint_url:
														__lowerCamelCase      :      List[Any]         ='''https://fki.tic.heia-fr.ch/static/img/a01-122-02-00.jpg'''  # industry
														# url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02-12.jpg" # have
														# url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02-10.jpg" # let
														# url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02.jpg"  #
														# url = "https://fki.tic.heia-fr.ch/static/img/a01-122.jpg"
							elif "printed" in checkpoint_url or "stage1" in checkpoint_url:
														__lowerCamelCase      :      List[Any]         ='''https://www.researchgate.net/profile/Dinh-Sang/publication/338099565/figure/fig8/AS:840413229350922@1577381536857/An-receipt-example-in-the-SROIE-2019-dataset_Q640.jpg'''
							__lowerCamelCase      :      int         =Image.open(requests.get(SCREAMING_SNAKE_CASE    ,						stream=SCREAMING_SNAKE_CASE   ).raw   ).convert('''RGB'''   )
							return im



@torch.no_grad()
def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							List[str]    ,						SCREAMING_SNAKE_CASE       :							Any   ):


							'''simple docstring'''







							__lowerCamelCase      :      Union[str, Any]         =ViTConfig(image_size=384    ,						qkv_bias=SCREAMING_SNAKE_CASE   )
							__lowerCamelCase      :      Optional[int]         =TrOCRConfig()

							# size of the architecture
							if "base" in checkpoint_url:
														__lowerCamelCase      :      Any         =768
							elif "large" in checkpoint_url:
														# use ViT-large encoder
														__lowerCamelCase      :      Optional[int]         =1024
														__lowerCamelCase      :      List[Any]         =4096
														__lowerCamelCase      :      Dict         =24
														__lowerCamelCase      :      Optional[Any]         =16
														__lowerCamelCase      :      Union[str, Any]         =1024
							else:
														raise ValueError('''Should either find \'base\' or \'large\' in checkpoint URL'''   )

							# the large-printed + stage1 checkpoints uses sinusoidal position embeddings, no layernorm afterwards
							if "large-printed" in checkpoint_url or "stage1" in checkpoint_url:
														__lowerCamelCase      :      Optional[int]         =False
														__lowerCamelCase      :      int         ='''relu'''
														__lowerCamelCase      :      Any         =1024
														__lowerCamelCase      :      str         =True
														__lowerCamelCase      :      List[Any]         =False
														__lowerCamelCase      :      Optional[int]         =False

							# load HuggingFace model
							__lowerCamelCase      :      Dict         =ViTModel(SCREAMING_SNAKE_CASE    ,						add_pooling_layer=SCREAMING_SNAKE_CASE   )
							__lowerCamelCase      :      Any         =TrOCRForCausalLM(SCREAMING_SNAKE_CASE   )
							__lowerCamelCase      :      List[Any]         =VisionEncoderDecoderModel(encoder=SCREAMING_SNAKE_CASE    ,						decoder=SCREAMING_SNAKE_CASE   )
							model.eval()

							# load state_dict of original model, rename some keys
							__lowerCamelCase      :      Optional[Any]         =torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE    ,						map_location='''cpu'''    ,						check_hash=SCREAMING_SNAKE_CASE   )['''model''']

							__lowerCamelCase      :      Optional[int]         =create_rename_keys(SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE   )
							for src, dest in rename_keys:
														rename_key(SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE   )
							read_in_q_k_v(SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE   )

							# remove parameters we don't need
							del state_dict["encoder.deit.head.weight"]
							del state_dict["encoder.deit.head.bias"]
							del state_dict["decoder.version"]

							# add prefix to decoder keys
							for key, val in state_dict.copy().items():
														__lowerCamelCase      :      List[Any]         =state_dict.pop(SCREAMING_SNAKE_CASE   )
														if key.startswith('''decoder'''   ) and "output_projection" not in key:
																					__lowerCamelCase      :      List[Any]         =val
														else:
																					__lowerCamelCase      :      Optional[Any]         =val

    # load state dict
							model.load_state_dict(SCREAMING_SNAKE_CASE   )

							# Check outputs on an image
							__lowerCamelCase      :      int         =ViTImageProcessor(size=encoder_config.image_size   )
							__lowerCamelCase      :      Dict         =RobertaTokenizer.from_pretrained('''roberta-large'''   )
							__lowerCamelCase      :      str         =TrOCRProcessor(SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE   )

							__lowerCamelCase      :      str         =processor(images=prepare_img(SCREAMING_SNAKE_CASE   )    ,						return_tensors='''pt'''   ).pixel_values

							# verify logits
							__lowerCamelCase      :      List[Any]         =torch.tensor([[model.config.decoder.decoder_start_token_id]]   )
							__lowerCamelCase      :      Optional[int]         =model(pixel_values=SCREAMING_SNAKE_CASE    ,						decoder_input_ids=SCREAMING_SNAKE_CASE   )
							__lowerCamelCase      :      Optional[Any]         =outputs.logits

							__lowerCamelCase      :      Optional[Any]         =torch.Size([1, 1, 50265]   )
							if "trocr-base-handwritten" in checkpoint_url:
														__lowerCamelCase      :      List[str]         =torch.tensor(
														    [-1.4_502, -4.6_683, -0.5_347, -2.9_291, 9.1_435, -3.0_571, 8.9_764, 1.7_560, 8.7_358, -1.5_311]   )
							elif "trocr-large-handwritten" in checkpoint_url:
														__lowerCamelCase      :      Any         =torch.tensor(
														    [-2.6_437, -1.3_129, -2.2_596, -5.3_455, 6.3_539, 1.7_604, 5.4_991, 1.4_702, 5.6_113, 2.0_170]   )
							elif "trocr-base-printed" in checkpoint_url:
														__lowerCamelCase      :      Optional[Any]         =torch.tensor(
														    [-5.6_816, -5.8_388, 1.1_398, -6.9_034, 6.8_505, -2.4_393, 1.2_284, -1.0_232, -1.9_661, -3.9_210]   )
							elif "trocr-large-printed" in checkpoint_url:
														__lowerCamelCase      :      Optional[Any]         =torch.tensor(
														    [-6.0_162, -7.0_959, 4.4_155, -5.1_063, 7.0_468, -3.1_631, 2.6_466, -0.3_081, -0.8_106, -1.7_535]   )

							if "stage1" not in checkpoint_url:
														assert logits.shape == expected_shape, "Shape of logits not as expected"
														assert torch.allclose(logits[0, 0, :10]    ,						SCREAMING_SNAKE_CASE    ,						atol=1E-3   ), "First elements of logits not as expected"

							Path(SCREAMING_SNAKE_CASE   ).mkdir(exist_ok=SCREAMING_SNAKE_CASE   )
							print(F'Saving model to {pytorch_dump_folder_path}'   )
							model.save_pretrained(SCREAMING_SNAKE_CASE   )
							print(F'Saving processor to {pytorch_dump_folder_path}'   )
							processor.save_pretrained(SCREAMING_SNAKE_CASE   )


if __name__ == "__main__":
							_UpperCamelCase							=      argparse.ArgumentParser()

							parser.add_argument(
							    '--checkpoint_url',
							    default='https://layoutlm.blob.core.windows.net/trocr/model_zoo/fairseq/trocr-base-handwritten.pt',
							    type=str,
							    help='URL to the original PyTorch checkpoint (.pth file).',
							)
							parser.add_argument(
							    '--pytorch_dump_folder_path', default=None, type=str, help='Path to the folder to output PyTorch model.'
							)
							_UpperCamelCase							=      parser.parse_args()
							convert_tr_ocr_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)




