'''simple docstring'''


from typing import Optional

import torch
import torch.utils.checkpoint
from torch import Tensor, nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...modeling_outputs import (
    BackboneOutput,
    BaseModelOutputWithNoAttention,
    BaseModelOutputWithPoolingAndNoAttention,
    ImageClassifierOutputWithNoAttention,
)
from ...modeling_utils import PreTrainedModel
from ...utils import (
    add_code_sample_docstrings,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
    replace_return_docstrings,
)
from ...utils.backbone_utils import BackboneMixin
from .configuration_resnet import ResNetConfig


__lowerCAmelCase	:       Optional[Any]						       =	logging.get_logger(__name__)

# General docstring
__lowerCAmelCase	:       Union[str, Any]						       =	"ResNetConfig"

# Base docstring
__lowerCAmelCase	:       str						       =	"microsoft/resnet-50"
__lowerCAmelCase	:       Dict						       =	[1, 2_048, 7, 7]

# Image classification docstring
__lowerCAmelCase	:       str						       =	"microsoft/resnet-50"
__lowerCAmelCase	:       Optional[Any]						       =	"tiger cat"

__lowerCAmelCase	:       Any						       =	[
    "microsoft/resnet-50",
    # See all resnet models at https://huggingface.co/models?filter=resnet
]






