



'''simple docstring'''

import os
import sys
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from huggingface_hub import HfFolder, delete_repo
from huggingface_hub.file_download import http_get
from requests.exceptions import HTTPError

from transformers import (
    AlbertTokenizer,
    AutoTokenizer,
    BertTokenizer,
    BertTokenizerFast,
    GPTaTokenizerFast,
    is_tokenizers_available,
)
from transformers.testing_utils import TOKEN, USER, is_staging_test, require_tokenizers
from transformers.tokenization_utils import Trie


sys.path.append(str(Path(__file__).parent.parent / """utils"""))

from test_module.custom_tokenization import CustomTokenizer  # noqa E402


if is_tokenizers_available():
	from test_module.custom_tokenization_fast import CustomTokenizerFast


class 	a__     (				unittest.TestCase	):


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

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						int      					= mock.Mock()
					SCREAMING_SNAKE_CASE  :						Tuple      					= 500
					SCREAMING_SNAKE_CASE  :						Dict      					= {}
					SCREAMING_SNAKE_CASE  :						Dict      					= HTTPError
					SCREAMING_SNAKE_CASE  :						Optional[Any]      					= {}

					# Download this model to make sure it's in the cache.
					SCREAMING_SNAKE_CASE  :						List[Any]      					= BertTokenizer.from_pretrained('''hf-internal-testing/tiny-random-bert'''      )

					# Under the mock environment we get a 500 error when trying to reach the tokenizer.
					with mock.patch('''requests.Session.request''',					return_value=lowerCamelCase__      ) as mock_head:
									SCREAMING_SNAKE_CASE  :						Any      					= BertTokenizer.from_pretrained('''hf-internal-testing/tiny-random-bert'''      )
									# This check we did call the fake head request
									mock_head.assert_called()


	@require_tokenizers
	def 				lowercase__  (self : str      )     ->			Any:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= mock.Mock()
					SCREAMING_SNAKE_CASE  :						Any      					= 500
					SCREAMING_SNAKE_CASE  :						str      					= {}
					SCREAMING_SNAKE_CASE  :						Optional[Any]      					= HTTPError
					SCREAMING_SNAKE_CASE  :						Tuple      					= {}

					# Download this model to make sure it's in the cache.
					SCREAMING_SNAKE_CASE  :						Optional[Any]      					= GPTaTokenizerFast.from_pretrained('''gpt2'''      )

					# Under the mock environment we get a 500 error when trying to reach the tokenizer.
					with mock.patch('''requests.Session.request''',					return_value=lowerCamelCase__      ) as mock_head:
									SCREAMING_SNAKE_CASE  :						str      					= GPTaTokenizerFast.from_pretrained('''gpt2'''      )
									# This check we did call the fake head request
									mock_head.assert_called()


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

					"""simple docstring"""

					try:
									SCREAMING_SNAKE_CASE  :						Dict      					= tempfile.mktemp()
									with open(lowerCamelCase__,					'''wb'''      ) as f:
													http_get('''https://huggingface.co/albert-base-v1/resolve/main/spiece.model''',					lowerCamelCase__      )

									SCREAMING_SNAKE_CASE  :						Dict      					= AlbertTokenizer.from_pretrained(lowerCamelCase__      )
					finally:
									os.remove(lowerCamelCase__      )

					# Supporting this legacy load introduced a weird bug where the tokenizer would load local files if they are in
					# the current folder and have the right name.
					if os.path.isfile('''tokenizer.json'''      ):
									# We skip the test if the user has a `tokenizer.json` in this folder to avoid deleting it.
									return
					try:
									with open('''tokenizer.json''',					'''wb'''      ) as f:
													http_get('''https://huggingface.co/hf-internal-testing/tiny-random-bert/blob/main/tokenizer.json''',					lowerCamelCase__      )
									SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''      )
									# The tiny random BERT has a vocab size of 1024, tiny gpt2 as a vocab size of 1000
									self.assertEqual(tokenizer.vocab_size,					1000      )
									# Tokenizer should depend on the remote checkpoint, not the local tokenizer.json file.

					finally:
									os.remove('''tokenizer.json'''      )






	def 				lowercase__  (self : str      )     ->			int:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						Any      					= AlbertTokenizer.from_pretrained('''https://huggingface.co/albert-base-v1/resolve/main/spiece.model'''      )



