



'''simple docstring'''

import inspect
import unittest
from math import floor

from transformers import CvtConfig
from transformers.file_utils import cached_property, is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_vision, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
	import torch

	from transformers import CvtForImageClassification, CvtModel
	from transformers.models.cvt.modeling_cvt import CVT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
	from PIL import Image

	from transformers import AutoImageProcessor


class 	a__     (				_lowercase	):


	def 				lowercase__  (self : Dict      )     ->			Optional[int]:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						str      					= self.config_class(**self.inputs_dict      )
					self.parent.assertTrue(hasattr(__UpperCAmelCase,					'''embed_dim'''      )      )
					self.parent.assertTrue(hasattr(__UpperCAmelCase,					'''num_heads'''      )      )



class 	a__     :


	def __init__(self : Dict,					__UpperCAmelCase : Union[str, Any],					__UpperCAmelCase : List[Any]=13,					__UpperCAmelCase : List[str]=64,					__UpperCAmelCase : Optional[Any]=3,					__UpperCAmelCase : int=[16, 48, 96],					__UpperCAmelCase : Tuple=[1, 3, 6],					__UpperCAmelCase : List[str]=[1, 2, 10],					__UpperCAmelCase : str=[7, 3, 3],					__UpperCAmelCase : List[str]=[4, 2, 2],					__UpperCAmelCase : Any=[2, 1, 1],					__UpperCAmelCase : int=[2, 2, 2],					__UpperCAmelCase : int=[False, False, True],					__UpperCAmelCase : Any=[0.0, 0.0, 0.0],					__UpperCAmelCase : Tuple=0.02,					__UpperCAmelCase : Optional[Any]=1e-12,					__UpperCAmelCase : Optional[int]=True,					__UpperCAmelCase : Tuple=True,					__UpperCAmelCase : List[str]=2,					)     ->			Optional[int]:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						Optional[int]      					= parent
					SCREAMING_SNAKE_CASE  :						List[Any]      					= batch_size
					SCREAMING_SNAKE_CASE  :						Optional[Any]      					= image_size
					SCREAMING_SNAKE_CASE  :						str      					= patch_sizes
					SCREAMING_SNAKE_CASE  :						List[Any]      					= patch_stride
					SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= patch_padding
					SCREAMING_SNAKE_CASE  :						Optional[int]      					= is_training
					SCREAMING_SNAKE_CASE  :						Tuple      					= use_labels
					SCREAMING_SNAKE_CASE  :						int      					= num_labels
					SCREAMING_SNAKE_CASE  :						int      					= num_channels
					SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= embed_dim
					SCREAMING_SNAKE_CASE  :						Dict      					= num_heads
					SCREAMING_SNAKE_CASE  :						Tuple      					= stride_kv
					SCREAMING_SNAKE_CASE  :						Optional[int]      					= depth
					SCREAMING_SNAKE_CASE  :						Optional[Any]      					= cls_token
					SCREAMING_SNAKE_CASE  :						Any      					= attention_drop_rate
					SCREAMING_SNAKE_CASE  :						Optional[Any]      					= initializer_range
					SCREAMING_SNAKE_CASE  :						Dict      					= layer_norm_eps


	def 				lowercase__  (self : Dict      )     ->			int:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						str      					= floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

					SCREAMING_SNAKE_CASE  :						int      					= None
					if self.use_labels:
									SCREAMING_SNAKE_CASE  :						Any      					= ids_tensor([self.batch_size],					self.num_labels      )

					SCREAMING_SNAKE_CASE  :						Optional[Any]      					= self.get_config()
					return config, pixel_values, labels


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

					"""simple docstring"""

					return CvtConfig(
					    image_size=self.image_size,					num_labels=self.num_labels,					num_channels=self.num_channels,					embed_dim=self.embed_dim,					num_heads=self.num_heads,					patch_sizes=self.patch_sizes,					patch_padding=self.patch_padding,					patch_stride=self.patch_stride,					stride_kv=self.stride_kv,					depth=self.depth,					cls_token=self.cls_token,					attention_drop_rate=self.attention_drop_rate,					initializer_range=self.initializer_range,					)


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

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						List[Any]      					= CvtModel(config=__UpperCAmelCase      )
					model.to(__UpperCAmelCase      )
					model.eval()
					SCREAMING_SNAKE_CASE  :						Dict      					= model(__UpperCAmelCase      )
					SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= (self.image_size, self.image_size)
					SCREAMING_SNAKE_CASE					,					SCREAMING_SNAKE_CASE  :						Any      					= image_size[0], image_size[1]
					for i in range(len(self.depth      )      ):
									SCREAMING_SNAKE_CASE  :						Any      					= floor(((height + 2 * self.patch_padding[i] - self.patch_sizes[i]) / self.patch_stride[i]) + 1      )
									SCREAMING_SNAKE_CASE  :						Optional[Any]      					= floor(((width + 2 * self.patch_padding[i] - self.patch_sizes[i]) / self.patch_stride[i]) + 1      )
					self.parent.assertEqual(result.last_hidden_state.shape,					(self.batch_size, self.embed_dim[-1], height, width)      )


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

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						Optional[int]      					= self.num_labels
					SCREAMING_SNAKE_CASE  :						Dict      					= CvtForImageClassification(__UpperCAmelCase      )
					model.to(__UpperCAmelCase      )
					model.eval()
					SCREAMING_SNAKE_CASE  :						List[Any]      					= model(__UpperCAmelCase,					labels=__UpperCAmelCase      )
					self.parent.assertEqual(result.logits.shape,					(self.batch_size, self.num_labels)      )






	def 				lowercase__  (self : Tuple      )     ->			Optional[int]:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= self.prepare_config_and_inputs()
					SCREAMING_SNAKE_CASE					,					SCREAMING_SNAKE_CASE					,					SCREAMING_SNAKE_CASE  :						List[str]      					= config_and_inputs
					SCREAMING_SNAKE_CASE  :						Optional[Any]      					= {'''pixel_values''': pixel_values}
					return config, inputs_dict