class 						A   (   nn.Module   ):





							def __init__(      self				:							List[Any]     ,		__a				:							int     ,		__a				:							int     ,		__a				:							int = 3     ,		__a				:							int = 1     ,		__a				:							str = "relu"	)							->	Tuple:
								super().__init__()
								__UpperCAmelCase							=			nn.Convad(
								    __a     ,		__a     ,		kernel_size=__a     ,		stride=__a     ,		padding=kernel_size // 2     ,		bias=__a	)
								__UpperCAmelCase							=			nn.BatchNormad(__a	)
								__UpperCAmelCase							=			ACTaFN[activation] if activation is not None else nn.Identity()






							def       snake_case__						(      self				:							Any     ,		__a				:							Tensor	)							->	Tensor:
								__UpperCAmelCase							=			self.convolution(__a	)
								__UpperCAmelCase							=			self.normalization(__a	)
								__UpperCAmelCase							=			self.activation(__a	)
								return hidden_state






class 						A   (   nn.Module   ):





							def __init__(      self				:							List[Any]     ,		__a				:							ResNetConfig	)							->	Tuple:
								super().__init__()
								__UpperCAmelCase							=			ResNetConvLayer(
								    config.num_channels     ,		config.embedding_size     ,		kernel_size=7     ,		stride=2     ,		activation=config.hidden_act	)
								__UpperCAmelCase							=			nn.MaxPoolad(kernel_size=3     ,		stride=2     ,		padding=1	)
								__UpperCAmelCase							=			config.num_channels






							def       snake_case__						(      self				:							Optional[Any]     ,		__a				:							Tensor	)							->	Tensor:
								__UpperCAmelCase							=			pixel_values.shape[1]
								if num_channels != self.num_channels:
									raise ValueError(
									    '''Make sure that the channel dimension of the pixel values match with the one set in the configuration.'''	)
								__UpperCAmelCase							=			self.embedder(__a	)
								__UpperCAmelCase							=			self.pooler(__a	)
								return embedding






class 						A   (   nn.Module   ):





							def __init__(      self				:							Any     ,		__a				:							int     ,		__a				:							int     ,		__a				:							int = 2	)							->	int:
								super().__init__()
								__UpperCAmelCase							=			nn.Convad(__a     ,		__a     ,		kernel_size=1     ,		stride=__a     ,		bias=__a	)
								__UpperCAmelCase							=			nn.BatchNormad(__a	)






							def       snake_case__						(      self				:							int     ,		__a				:							Tensor	)							->	Tensor:
								__UpperCAmelCase							=			self.convolution(__a	)
								__UpperCAmelCase							=			self.normalization(__a	)
								return hidden_state






class 						A   (   nn.Module   ):





							def __init__(      self				:							Optional[Any]     ,		__a				:							int     ,		__a				:							int     ,		__a				:							int = 1     ,		__a				:							str = "relu"	)							->	Any:
								super().__init__()
								__UpperCAmelCase							=			in_channels != out_channels or stride != 1
								__UpperCAmelCase							=			(
								    ResNetShortCut(__a     ,		__a     ,		stride=__a	) if should_apply_shortcut else nn.Identity()
								)
								__UpperCAmelCase							=			nn.Sequential(
								    ResNetConvLayer(__a     ,		__a     ,		stride=__a	)     ,		ResNetConvLayer(__a     ,		__a     ,		activation=__a	)     ,		)
								__UpperCAmelCase							=			ACTaFN[activation]






							def       snake_case__						(      self				:							Optional[int]     ,		__a				:							List[Any]	)							->	int:
								__UpperCAmelCase							=			hidden_state
								__UpperCAmelCase							=			self.layer(__a	)
								__UpperCAmelCase							=			self.shortcut(__a	)
								hidden_state += residual
								__UpperCAmelCase							=			self.activation(__a	)
								return hidden_state






class 						A   (   nn.Module   ):





							def __init__(      self				:							Optional[int]     ,		__a				:							int     ,		__a				:							int     ,		__a				:							int = 1     ,		__a				:							str = "relu"     ,		__a				:							int = 4	)							->	Optional[Any]:
								super().__init__()
								__UpperCAmelCase							=			in_channels != out_channels or stride != 1
								__UpperCAmelCase							=			out_channels // reduction
								__UpperCAmelCase							=			(
								    ResNetShortCut(__a     ,		__a     ,		stride=__a	) if should_apply_shortcut else nn.Identity()
								)
								__UpperCAmelCase							=			nn.Sequential(
								    ResNetConvLayer(__a     ,		__a     ,		kernel_size=1	)     ,		ResNetConvLayer(__a     ,		__a     ,		stride=__a	)     ,		ResNetConvLayer(__a     ,		__a     ,		kernel_size=1     ,		activation=__a	)     ,		)
								__UpperCAmelCase							=			ACTaFN[activation]






							def       snake_case__						(      self				:							Any     ,		__a				:							Optional[int]	)							->	int:
								__UpperCAmelCase							=			hidden_state
								__UpperCAmelCase							=			self.layer(__a	)
								__UpperCAmelCase							=			self.shortcut(__a	)
								hidden_state += residual
								__UpperCAmelCase							=			self.activation(__a	)
								return hidden_state






class 						A   (   nn.Module   ):





							def __init__(      self				:							int     ,		__a				:							ResNetConfig     ,		__a				:							int     ,		__a				:							int     ,		__a				:							int = 2     ,		__a				:							int = 2     ,		)							->	List[Any]:
								super().__init__()

								__UpperCAmelCase							=			ResNetBottleNeckLayer if config.layer_type == '''bottleneck''' else ResNetBasicLayer

								__UpperCAmelCase							=			nn.Sequential(
								    # downsampling is done in the first layer with stride of 2
								    layer(__a     ,		__a     ,		stride=__a     ,		activation=config.hidden_act	)     ,		*[layer(__a     ,		__a     ,		activation=config.hidden_act	) for _ in range(depth - 1	)]     ,		)






							def       snake_case__						(      self				:							str     ,		__a				:							Tensor	)							->	Tensor:
								__UpperCAmelCase							=			input
								for layer in self.layers:
									__UpperCAmelCase							=			layer(__a	)
								return hidden_state






class 						A   (   nn.Module   ):





							def __init__(      self				:							Optional[Any]     ,		__a				:							ResNetConfig	)							->	Any:
								super().__init__()
								__UpperCAmelCase							=			nn.ModuleList([]	)
								# based on `downsample_in_first_stage` the first layer of the first stage may or may not downsample the input
								self.stages.append(
								    ResNetStage(
								        __a     ,		config.embedding_size     ,		config.hidden_sizes[0]     ,		stride=2 if config.downsample_in_first_stage else 1     ,		depth=config.depths[0]     ,		)	)
								__UpperCAmelCase							=			zip(config.hidden_sizes     ,		config.hidden_sizes[1:]	)
								for (in_channels, out_channels), depth in zip(__a     ,		config.depths[1:]	):
									self.stages.append(ResNetStage(__a     ,		__a     ,		__a     ,		depth=__a	)	)






							def       snake_case__						(      self				:							Dict     ,		__a				:							Tensor     ,		__a				:							bool = False     ,		__a				:							bool = True	)							->	BaseModelOutputWithNoAttention:
								__UpperCAmelCase							=			() if output_hidden_states else None

								for stage_module in self.stages:
									if output_hidden_states:
										__UpperCAmelCase							=			hidden_states + (hidden_state,)

									__UpperCAmelCase							=			stage_module(__a	)

								if output_hidden_states:
									__UpperCAmelCase							=			hidden_states + (hidden_state,)

								if not return_dict:
									return tuple(v for v in [hidden_state, hidden_states] if v is not None	)

								return BaseModelOutputWithNoAttention(
								    last_hidden_state=__a     ,		hidden_states=__a     ,		)






class 						A   (   UpperCAmelCase   ):
							a_        =      ResNetConfig
							a_        =      '''resnet'''
							a_        =      '''pixel_values'''
							a_        =      True





							def       snake_case__						(      self				:							List[Any]     ,		__a				:							List[Any]	)							->	List[str]:
								if isinstance(__a     ,		nn.Convad	):
									nn.init.kaiming_normal_(module.weight     ,		mode='''fan_out'''     ,		nonlinearity='''relu'''	)
								elif isinstance(__a     ,		(nn.BatchNormad, nn.GroupNorm)	):
									nn.init.constant_(module.weight     ,		1	)
									nn.init.constant_(module.bias     ,		0	)






							def       snake_case__						(      self				:							List[Any]     ,		__a				:							Optional[int]     ,		__a				:							Optional[Any]=False	)							->	Dict:
								if isinstance(__a     ,		__a	):
									__UpperCAmelCase							=			value

__lowerCAmelCase	:       int						       =	r"\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it\n    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and\n    behavior.\n\n    Parameters:\n        config ([`ResNetConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n"

__lowerCAmelCase	:       List[Any]						       =	r"\n    Args:\n        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`ConvNextImageProcessor.__call__`] for details.\n\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for\n            more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n"






@add_start_docstrings(
    '''The bare ResNet model outputting raw features without any specific head on top.'''       ,	UpperCAmelCase       ,	)
class 						A   (   UpperCAmelCase   ):





							def __init__(      self				:							List[str]     ,		__a				:							Any	)							->	Dict:
								super().__init__(__a	)
								__UpperCAmelCase							=			config
								__UpperCAmelCase							=			ResNetEmbeddings(__a	)
								__UpperCAmelCase							=			ResNetEncoder(__a	)
								__UpperCAmelCase							=			nn.AdaptiveAvgPoolad((1, 1)	)
								# Initialize weights and apply final processing
								self.post_init()






							@add_start_docstrings_to_model_forward(__a	)
							@add_code_sample_docstrings(
							    checkpoint=_CHECKPOINT_FOR_DOC     ,		output_type=__a     ,		config_class=_CONFIG_FOR_DOC     ,		modality='''vision'''     ,		expected_output=_EXPECTED_OUTPUT_SHAPE     ,		)
							def       snake_case__						(      self				:							List[str]     ,		__a				:							Tensor     ,		__a				:							Optional[bool] = None     ,		__a				:							Optional[bool] = None	)							->	BaseModelOutputWithPoolingAndNoAttention:
								__UpperCAmelCase							=			(
								    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
								)
								__UpperCAmelCase							=			return_dict if return_dict is not None else self.config.use_return_dict

								__UpperCAmelCase							=			self.embedder(__a	)

								__UpperCAmelCase							=			self.encoder(
								    __a     ,		output_hidden_states=__a     ,		return_dict=__a	)

								__UpperCAmelCase							=			encoder_outputs[0]

								__UpperCAmelCase							=			self.pooler(__a	)

								if not return_dict:
									return (last_hidden_state, pooled_output) + encoder_outputs[1:]

								return BaseModelOutputWithPoolingAndNoAttention(
								    last_hidden_state=__a     ,		pooler_output=__a     ,		hidden_states=encoder_outputs.hidden_states     ,		)






@add_start_docstrings(
    '''
    ResNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
    ImageNet.
    '''       ,	UpperCAmelCase       ,	)
class 						A   (   UpperCAmelCase   ):





							def __init__(      self				:							Optional[int]     ,		__a				:							Tuple	)							->	Union[str, Any]:
								super().__init__(__a	)
								__UpperCAmelCase							=			config.num_labels
								__UpperCAmelCase							=			ResNetModel(__a	)
								# classification head
								__UpperCAmelCase							=			nn.Sequential(
								    nn.Flatten()     ,		nn.Linear(config.hidden_sizes[-1]     ,		config.num_labels	) if config.num_labels > 0 else nn.Identity()     ,		)
								# initialize weights and apply final processing
								self.post_init()






							@add_start_docstrings_to_model_forward(__a	)
							@add_code_sample_docstrings(
							    checkpoint=_IMAGE_CLASS_CHECKPOINT     ,		output_type=__a     ,		config_class=_CONFIG_FOR_DOC     ,		expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT     ,		)
							def       snake_case__						(      self				:							Union[str, Any]     ,		__a				:							Optional[torch.FloatTensor] = None     ,		__a				:							Optional[torch.LongTensor] = None     ,		__a				:							Optional[bool] = None     ,		__a				:							Optional[bool] = None     ,		)							->	ImageClassifierOutputWithNoAttention:
								__UpperCAmelCase							=			return_dict if return_dict is not None else self.config.use_return_dict

								__UpperCAmelCase							=			self.resnet(__a     ,		output_hidden_states=__a     ,		return_dict=__a	)

								__UpperCAmelCase							=			outputs.pooler_output if return_dict else outputs[1]

								__UpperCAmelCase							=			self.classifier(__a	)

								__UpperCAmelCase							=			None

								if labels is not None:
									if self.config.problem_type is None:
										if self.num_labels == 1:
											__UpperCAmelCase							=			'''regression'''
										elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
											__UpperCAmelCase							=			'''single_label_classification'''
										else:
											__UpperCAmelCase							=			'''multi_label_classification'''
									if self.config.problem_type == "regression":
										__UpperCAmelCase							=			MSELoss()
										if self.num_labels == 1:
											__UpperCAmelCase							=			loss_fct(logits.squeeze()     ,		labels.squeeze()	)
										else:
											__UpperCAmelCase							=			loss_fct(__a     ,		__a	)
									elif self.config.problem_type == "single_label_classification":
										__UpperCAmelCase							=			CrossEntropyLoss()
										__UpperCAmelCase							=			loss_fct(logits.view(-1     ,		self.num_labels	)     ,		labels.view(-1	)	)
									elif self.config.problem_type == "multi_label_classification":
										__UpperCAmelCase							=			BCEWithLogitsLoss()
										__UpperCAmelCase							=			loss_fct(__a     ,		__a	)

								if not return_dict:
									__UpperCAmelCase							=			(logits,) + outputs[2:]
									return (loss,) + output if loss is not None else output

								return ImageClassifierOutputWithNoAttention(loss=__a     ,		logits=__a     ,		hidden_states=outputs.hidden_states	)






@add_start_docstrings(
    '''
    ResNet backbone, to be used with frameworks like DETR and MaskFormer.
    '''       ,	UpperCAmelCase       ,	)
class 						A   (   UpperCAmelCase       ,	UpperCAmelCase   ):





							def __init__(      self				:							str     ,		__a				:							str	)							->	Optional[Any]:
								super().__init__(__a	)
								super()._init_backbone(__a	)

								__UpperCAmelCase							=			[config.embedding_size] + config.hidden_sizes
								__UpperCAmelCase							=			ResNetEmbeddings(__a	)
								__UpperCAmelCase							=			ResNetEncoder(__a	)

								# initialize weights and apply final processing
								self.post_init()






							@add_start_docstrings_to_model_forward(__a	)
							@replace_return_docstrings(output_type=__a     ,		config_class=_CONFIG_FOR_DOC	)
							def       snake_case__						(      self				:							Dict     ,		__a				:							Tensor     ,		__a				:							Optional[bool] = None     ,		__a				:							Optional[bool] = None	)							->	BackboneOutput:
								__UpperCAmelCase							=			return_dict if return_dict is not None else self.config.use_return_dict
								__UpperCAmelCase							=			(
								    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
								)

								__UpperCAmelCase							=			self.embedder(__a	)

								__UpperCAmelCase							=			self.encoder(__a     ,		output_hidden_states=__a     ,		return_dict=__a	)

								__UpperCAmelCase							=			outputs.hidden_states

								__UpperCAmelCase							=			()
								for idx, stage in enumerate(self.stage_names	):
									if stage in self.out_features:
										feature_maps += (hidden_states[idx],)

								if not return_dict:
									__UpperCAmelCase							=			(feature_maps,)
									if output_hidden_states:
										output += (outputs.hidden_states,)
									return output

								return BackboneOutput(
								    feature_maps=__a     ,		hidden_states=outputs.hidden_states if output_hidden_states else None     ,		attentions=__a     ,		)

'''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()
__lowerCAmelCase	:       Dict						       =	logging.get_logger(__name__)





def       lowerCAmelCase					(		UpperCamelCase__				:   Optional[Any]    ,					UpperCamelCase__				:   Union[str, Any]       ):


	"""simple docstring"""
	__UpperCAmelCase							=			[]
	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					(		UpperCamelCase__				:   Optional[Any]    ,					UpperCamelCase__				:   List[Any]       ):


	"""simple docstring"""
	for i in range(encoder_config.num_hidden_layers       ):
		# queries, keys and values (only weights, no biases)
		__UpperCAmelCase							=			state_dict.pop(f"""encoder.deit.blocks.{i}.attn.qkv.weight"""       )

		__UpperCAmelCase							=			in_proj_weight[
		    : encoder_config.hidden_size, :
		]
		__UpperCAmelCase							=			in_proj_weight[
		    encoder_config.hidden_size : encoder_config.hidden_size * 2, :
		]
		__UpperCAmelCase							=			in_proj_weight[
		    -encoder_config.hidden_size :, :
		]





def       lowerCAmelCase					(		UpperCamelCase__				:   Any    ,					UpperCamelCase__				:   Optional[Any]    ,					UpperCamelCase__				:   List[str]       ):


	"""simple docstring"""
	__UpperCAmelCase							=			dct.pop(UpperCamelCase__       )
	__UpperCAmelCase							=			val





def       lowerCAmelCase					(		UpperCamelCase__				:   Dict       ):


	"""simple docstring"""
	if "handwritten" in checkpoint_url:
		__UpperCAmelCase							=			'''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:
		__UpperCAmelCase							=			'''https://www.researchgate.net/profile/Dinh-Sang/publication/338099565/figure/fig8/AS:840413229350922@1577381536857/An-receipt-example-in-the-SROIE-2019-dataset_Q640.jpg'''
	__UpperCAmelCase							=			Image.open(requests.get(UpperCamelCase__    ,					stream=UpperCamelCase__       ).raw       ).convert('''RGB'''       )
	return im





@torch.no_grad()
def       lowerCAmelCase					(		UpperCamelCase__				:   Dict    ,					UpperCamelCase__				:   Union[str, Any]       ):


	"""simple docstring"""
	__UpperCAmelCase							=			ViTConfig(image_size=3_8_4    ,					qkv_bias=UpperCamelCase__       )
	__UpperCAmelCase							=			TrOCRConfig()

	# size of the architecture
	if "base" in checkpoint_url:
		__UpperCAmelCase							=			7_6_8
	elif "large" in checkpoint_url:
		# use ViT-large encoder
		__UpperCAmelCase							=			1_0_2_4
		__UpperCAmelCase							=			4_0_9_6
		__UpperCAmelCase							=			2_4
		__UpperCAmelCase							=			1_6
		__UpperCAmelCase							=			1_0_2_4
	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:
		__UpperCAmelCase							=			False
		__UpperCAmelCase							=			'''relu'''
		__UpperCAmelCase							=			1_0_2_4
		__UpperCAmelCase							=			True
		__UpperCAmelCase							=			False
		__UpperCAmelCase							=			False

	# load HuggingFace model
	__UpperCAmelCase							=			ViTModel(UpperCamelCase__    ,					add_pooling_layer=UpperCamelCase__       )
	__UpperCAmelCase							=			TrOCRForCausalLM(UpperCamelCase__       )
	__UpperCAmelCase							=			VisionEncoderDecoderModel(encoder=UpperCamelCase__    ,					decoder=UpperCamelCase__       )
	model.eval()

	# load state_dict of original model, rename some keys
	__UpperCAmelCase							=			torch.hub.load_state_dict_from_url(UpperCamelCase__    ,					map_location='''cpu'''    ,					check_hash=UpperCamelCase__       )['''model''']

	__UpperCAmelCase							=			create_rename_keys(UpperCamelCase__    ,					UpperCamelCase__       )
	for src, dest in rename_keys:
		rename_key(UpperCamelCase__    ,					UpperCamelCase__    ,					UpperCamelCase__       )
	read_in_q_k_v(UpperCamelCase__    ,					UpperCamelCase__       )

	# 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():
		__UpperCAmelCase							=			state_dict.pop(UpperCamelCase__       )
		if key.startswith('''decoder'''       ) and "output_projection" not in key:
			__UpperCAmelCase							=			val
		else:
			__UpperCAmelCase							=			val

    # load state dict
	model.load_state_dict(UpperCamelCase__       )

	# Check outputs on an image
	__UpperCAmelCase							=			ViTImageProcessor(size=encoder_config.image_size       )
	__UpperCAmelCase							=			RobertaTokenizer.from_pretrained('''roberta-large'''       )
	__UpperCAmelCase							=			TrOCRProcessor(UpperCamelCase__    ,					UpperCamelCase__       )

	__UpperCAmelCase							=			processor(images=prepare_img(UpperCamelCase__       )    ,					return_tensors='''pt'''       ).pixel_values

	# verify logits
	__UpperCAmelCase							=			torch.tensor([[model.config.decoder.decoder_start_token_id]]       )
	__UpperCAmelCase							=			model(pixel_values=UpperCamelCase__    ,					decoder_input_ids=UpperCamelCase__       )
	__UpperCAmelCase							=			outputs.logits

	__UpperCAmelCase							=			torch.Size([1, 1, 5_0_2_6_5]       )
	if "trocr-base-handwritten" in checkpoint_url:
		__UpperCAmelCase							=			torch.tensor(
		    [-1.45_02, -4.66_83, -0.53_47, -2.92_91, 9.14_35, -3.05_71, 8.97_64, 1.75_60, 8.73_58, -1.53_11]       )
	elif "trocr-large-handwritten" in checkpoint_url:
		__UpperCAmelCase							=			torch.tensor(
		    [-2.64_37, -1.31_29, -2.25_96, -5.34_55, 6.35_39, 1.76_04, 5.49_91, 1.47_02, 5.61_13, 2.01_70]       )
	elif "trocr-base-printed" in checkpoint_url:
		__UpperCAmelCase							=			torch.tensor(
		    [-5.68_16, -5.83_88, 1.13_98, -6.90_34, 6.85_05, -2.43_93, 1.22_84, -1.02_32, -1.96_61, -3.92_10]       )
	elif "trocr-large-printed" in checkpoint_url:
		__UpperCAmelCase							=			torch.tensor(
		    [-6.01_62, -7.09_59, 4.41_55, -5.10_63, 7.04_68, -3.16_31, 2.64_66, -0.30_81, -0.81_06, -1.75_35]       )

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

	Path(UpperCamelCase__       ).mkdir(exist_ok=UpperCamelCase__       )
	print(f"""Saving model to {pytorch_dump_folder_path}"""       )
	model.save_pretrained(UpperCamelCase__       )
	print(f"""Saving processor to {pytorch_dump_folder_path}"""       )
	processor.save_pretrained(UpperCamelCase__       )


if __name__ == "__main__":
					__lowerCAmelCase	:       Any						       =	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."
					)
					__lowerCAmelCase	:       Optional[int]						       =	parser.parse_args()
					convert_tr_ocr_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)
