



'''simple docstring'''

import re
from typing import Callable, List, Optional, Union

import tensorflow as tf


try:
	from tensorflow.keras.optimizers.legacy import Adam
except ImportError:
	from tensorflow.keras.optimizers import Adam


class 	a__     (				tf.keras.optimizers.schedules.LearningRateSchedule	):


	def __init__(self : List[Any],					__UpperCAmelCase : float,					__UpperCAmelCase : Callable,					__UpperCAmelCase : int,					__UpperCAmelCase : float = 1.0,					__UpperCAmelCase : str = None,					)     ->			Any:

					"""simple docstring"""

					super().__init__()
					SCREAMING_SNAKE_CASE  :						Tuple      					= initial_learning_rate
					SCREAMING_SNAKE_CASE  :						int      					= warmup_steps
					SCREAMING_SNAKE_CASE  :						int      					= power
					SCREAMING_SNAKE_CASE  :						Tuple      					= decay_schedule_fn
					SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= name


	def __call__(self : Dict,					__UpperCAmelCase : List[Any]      )     ->			Optional[int]:

					"""simple docstring"""

					with tf.name_scope(self.name or '''WarmUp'''      ) as name:
									# Implements polynomial warmup. i.e., if global_step < warmup_steps, the
									# learning rate will be `global_step/num_warmup_steps * init_lr`.
									SCREAMING_SNAKE_CASE  :						Tuple      					= tf.cast(__UpperCAmelCase,					tf.floataa      )
									SCREAMING_SNAKE_CASE  :						Any      					= tf.cast(self.warmup_steps,					tf.floataa      )
									SCREAMING_SNAKE_CASE  :						Any      					= global_step_float / warmup_steps_float
									SCREAMING_SNAKE_CASE  :						Tuple      					= self.initial_learning_rate * tf.math.pow(__UpperCAmelCase,					self.power      )
									return tf.cond(
									    global_step_float < warmup_steps_float,					lambda: warmup_learning_rate,					lambda: self.decay_schedule_fn(step - self.warmup_steps      ),					name=__UpperCAmelCase,					)






	def 				lowercase__  (self : Tuple      )     ->			str:

					"""simple docstring"""

					return {
					    "initial_learning_rate": self.initial_learning_rate,
					    "decay_schedule_fn": self.decay_schedule_fn,
					    "warmup_steps": self.warmup_steps,
					    "power": self.power,
					    "name": self.name,
					}






def       __lowercase  (_SCREAMING_SNAKE_CASE      :float    ,	_SCREAMING_SNAKE_CASE      :int    ,	_SCREAMING_SNAKE_CASE      :int    ,	_SCREAMING_SNAKE_CASE      :float = 0.0    ,	_SCREAMING_SNAKE_CASE      :float = 0.9    ,	_SCREAMING_SNAKE_CASE      :float = 0.999    ,	_SCREAMING_SNAKE_CASE      :float = 1E-8    ,	_SCREAMING_SNAKE_CASE      :Optional[float] = None    ,	_SCREAMING_SNAKE_CASE      :Optional[float] = None    ,	_SCREAMING_SNAKE_CASE      :float = 0.0    ,	_SCREAMING_SNAKE_CASE      :float = 1.0    ,	_SCREAMING_SNAKE_CASE      :Optional[List[str]] = None    ,	):
				SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= tf.keras.optimizers.schedules.PolynomialDecay(
				    initial_learning_rate=_SCREAMING_SNAKE_CASE    ,	decay_steps=num_train_steps - num_warmup_steps    ,	end_learning_rate=init_lr * min_lr_ratio    ,	power=_SCREAMING_SNAKE_CASE    ,	)
				if num_warmup_steps:
								SCREAMING_SNAKE_CASE  :						Any      					= WarmUp(
								    initial_learning_rate=_SCREAMING_SNAKE_CASE    ,	decay_schedule_fn=_SCREAMING_SNAKE_CASE    ,	warmup_steps=_SCREAMING_SNAKE_CASE    ,	)
				if weight_decay_rate > 0.0:
								SCREAMING_SNAKE_CASE  :						Optional[Any]      					= AdamWeightDecay(
								    learning_rate=_SCREAMING_SNAKE_CASE    ,	weight_decay_rate=_SCREAMING_SNAKE_CASE    ,	beta_a=_SCREAMING_SNAKE_CASE    ,	beta_a=_SCREAMING_SNAKE_CASE    ,	epsilon=_SCREAMING_SNAKE_CASE    ,	clipnorm=_SCREAMING_SNAKE_CASE    ,	global_clipnorm=_SCREAMING_SNAKE_CASE    ,	exclude_from_weight_decay=['''LayerNorm''', '''layer_norm''', '''bias''']    ,	include_in_weight_decay=_SCREAMING_SNAKE_CASE    ,	)
				else:
								SCREAMING_SNAKE_CASE  :						int      					= tf.keras.optimizers.Adam(
								    learning_rate=_SCREAMING_SNAKE_CASE    ,	beta_a=_SCREAMING_SNAKE_CASE    ,	beta_a=_SCREAMING_SNAKE_CASE    ,	epsilon=_SCREAMING_SNAKE_CASE    ,	clipnorm=_SCREAMING_SNAKE_CASE    ,	global_clipnorm=_SCREAMING_SNAKE_CASE    ,	)
				# We return the optimizer and the LR scheduler in order to better track the
				# evolution of the LR independently of the optimizer.
				return optimizer, lr_schedule


