





"""simple docstring"""





import unittest

from transformers import TrOCRConfig
from transformers.testing_utils import is_torch_available, require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch

							from transformers.models.trocr.modeling_trocr import TrOCRDecoder, TrOCRForCausalLM
@require_torch
class 		SCREAMING_SNAKE_CASE_							:
				"""simple docstring"""




				def __init__(     self  :int						,    __lowercase  :Union[str, Any]						,    __lowercase  :Optional[int]=99						,    __lowercase  :Any=13						,    __lowercase  :int=16						,    __lowercase  :Dict=7						,    __lowercase  :List[Any]=True						,    __lowercase  :List[str]=True						,    __lowercase  :List[Any]=True						,    __lowercase  :Optional[int]=False						,    __lowercase  :int=True						,    __lowercase  :Union[str, Any]=2						,    __lowercase  :List[str]=32						,    __lowercase  :Optional[int]=4						,    __lowercase  :Dict=4						,    __lowercase  :Any=30						,    __lowercase  :List[Any]=0						,    __lowercase  :Any=1						,    __lowercase  :Optional[Any]=2						,    __lowercase  :Tuple=None						,    ):
											__lowerCamelCase      :      Tuple         =parent
											__lowerCamelCase      :      List[str]         =batch_size
											__lowerCamelCase      :      List[str]         =decoder_seq_length
											# For common tests
											__lowerCamelCase      :      int         =self.decoder_seq_length
											__lowerCamelCase      :      str         =is_training
											__lowerCamelCase      :      int         =use_attention_mask
											__lowerCamelCase      :      Any         =use_labels

											__lowerCamelCase      :      Union[str, Any]         =vocab_size
											__lowerCamelCase      :      Any         =d_model
											__lowerCamelCase      :      List[str]         =d_model
											__lowerCamelCase      :      List[str]         =decoder_layers
											__lowerCamelCase      :      Optional[int]         =decoder_layers
											__lowerCamelCase      :      Dict         =decoder_ffn_dim
											__lowerCamelCase      :      Optional[Any]         =decoder_attention_heads
											__lowerCamelCase      :      List[str]         =decoder_attention_heads
											__lowerCamelCase      :      List[Any]         =eos_token_id
											__lowerCamelCase      :      List[str]         =bos_token_id
											__lowerCamelCase      :      Union[str, Any]         =pad_token_id
											__lowerCamelCase      :      Union[str, Any]         =decoder_start_token_id
											__lowerCamelCase      :      Any         =use_cache
											__lowerCamelCase      :      Dict         =max_position_embeddings

											__lowerCamelCase      :      Tuple         =None
											__lowerCamelCase      :      List[Any]         =decoder_seq_length
											__lowerCamelCase      :      int         =2
											__lowerCamelCase      :      Optional[Any]         =1




				def 					__lowercase  (     self  :str			):
											__lowerCamelCase      :      str         =ids_tensor([self.batch_size, self.decoder_seq_length]						,    self.vocab_size			)

											__lowerCamelCase      :      Dict         =None
											if self.use_attention_mask:
																		__lowerCamelCase      :      Optional[Any]         =ids_tensor([self.batch_size, self.decoder_seq_length]						,    vocab_size=2			)

											__lowerCamelCase      :      Any         =None
											if self.use_labels:
																		__lowerCamelCase      :      str         =ids_tensor([self.batch_size, self.decoder_seq_length]						,    self.vocab_size			)

											__lowerCamelCase      :      Tuple         =TrOCRConfig(
											    vocab_size=self.vocab_size						,    d_model=self.d_model						,    decoder_layers=self.decoder_layers						,    decoder_ffn_dim=self.decoder_ffn_dim						,    decoder_attention_heads=self.decoder_attention_heads						,    eos_token_id=self.eos_token_id						,    bos_token_id=self.bos_token_id						,    use_cache=self.use_cache						,    pad_token_id=self.pad_token_id						,    decoder_start_token_id=self.decoder_start_token_id						,    max_position_embeddings=self.max_position_embeddings						,    )

											return (config, input_ids, attention_mask, lm_labels)




				def 					__lowercase  (     self  :Optional[int]						,    __lowercase  :int						,    __lowercase  :Any						,    __lowercase  :int						,    __lowercase  :int						,    ):
											__lowerCamelCase      :      List[str]         =True
											__lowerCamelCase      :      Dict         =TrOCRDecoder(config=__lowercase			).to(__lowercase			).eval()
											__lowerCamelCase      :      List[str]         =input_ids[:2]

											input_ids[input_ids == 0] += 1
											# first forward pass
											__lowerCamelCase      :      Optional[int]         =model(__lowercase						,    use_cache=__lowercase			)
											__lowerCamelCase      :      Union[str, Any]         =model(__lowercase			)
											__lowerCamelCase      :      str         =model(__lowercase						,    use_cache=__lowercase			)

											self.parent.assertTrue(len(__lowercase			) == len(__lowercase			)			)
											self.parent.assertTrue(len(__lowercase			) == len(__lowercase			) + 1			)

											__lowerCamelCase      :      Tuple         =outputs['''past_key_values''']

											# create hypothetical next token and extent to next_input_ids
											__lowerCamelCase      :      Optional[int]         =ids_tensor((2, 1)						,    config.vocab_size - 1			) + 1

											# append to next input_ids and
											__lowerCamelCase      :      Dict         =torch.cat([input_ids, next_tokens]						,    dim=-1			)

											__lowerCamelCase      :      str         =model(__lowercase			)['''last_hidden_state''']
											__lowerCamelCase      :      Optional[Any]         =model(__lowercase						,    past_key_values=__lowercase			)['''last_hidden_state''']

											# select random slice
											__lowerCamelCase      :      int         =ids_tensor((1,)						,    output_from_past.shape[-1]			).item()
											__lowerCamelCase      :      str         =output_from_no_past[:, next_input_ids.shape[-1] - 1, random_slice_idx].detach()
											__lowerCamelCase      :      Dict         =output_from_past[:, 0, random_slice_idx].detach()

											# test that outputs are equal for slice
											assert torch.allclose(__lowercase						,    __lowercase						,    atol=1e-3			)




				def 					__lowercase  (     self  :Dict			):
											__lowerCamelCase      :      Tuple         =self.prepare_config_and_inputs()
											__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase      :      str         =config_and_inputs

											__lowerCamelCase      :      str         ={'''input_ids''': input_ids, '''attention_mask''': attention_mask}
											return config, inputs_dict