@require_torch
class 	a__     (				_lowercase,			_lowercase,			unittest.TestCase	):
	__magic_name__		:    Any	    =     (CvtModel, CvtForImageClassification) if is_torch_available() else ()
	__magic_name__		:    List[str]	    =     (
	    {"feature-extraction": CvtModel, "image-classification": CvtForImageClassification}
	    if is_torch_available()
	    else {}
	)

	__magic_name__		:    List[Any]	    =     False
	__magic_name__		:    Union[str, Any]	    =     False
	__magic_name__		:    List[Any]	    =     False
	__magic_name__		:    Optional[Any]	    =     False
	__magic_name__		:    List[Any]	    =     False


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

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= CvtModelTester(self      )
					SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= ConfigTester(self,					config_class=__UpperCAmelCase,					has_text_modality=__UpperCAmelCase,					hidden_size=37      )


	def 				lowercase__  (self : str      )     ->			int:

					"""simple docstring"""

					self.create_and_test_config_common_properties()
					self.config_tester.create_and_test_config_to_json_string()
					self.config_tester.create_and_test_config_to_json_file()
					self.config_tester.create_and_test_config_from_and_save_pretrained()
					self.config_tester.create_and_test_config_with_num_labels()
					self.config_tester.check_config_can_be_init_without_params()
					self.config_tester.check_config_arguments_init()


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

					"""simple docstring"""

					return


	@unittest.skip(reason='''Cvt does not output attentions'''      )
	def 				lowercase__  (self : Union[str, Any]      )     ->			Union[str, Any]:

					"""simple docstring"""

					pass


	@unittest.skip(reason='''Cvt does not use inputs_embeds'''      )
	def 				lowercase__  (self : Dict      )     ->			Tuple:

					"""simple docstring"""

					pass


	@unittest.skip(reason='''Cvt does not support input and output embeddings'''      )
	def 				lowercase__  (self : Tuple      )     ->			int:

					"""simple docstring"""

					pass


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

					"""simple docstring"""

					SCREAMING_SNAKE_CASE					,					SCREAMING_SNAKE_CASE  :						Any      					= self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
									SCREAMING_SNAKE_CASE  :						Tuple      					= model_class(__UpperCAmelCase      )
									SCREAMING_SNAKE_CASE  :						Optional[int]      					= inspect.signature(model.forward      )
									# signature.parameters is an OrderedDict => so arg_names order is deterministic
									SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= [*signature.parameters.keys()]

									SCREAMING_SNAKE_CASE  :						Any      					= ['''pixel_values''']
									self.assertListEqual(arg_names[:1],					__UpperCAmelCase      )


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

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						str      					= self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*__UpperCAmelCase      )


	def 				lowercase__  (self : str      )     ->			Tuple:

					"""simple docstring"""

					def check_hidden_states_output(__UpperCAmelCase : List[Any],					__UpperCAmelCase : Optional[Any],					__UpperCAmelCase : Dict      ):
									SCREAMING_SNAKE_CASE  :						Optional[Any]      					= model_class(__UpperCAmelCase      )
									model.to(__UpperCAmelCase      )
									model.eval()

									with torch.no_grad():
													SCREAMING_SNAKE_CASE  :						Any      					= model(**self._prepare_for_class(__UpperCAmelCase,					__UpperCAmelCase      )      )

									SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= outputs.hidden_states

									SCREAMING_SNAKE_CASE  :						List[str]      					= len(self.model_tester.depth      )
									self.assertEqual(len(__UpperCAmelCase      ),					__UpperCAmelCase      )

									# verify the first hidden states (first block)
									self.assertListEqual(
									    list(hidden_states[0].shape[-3:]      ),					[
									        self.model_tester.embed_dim[0],
									        self.model_tester.image_size // 4,
									        self.model_tester.image_size // 4,
									    ],					)

					SCREAMING_SNAKE_CASE					,					SCREAMING_SNAKE_CASE  :						int      					= self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
									SCREAMING_SNAKE_CASE  :						List[Any]      					= True
									check_hidden_states_output(__UpperCAmelCase,					__UpperCAmelCase,					__UpperCAmelCase      )

									# check that output_hidden_states also work using config
									del inputs_dict["output_hidden_states"]
									SCREAMING_SNAKE_CASE  :						Any      					= True

									check_hidden_states_output(__UpperCAmelCase,					__UpperCAmelCase,					__UpperCAmelCase      )


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

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						List[Any]      					= self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_image_classification(*__UpperCAmelCase      )


	@unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''      )
	def 				lowercase__  (self : Optional[int]      )     ->			str:

					"""simple docstring"""

					pass






	@slow
	def 				lowercase__  (self : Union[str, Any]      )     ->			int:

					"""simple docstring"""

					for model_name in CVT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									SCREAMING_SNAKE_CASE  :						Any      					= CvtModel.from_pretrained(__UpperCAmelCase      )
									self.assertIsNotNone(__UpperCAmelCase      )