class 	a__     (				_lowercase	):


	def __init__(self : str,					__UpperCAmelCase : Union[float, tf.keras.optimizers.schedules.LearningRateSchedule] = 0.001,					__UpperCAmelCase : float = 0.9,					__UpperCAmelCase : float = 0.999,					__UpperCAmelCase : float = 1e-7,					__UpperCAmelCase : bool = False,					__UpperCAmelCase : float = 0.0,					__UpperCAmelCase : Optional[List[str]] = None,					__UpperCAmelCase : Optional[List[str]] = None,					__UpperCAmelCase : str = "AdamWeightDecay",					**__UpperCAmelCase : Tuple,					)     ->			List[str]:

					"""simple docstring"""

					super().__init__(__UpperCAmelCase,					__UpperCAmelCase,					__UpperCAmelCase,					__UpperCAmelCase,					__UpperCAmelCase,					__UpperCAmelCase,					**__UpperCAmelCase      )
					SCREAMING_SNAKE_CASE  :						Optional[int]      					= weight_decay_rate
					SCREAMING_SNAKE_CASE  :						Dict      					= include_in_weight_decay
					SCREAMING_SNAKE_CASE  :						Optional[Any]      					= exclude_from_weight_decay


	@classmethod
	def 				lowercase__  (cls : Union[str, Any],					__UpperCAmelCase : Union[str, Any]      )     ->			int:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						Any      					= {'''WarmUp''': WarmUp}
					return super(__UpperCAmelCase,					cls      ).from_config(__UpperCAmelCase,					custom_objects=__UpperCAmelCase      )


	def 				lowercase__  (self : str,					__UpperCAmelCase : List[str],					__UpperCAmelCase : List[str],					__UpperCAmelCase : Tuple      )     ->			str:

					"""simple docstring"""

					super(__UpperCAmelCase,					self      )._prepare_local(__UpperCAmelCase,					__UpperCAmelCase,					__UpperCAmelCase      )
					SCREAMING_SNAKE_CASE  :						Dict      					= tf.constant(
					    self.weight_decay_rate,					name='''adam_weight_decay_rate'''      )


	def 				lowercase__  (self : Optional[Any],					__UpperCAmelCase : Any,					__UpperCAmelCase : Any,					__UpperCAmelCase : Tuple      )     ->			str:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						Dict      					= self._do_use_weight_decay(var.name      )
					if do_decay:
									return var.assign_sub(
									    learning_rate * var * apply_state[(var.device, var.dtype.base_dtype)]['''weight_decay_rate'''],					use_locking=self._use_locking,					)
					return tf.no_op()


	def 				lowercase__  (self : Union[str, Any],					__UpperCAmelCase : int,					__UpperCAmelCase : List[str]=None,					**__UpperCAmelCase : int      )     ->			Optional[int]:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						Any      					= list(zip(*__UpperCAmelCase      )      )
					return super(__UpperCAmelCase,					self      ).apply_gradients(zip(__UpperCAmelCase,					__UpperCAmelCase      ),					name=__UpperCAmelCase,					**__UpperCAmelCase      )


	def 				lowercase__  (self : Tuple,					__UpperCAmelCase : int,					__UpperCAmelCase : Union[str, Any],					__UpperCAmelCase : Dict      )     ->			Union[str, Any]:

					"""simple docstring"""

					if apply_state is None:
									return self._decayed_lr_t[var_dtype], {}

					SCREAMING_SNAKE_CASE  :						List[Any]      					= apply_state or {}
					SCREAMING_SNAKE_CASE  :						Dict      					= apply_state.get((var_device, var_dtype)      )
					if coefficients is None:
									SCREAMING_SNAKE_CASE  :						Optional[Any]      					= self._fallback_apply_state(__UpperCAmelCase,					__UpperCAmelCase      )
									SCREAMING_SNAKE_CASE  :						List[Any]      					= coefficients

					return coefficients["lr_t"], {"apply_state": apply_state}


	def 				lowercase__  (self : Tuple,					__UpperCAmelCase : int,					__UpperCAmelCase : Union[str, Any],					__UpperCAmelCase : List[str]=None      )     ->			Optional[Any]:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						int      					= self._get_lr(var.device,					var.dtype.base_dtype,					__UpperCAmelCase      )
					SCREAMING_SNAKE_CASE  :						Dict      					= self._decay_weights_op(__UpperCAmelCase,					__UpperCAmelCase,					__UpperCAmelCase      )
					with tf.control_dependencies([decay]      ):
									return super(__UpperCAmelCase,					self      )._resource_apply_dense(__UpperCAmelCase,					__UpperCAmelCase,					**__UpperCAmelCase      )


	def 				lowercase__  (self : List[str],					__UpperCAmelCase : Union[str, Any],					__UpperCAmelCase : Any,					__UpperCAmelCase : int,					__UpperCAmelCase : Optional[int]=None      )     ->			Union[str, Any]:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						List[str]      					= self._get_lr(var.device,					var.dtype.base_dtype,					__UpperCAmelCase      )
					SCREAMING_SNAKE_CASE  :						Optional[int]      					= self._decay_weights_op(__UpperCAmelCase,					__UpperCAmelCase,					__UpperCAmelCase      )
					with tf.control_dependencies([decay]      ):
									return super(__UpperCAmelCase,					self      )._resource_apply_sparse(__UpperCAmelCase,					__UpperCAmelCase,					__UpperCAmelCase,					**__UpperCAmelCase      )


	def 				lowercase__  (self : List[str]      )     ->			Optional[Any]:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						int      					= super().get_config()
					config.update({'''weight_decay_rate''': self.weight_decay_rate}      )
					return config






	def 				lowercase__  (self : List[str],					__UpperCAmelCase : Union[str, Any]      )     ->			Tuple:

					"""simple docstring"""

					if self.weight_decay_rate == 0:
									return False

					if self._include_in_weight_decay:
									for r in self._include_in_weight_decay:
													if re.search(__UpperCAmelCase,					__UpperCAmelCase      ) is not None:
																	return True

					if self._exclude_from_weight_decay:
									for r in self._exclude_from_weight_decay:
													if re.search(__UpperCAmelCase,					__UpperCAmelCase      ) is not None:
																	return False
					return True