@require_torch
class 		SCREAMING_SNAKE_CASE_							(    snake_case__       ,  snake_case__       ,  snake_case__       ,  unittest.TestCase   ):
				"""simple docstring"""
				__snake_case   :       Tuple	 =		(TrOCRDecoder, TrOCRForCausalLM) if is_torch_available() else ()
				__snake_case   :       Optional[Any]	 =		(TrOCRForCausalLM,) if is_torch_available() else ()
				__snake_case   :       Optional[int]	 =		{"""text-generation""": TrOCRForCausalLM} if is_torch_available() else {}
				__snake_case   :       Any	 =		True
				__snake_case   :       List[str]	 =		False




				def 					__lowercase  (     self  :Any			):
											__lowerCamelCase      :      str         =TrOCRStandaloneDecoderModelTester(self						,    is_training=__lowercase			)
											__lowerCamelCase      :      Dict         =ConfigTester(self						,    config_class=__lowercase			)




				def 					__lowercase  (     self  :str			):
											pass




				def 					__lowercase  (     self  :List[Any]			):
											pass




				def 					__lowercase  (     self  :Optional[Any]			):
											pass




				def 					__lowercase  (     self  :str			):
											self.config_tester.run_common_tests()




				def 					__lowercase  (     self  :Dict			):
											__lowerCamelCase      :      List[Any]         =self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_decoder_model_past(*__lowercase			)




				def 					__lowercase  (     self  :Union[str, Any]			):
											return




				@unittest.skip('''The model doesn\'t support left padding'''			)  # and it's not used enough to be worth fixing :)
				def 					__lowercase  (     self  :List[str]			):
											pass











"""simple docstring"""





import argparse
from pathlib import Path

import requests
import torch
from PIL import Image

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


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



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


							'''simple docstring'''







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

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

							return rename_keys



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


							'''simple docstring'''







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

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



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


							'''simple docstring'''







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



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							List[Any]   ):


							'''simple docstring'''







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



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


							'''simple docstring'''







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

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

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

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

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

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

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

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

    # load state dict
							model.load_state_dict(SCREAMING_SNAKE_CASE   )

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

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

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

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

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

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


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

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