@is_staging_test
class 	a__     (				unittest.TestCase	):
	__magic_name__		:    List[Any]	    =     ["[UNK]", "[CLS]", "[SEP]", "[PAD]", "[MASK]", "bla", "blou"]


	@classmethod
	def 				lowercase__  (cls : List[str]      )     ->			int:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						List[str]      					= TOKEN
					HfFolder.save_token(lowerCamelCase__      )


	@classmethod
	def 				lowercase__  (cls : Any      )     ->			Optional[Any]:

					"""simple docstring"""

					try:
									delete_repo(token=cls._token,					repo_id='''test-tokenizer'''      )
					except HTTPError:
									pass

					try:
									delete_repo(token=cls._token,					repo_id='''valid_org/test-tokenizer-org'''      )
					except HTTPError:
									pass

					try:
									delete_repo(token=cls._token,					repo_id='''test-dynamic-tokenizer'''      )
					except HTTPError:
									pass


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

					"""simple docstring"""

					with tempfile.TemporaryDirectory() as tmp_dir:
									SCREAMING_SNAKE_CASE  :						int      					= os.path.join(lowerCamelCase__,					'''vocab.txt'''      )
									with open(lowerCamelCase__,					'''w''',					encoding='''utf-8'''      ) as vocab_writer:
													vocab_writer.write(''''''.join([x + '''\n''' for x in self.vocab_tokens]      )      )
									SCREAMING_SNAKE_CASE  :						Tuple      					= BertTokenizer(lowerCamelCase__      )

					tokenizer.push_to_hub('''test-tokenizer''',					use_auth_token=self._token      )
					SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= BertTokenizer.from_pretrained(F'''{USER}/test-tokenizer'''      )
					self.assertDictEqual(new_tokenizer.vocab,					tokenizer.vocab      )

					# Reset repo
					delete_repo(token=self._token,					repo_id='''test-tokenizer'''      )

					# Push to hub via save_pretrained
					with tempfile.TemporaryDirectory() as tmp_dir:
									tokenizer.save_pretrained(lowerCamelCase__,					repo_id='''test-tokenizer''',					push_to_hub=lowerCamelCase__,					use_auth_token=self._token      )

					SCREAMING_SNAKE_CASE  :						str      					= BertTokenizer.from_pretrained(F'''{USER}/test-tokenizer'''      )
					self.assertDictEqual(new_tokenizer.vocab,					tokenizer.vocab      )


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

					"""simple docstring"""

					with tempfile.TemporaryDirectory() as tmp_dir:
									SCREAMING_SNAKE_CASE  :						Dict      					= os.path.join(lowerCamelCase__,					'''vocab.txt'''      )
									with open(lowerCamelCase__,					'''w''',					encoding='''utf-8'''      ) as vocab_writer:
													vocab_writer.write(''''''.join([x + '''\n''' for x in self.vocab_tokens]      )      )
									SCREAMING_SNAKE_CASE  :						List[str]      					= BertTokenizer(lowerCamelCase__      )

					tokenizer.push_to_hub('''valid_org/test-tokenizer-org''',					use_auth_token=self._token      )
					SCREAMING_SNAKE_CASE  :						Optional[Any]      					= BertTokenizer.from_pretrained('''valid_org/test-tokenizer-org'''      )
					self.assertDictEqual(new_tokenizer.vocab,					tokenizer.vocab      )

					# Reset repo
					delete_repo(token=self._token,					repo_id='''valid_org/test-tokenizer-org'''      )

					# Push to hub via save_pretrained
					with tempfile.TemporaryDirectory() as tmp_dir:
									tokenizer.save_pretrained(
									    lowerCamelCase__,					repo_id='''valid_org/test-tokenizer-org''',					push_to_hub=lowerCamelCase__,					use_auth_token=self._token      )

					SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= BertTokenizer.from_pretrained('''valid_org/test-tokenizer-org'''      )
					self.assertDictEqual(new_tokenizer.vocab,					tokenizer.vocab      )






	@require_tokenizers
	def 				lowercase__  (self : List[str]      )     ->			int:

					"""simple docstring"""

					CustomTokenizer.register_for_auto_class()
					with tempfile.TemporaryDirectory() as tmp_dir:
									SCREAMING_SNAKE_CASE  :						List[Any]      					= os.path.join(lowerCamelCase__,					'''vocab.txt'''      )
									with open(lowerCamelCase__,					'''w''',					encoding='''utf-8'''      ) as vocab_writer:
													vocab_writer.write(''''''.join([x + '''\n''' for x in self.vocab_tokens]      )      )
									SCREAMING_SNAKE_CASE  :						Any      					= CustomTokenizer(lowerCamelCase__      )

					# No fast custom tokenizer
					tokenizer.push_to_hub('''test-dynamic-tokenizer''',					use_auth_token=self._token      )

					SCREAMING_SNAKE_CASE  :						str      					= AutoTokenizer.from_pretrained(F'''{USER}/test-dynamic-tokenizer''',					trust_remote_code=lowerCamelCase__      )
					# Can't make an isinstance check because the new_model.config is from the CustomTokenizer class of a dynamic module
					self.assertEqual(tokenizer.__class__.__name__,					'''CustomTokenizer'''      )

					# Fast and slow custom tokenizer
					CustomTokenizerFast.register_for_auto_class()
					with tempfile.TemporaryDirectory() as tmp_dir:
									SCREAMING_SNAKE_CASE  :						List[Any]      					= os.path.join(lowerCamelCase__,					'''vocab.txt'''      )
									with open(lowerCamelCase__,					'''w''',					encoding='''utf-8'''      ) as vocab_writer:
													vocab_writer.write(''''''.join([x + '''\n''' for x in self.vocab_tokens]      )      )

									SCREAMING_SNAKE_CASE  :						List[Any]      					= BertTokenizerFast.from_pretrained(lowerCamelCase__      )
									bert_tokenizer.save_pretrained(lowerCamelCase__      )
									SCREAMING_SNAKE_CASE  :						Optional[int]      					= CustomTokenizerFast.from_pretrained(lowerCamelCase__      )

					tokenizer.push_to_hub('''test-dynamic-tokenizer''',					use_auth_token=self._token      )

					SCREAMING_SNAKE_CASE  :						str      					= AutoTokenizer.from_pretrained(F'''{USER}/test-dynamic-tokenizer''',					trust_remote_code=lowerCamelCase__      )
					# Can't make an isinstance check because the new_model.config is from the FakeConfig class of a dynamic module
					self.assertEqual(tokenizer.__class__.__name__,					'''CustomTokenizerFast'''      )
					SCREAMING_SNAKE_CASE  :						int      					= AutoTokenizer.from_pretrained(
					    F'''{USER}/test-dynamic-tokenizer''',					use_fast=lowerCamelCase__,					trust_remote_code=lowerCamelCase__      )
					# Can't make an isinstance check because the new_model.config is from the FakeConfig class of a dynamic module
					self.assertEqual(tokenizer.__class__.__name__,					'''CustomTokenizer'''      )