class 	a__     (				_lowercase	):


	def __init__(self : List[Any]      )     ->			Union[str, Any]:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						int      					= []
					SCREAMING_SNAKE_CASE  :						Optional[Any]      					= None


	@property
	def 				lowercase__  (self : Optional[Any]      )     ->			List[str]:

					"""simple docstring"""

					if self._accum_steps is None:
									SCREAMING_SNAKE_CASE  :						List[Any]      					= tf.Variable(
									    tf.constant(0,					dtype=tf.intaa      ),					trainable=__UpperCAmelCase,					synchronization=tf.VariableSynchronization.ON_READ,					aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,					)

					return self._accum_steps.value()


	@property
	def 				lowercase__  (self : List[Any]      )     ->			int:

					"""simple docstring"""

					if not self._gradients:
									raise ValueError('''The accumulator should be called first to initialize the gradients'''      )
					return [gradient.value() if gradient is not None else gradient for gradient in self._gradients]


	def __call__(self : Optional[Any],					__UpperCAmelCase : Optional[Any]      )     ->			Union[str, Any]:

					"""simple docstring"""

					if not self._gradients:
									SCREAMING_SNAKE_CASE  :						Optional[Any]      					= self.step  # Create the step variable.
									self._gradients.extend(
									    [
									        tf.Variable(
									            tf.zeros_like(__UpperCAmelCase      ),					trainable=__UpperCAmelCase,					synchronization=tf.VariableSynchronization.ON_READ,					aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,					)
									        if gradient is not None
									        else gradient
									        for gradient in gradients
									    ]      )
					if len(__UpperCAmelCase      ) != len(self._gradients      ):
									raise ValueError(F'''Expected {len(self._gradients      )} gradients, but got {len(__UpperCAmelCase      )}'''      )

					for accum_gradient, gradient in zip(self._gradients,					__UpperCAmelCase      ):
									if accum_gradient is not None and gradient is not None:
													accum_gradient.assign_add(__UpperCAmelCase      )

					self._accum_steps.assign_add(1      )






	def 				lowercase__  (self : List[Any]      )     ->			List[str]:

					"""simple docstring"""

					if not self._gradients:
									return
					self._accum_steps.assign(0      )
					for gradient in self._gradients:
									if gradient is not None:
													gradient.assign(tf.zeros_like(__UpperCAmelCase      )      )









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