def       __lowercase  ():
				SCREAMING_SNAKE_CASE  :						Dict      					= Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''       )
				return image




@require_torch
@require_vision
class 	a__     (				unittest.TestCase	):


	@cached_property
	def 				lowercase__  (self : Tuple      )     ->			str:

					"""simple docstring"""

					return AutoImageProcessor.from_pretrained(CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0]      )






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

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						List[Any]      					= CvtForImageClassification.from_pretrained(CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0]      ).to(__UpperCAmelCase      )

					SCREAMING_SNAKE_CASE  :						Optional[int]      					= self.default_image_processor
					SCREAMING_SNAKE_CASE  :						Tuple      					= prepare_img()
					SCREAMING_SNAKE_CASE  :						List[str]      					= image_processor(images=__UpperCAmelCase,					return_tensors='''pt'''      ).to(__UpperCAmelCase      )

					# forward pass
					with torch.no_grad():
									SCREAMING_SNAKE_CASE  :						str      					= model(**__UpperCAmelCase      )

					# verify the logits
					SCREAMING_SNAKE_CASE  :						List[Any]      					= torch.Size((1, 1000)      )
					self.assertEqual(outputs.logits.shape,					__UpperCAmelCase      )

					SCREAMING_SNAKE_CASE  :						Tuple      					= torch.tensor([0.9285, 0.9015, -0.3150]      ).to(__UpperCAmelCase      )

					self.assertTrue(torch.allclose(outputs.logits[0, :3],					__UpperCAmelCase,					atol=1e-4      )      )









'''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)