class 	a__     (				unittest.TestCase	):


	def 				lowercase__  (self : str      )     ->			Dict:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						Any      					= Trie()
					trie.add('''Hello 友達'''      )
					self.assertEqual(trie.data,					{'''H''': {'''e''': {'''l''': {'''l''': {'''o''': {''' ''': {'''友''': {'''達''': {'''''': 1}}}}}}}}}      )
					trie.add('''Hello'''      )
					trie.data
					self.assertEqual(trie.data,					{'''H''': {'''e''': {'''l''': {'''l''': {'''o''': {'''''': 1, ''' ''': {'''友''': {'''達''': {'''''': 1}}}}}}}}}      )


	def 				lowercase__  (self : str      )     ->			str:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						Dict      					= Trie()
					self.assertEqual(trie.split('''[CLS] This is a extra_id_100'''      ),					['''[CLS] This is a extra_id_100''']      )
					trie.add('''[CLS]'''      )
					trie.add('''extra_id_1'''      )
					trie.add('''extra_id_100'''      )
					self.assertEqual(trie.split('''[CLS] This is a extra_id_100'''      ),					['''[CLS]''', ''' This is a ''', '''extra_id_100''']      )


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

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						int      					= Trie()
					trie.add('''A'''      )
					self.assertEqual(trie.split('''ABC'''      ),					['''A''', '''BC''']      )
					self.assertEqual(trie.split('''BCA'''      ),					['''BC''', '''A''']      )


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

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						Any      					= Trie()
					trie.add('''TOKEN]'''      )
					trie.add('''[SPECIAL_TOKEN]'''      )
					self.assertEqual(trie.split('''This is something [SPECIAL_TOKEN]'''      ),					['''This is something ''', '''[SPECIAL_TOKEN]''']      )


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

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						List[Any]      					= Trie()
					trie.add('''A'''      )
					trie.add('''P'''      )
					trie.add('''[SPECIAL_TOKEN]'''      )
					self.assertEqual(trie.split('''This is something [SPECIAL_TOKEN]'''      ),					['''This is something ''', '''[SPECIAL_TOKEN]''']      )


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

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						List[Any]      					= Trie()
					trie.add('''AB'''      )
					trie.add('''B'''      )
					trie.add('''C'''      )
					self.assertEqual(trie.split('''ABC'''      ),					['''AB''', '''C''']      )


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

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						int      					= Trie()
					trie.add('''ABC'''      )
					trie.add('''B'''      )
					trie.add('''CD'''      )
					self.assertEqual(trie.split('''ABCD'''      ),					['''ABC''', '''D''']      )






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

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						Dict      					= Trie()
					SCREAMING_SNAKE_CASE  :						Optional[Any]      					= trie.cut_text('''ABC''',					[0, 0, 2, 1, 2, 3]      )
					self.assertEqual(lowerCamelCase__,					['''AB''', '''C''']      )









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