def       __lowercase  (_SCREAMING_SNAKE_CASE      :Dict    ,	_SCREAMING_SNAKE_CASE      :int       ):
				SCREAMING_SNAKE_CASE  :						Dict      					= []
				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       __lowercase  (_SCREAMING_SNAKE_CASE      :List[Any]    ,	_SCREAMING_SNAKE_CASE      :Tuple       ):
				for i in range(encoder_config.num_hidden_layers       ):
								# queries, keys and values (only weights, no biases)
								SCREAMING_SNAKE_CASE  :						str      					= state_dict.pop(F'''encoder.deit.blocks.{i}.attn.qkv.weight'''       )

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






def       __lowercase  (_SCREAMING_SNAKE_CASE      :int    ,	_SCREAMING_SNAKE_CASE      :str    ,	_SCREAMING_SNAKE_CASE      :Optional[int]       ):
				SCREAMING_SNAKE_CASE  :						List[str]      					= dct.pop(_SCREAMING_SNAKE_CASE       )
				SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= val






def       __lowercase  (_SCREAMING_SNAKE_CASE      :List[str]       ):
				if "handwritten" in checkpoint_url:
								SCREAMING_SNAKE_CASE  :						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:
								SCREAMING_SNAKE_CASE  :						Union[str, 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'''
				SCREAMING_SNAKE_CASE  :						List[Any]      					= Image.open(requests.get(_SCREAMING_SNAKE_CASE    ,	stream=_SCREAMING_SNAKE_CASE       ).raw       ).convert('''RGB'''       )
				return im






@torch.no_grad()
def       __lowercase  (_SCREAMING_SNAKE_CASE      :List[Any]    ,	_SCREAMING_SNAKE_CASE      :Optional[int]       ):
				SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= ViTConfig(image_size=3_84    ,	qkv_bias=_SCREAMING_SNAKE_CASE       )
				SCREAMING_SNAKE_CASE  :						Any      					= TrOCRConfig()

				# size of the architecture
				if "base" in checkpoint_url:
								SCREAMING_SNAKE_CASE  :						List[Any]      					= 7_68
				elif "large" in checkpoint_url:
								# use ViT-large encoder
								SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= 10_24
								SCREAMING_SNAKE_CASE  :						Optional[Any]      					= 40_96
								SCREAMING_SNAKE_CASE  :						List[str]      					= 24
								SCREAMING_SNAKE_CASE  :						Any      					= 16
								SCREAMING_SNAKE_CASE  :						Optional[Any]      					= 10_24
				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:
								SCREAMING_SNAKE_CASE  :						Optional[Any]      					= False
								SCREAMING_SNAKE_CASE  :						Optional[Any]      					= '''relu'''
								SCREAMING_SNAKE_CASE  :						Tuple      					= 10_24
								SCREAMING_SNAKE_CASE  :						int      					= True
								SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= False
								SCREAMING_SNAKE_CASE  :						Any      					= False

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

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

				SCREAMING_SNAKE_CASE  :						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():
								SCREAMING_SNAKE_CASE  :						List[Any]      					= state_dict.pop(_SCREAMING_SNAKE_CASE       )
								if key.startswith('''decoder'''       ) and "output_projection" not in key:
												SCREAMING_SNAKE_CASE  :						Tuple      					= val
								else:
												SCREAMING_SNAKE_CASE  :						int      					= val

    # load state dict
				model.load_state_dict(_SCREAMING_SNAKE_CASE       )

				# Check outputs on an image
				SCREAMING_SNAKE_CASE  :						List[Any]      					= ViTImageProcessor(size=encoder_config.image_size       )
				SCREAMING_SNAKE_CASE  :						Tuple      					= RobertaTokenizer.from_pretrained('''roberta-large'''       )
				SCREAMING_SNAKE_CASE  :						Dict      					= TrOCRProcessor(_SCREAMING_SNAKE_CASE    ,	_SCREAMING_SNAKE_CASE       )

				SCREAMING_SNAKE_CASE  :						Tuple      					= processor(images=prepare_img(_SCREAMING_SNAKE_CASE       )    ,	return_tensors='''pt'''       ).pixel_values

				# verify logits
				SCREAMING_SNAKE_CASE  :						str      					= torch.tensor([[model.config.decoder.decoder_start_token_id]]       )
				SCREAMING_SNAKE_CASE  :						Dict      					= model(pixel_values=_SCREAMING_SNAKE_CASE    ,	decoder_input_ids=_SCREAMING_SNAKE_CASE       )
				SCREAMING_SNAKE_CASE  :						int      					= outputs.logits

				SCREAMING_SNAKE_CASE  :						Any      					= torch.Size([1, 1, 5_02_65]       )
				if "trocr-base-handwritten" in checkpoint_url:
								SCREAMING_SNAKE_CASE  :						Optional[int]      					= torch.tensor(
								    [-1.4502, -4.6683, -0.5347, -2.9291, 9.1435, -3.0571, 8.9764, 1.7560, 8.7358, -1.5311]       )
				elif "trocr-large-handwritten" in checkpoint_url:
								SCREAMING_SNAKE_CASE  :						Any      					= torch.tensor(
								    [-2.6437, -1.3129, -2.2596, -5.3455, 6.3539, 1.7604, 5.4991, 1.4702, 5.6113, 2.0170]       )
				elif "trocr-base-printed" in checkpoint_url:
								SCREAMING_SNAKE_CASE  :						str      					= torch.tensor(
								    [-5.6816, -5.8388, 1.1398, -6.9034, 6.8505, -2.4393, 1.2284, -1.0232, -1.9661, -3.9210]       )
				elif "trocr-large-printed" in checkpoint_url:
								SCREAMING_SNAKE_CASE  :						Optional[int]      					= torch.tensor(
								    [-6.0162, -7.0959, 4.4155, -5.1063, 7.0468, -3.1631, 2.6466, -0.3081, -0.8106, -1.7535]       )

				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__":
	snake_case_			=     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."""
	)
	snake_case_			=     parser.parse_args()
	convert_tr_ocr_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)




