





"""simple docstring"""





import inspect

import jax
import jax.lax as lax
import jax.numpy as jnp

from ..utils import add_start_docstrings
from ..utils.logging import get_logger


_UpperCamelCase							=      get_logger(__name__)


_UpperCamelCase							=      r'\n    Args:\n        input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):\n            Indices of input sequence tokens in the vocabulary.\n\n            Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and\n            [`PreTrainedTokenizer.__call__`] for details.\n\n            [What are input IDs?](../glossary#input-ids)\n        scores (`jnp.ndarray` of shape `(batch_size, config.vocab_size)`):\n            Prediction scores of a language modeling head. These can be logits for each vocabulary when not using beam\n            search or log softmax for each vocabulary token when using beam search\n        kwargs (`Dict[str, Any]`, *optional*):\n            Additional logits processor specific kwargs.\n\n    Return:\n        `jnp.ndarray` of shape `(batch_size, config.vocab_size)`: The processed prediction scores.\n\n'
class 		SCREAMING_SNAKE_CASE_							:
				"""simple docstring"""




				@add_start_docstrings(__lowercase			)
				def __call__(     self  :Union[str, Any]						,    __lowercase  :jnp.ndarray						,    __lowercase  :jnp.ndarray			):
											raise NotImplementedError(
											    f'{self.__class__} is an abstract class. Only classes inheriting this class can be called.'			)

class 		SCREAMING_SNAKE_CASE_							:
				"""simple docstring"""




				@add_start_docstrings(__lowercase			)
				def __call__(     self  :Optional[int]						,    __lowercase  :jnp.ndarray						,    __lowercase  :jnp.ndarray			):
											raise NotImplementedError(
											    f'{self.__class__} is an abstract class. Only classes inheriting this class can be called.'			)

class 		SCREAMING_SNAKE_CASE_							(    snake_case__   ):
				"""simple docstring"""




				@add_start_docstrings(__lowercase			)
				def __call__(     self  :Union[str, Any]						,    __lowercase  :jnp.ndarray						,    __lowercase  :jnp.ndarray						,    __lowercase  :int						,    **__lowercase  :Optional[Any]			):
											for processor in self:
																		__lowerCamelCase      :      Union[str, Any]         =inspect.signature(processor.__call__			).parameters
																		if len(__lowercase			) > 3:
																									if not all(arg in kwargs for arg in list(function_args.keys()			)[2:]			):
																																raise ValueError(
																																    f'Make sure that all the required parameters: {list(function_args.keys()			)} for '
																																    f'{processor.__class__} are passed to the logits processor.'			)
																									__lowerCamelCase      :      Optional[Any]         =processor(__lowercase						,    __lowercase						,    __lowercase						,    **__lowercase			)
																		else:
																									__lowerCamelCase      :      List[Any]         =processor(__lowercase						,    __lowercase						,    __lowercase			)
											return scores

class 		SCREAMING_SNAKE_CASE_							(    snake_case__   ):
				"""simple docstring"""




				def __init__(     self  :Tuple						,    __lowercase  :float			):
											if not isinstance(__lowercase						,    __lowercase			) or not (temperature > 0):
																		raise ValueError(f'`temperature` has to be a strictly positive float, but is {temperature}'			)

											__lowerCamelCase      :      Any         =temperature




				def __call__(     self  :Union[str, Any]						,    __lowercase  :jnp.ndarray						,    __lowercase  :jnp.ndarray						,    __lowercase  :int			):
											__lowerCamelCase      :      Union[str, Any]         =scores / self.temperature
											return scores

class 		SCREAMING_SNAKE_CASE_							(    snake_case__   ):
				"""simple docstring"""




				def __init__(     self  :Tuple						,    __lowercase  :float						,    __lowercase  :float = -float('''Inf'''			)						,    __lowercase  :int = 1			):
											if not isinstance(__lowercase						,    __lowercase			) or (top_p < 0 or top_p > 1.0):
																		raise ValueError(f'`top_p` has to be a float > 0 and < 1, but is {top_p}'			)
											if not isinstance(__lowercase						,    __lowercase			) or (min_tokens_to_keep < 1):
																		raise ValueError(f'`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}'			)

											__lowerCamelCase      :      Union[str, Any]         =top_p
											__lowerCamelCase      :      Any         =filter_value
											__lowerCamelCase      :      str         =min_tokens_to_keep




				def __call__(     self  :List[Any]						,    __lowercase  :jnp.ndarray						,    __lowercase  :jnp.ndarray						,    __lowercase  :int			):
											__lowerCamelCase    ,						__lowerCamelCase      :      Union[str, Any]         =lax.top_k(__lowercase						,    scores.shape[-1]			)

											__lowerCamelCase      :      Tuple         =jnp.full_like(__lowercase						,    self.filter_value			)
											__lowerCamelCase      :      List[str]         =jax.nn.softmax(__lowercase						,    axis=-1			).cumsum(axis=-1			)
											__lowerCamelCase      :      Optional[int]         =cumulative_probs < self.top_p

											# include the token that is higher than top_p as well
											__lowerCamelCase      :      Optional[Any]         =jnp.roll(__lowercase						,    1			)
											score_mask |= score_mask.at[:, 0].set(__lowercase			)

											# min tokens to keep
											__lowerCamelCase      :      Tuple         =score_mask.at[:, : self.min_tokens_to_keep].set(__lowercase			)

											__lowerCamelCase      :      int         =jnp.where(__lowercase						,    __lowercase						,    __lowercase			)
											__lowerCamelCase      :      Optional[Any]         =jax.lax.sort_key_val(__lowercase						,    __lowercase			)[-1]

											return next_scores

class 		SCREAMING_SNAKE_CASE_							(    snake_case__   ):
				"""simple docstring"""




				def __init__(     self  :List[str]						,    __lowercase  :int						,    __lowercase  :float = -float('''Inf'''			)						,    __lowercase  :int = 1			):
											if not isinstance(__lowercase						,    __lowercase			) or top_k <= 0:
																		raise ValueError(f'`top_k` has to be a strictly positive integer, but is {top_k}'			)

											__lowerCamelCase      :      Optional[int]         =max(__lowercase						,    __lowercase			)
											__lowerCamelCase      :      int         =filter_value




				def __call__(     self  :int						,    __lowercase  :jnp.ndarray						,    __lowercase  :jnp.ndarray						,    __lowercase  :int			):
											__lowerCamelCase    ,						__lowerCamelCase      :      List[Any]         =scores.shape
											__lowerCamelCase      :      Optional[int]         =jnp.full(batch_size * vocab_size						,    self.filter_value			)

											__lowerCamelCase      :      Optional[Any]         =min(self.top_k						,    scores.shape[-1]			)  # Safety check
											__lowerCamelCase    ,						__lowerCamelCase      :      Optional[int]         =lax.top_k(__lowercase						,    __lowercase			)
											__lowerCamelCase      :      List[str]         =jnp.broadcast_to((jnp.arange(__lowercase			) * vocab_size)[:, None]						,    (batch_size, topk)			).flatten()
											__lowerCamelCase      :      Optional[Any]         =topk_scores.flatten()
											__lowerCamelCase      :      Optional[int]         =topk_indices.flatten() + shift

											__lowerCamelCase      :      Tuple         =next_scores_flat.at[topk_indices_flat].set(__lowercase			)
											__lowerCamelCase      :      Optional[Any]         =next_scores_flat.reshape(__lowercase						,    __lowercase			)
											return next_scores

class 		SCREAMING_SNAKE_CASE_							(    snake_case__   ):
				"""simple docstring"""




				def __init__(     self  :List[str]						,    __lowercase  :int			):
											__lowerCamelCase      :      int         =bos_token_id




				def __call__(     self  :Union[str, Any]						,    __lowercase  :jnp.ndarray						,    __lowercase  :jnp.ndarray						,    __lowercase  :int			):
											__lowerCamelCase      :      List[str]         =jnp.full(scores.shape						,    -float('''inf'''			)			)

											__lowerCamelCase      :      Union[str, Any]         =1 - jnp.bool_(cur_len - 1			)

											__lowerCamelCase      :      Tuple         =jnp.where(__lowercase						,    new_scores.at[:, self.bos_token_id].set(0			)						,    __lowercase			)

											return scores

class 		SCREAMING_SNAKE_CASE_							(    snake_case__   ):
				"""simple docstring"""




				def __init__(     self  :str						,    __lowercase  :int						,    __lowercase  :int			):
											__lowerCamelCase      :      List[Any]         =max_length
											__lowerCamelCase      :      List[Any]         =eos_token_id




				def __call__(     self  :str						,    __lowercase  :jnp.ndarray						,    __lowercase  :jnp.ndarray						,    __lowercase  :int			):
											__lowerCamelCase      :      str         =jnp.full(scores.shape						,    -float('''inf'''			)			)

											__lowerCamelCase      :      Tuple         =1 - jnp.bool_(cur_len - self.max_length + 1			)

											__lowerCamelCase      :      Optional[int]         =jnp.where(__lowercase						,    new_scores.at[:, self.eos_token_id].set(0			)						,    __lowercase			)

											return scores

class 		SCREAMING_SNAKE_CASE_							(    snake_case__   ):
				"""simple docstring"""




				def __init__(     self  :int						,    __lowercase  :int						,    __lowercase  :int			):
											if not isinstance(__lowercase						,    __lowercase			) or min_length < 0:
																		raise ValueError(f'`min_length` has to be a positive integer, but is {min_length}'			)

											if not isinstance(__lowercase						,    __lowercase			) or eos_token_id < 0:
																		raise ValueError(f'`eos_token_id` has to be a positive integer, but is {eos_token_id}'			)

											__lowerCamelCase      :      List[Any]         =min_length
											__lowerCamelCase      :      Optional[int]         =eos_token_id




				def __call__(     self  :Optional[Any]						,    __lowercase  :jnp.ndarray						,    __lowercase  :jnp.ndarray						,    __lowercase  :int			):
											# create boolean flag to decide if min length penalty should be applied
											__lowerCamelCase      :      Dict         =1 - jnp.clip(cur_len - self.min_length						,    0						,    1			)

											__lowerCamelCase      :      Optional[int]         =jnp.where(__lowercase						,    scores.at[:, self.eos_token_id].set(-float('''inf'''			)			)						,    __lowercase			)

											return scores

class 		SCREAMING_SNAKE_CASE_							(    snake_case__   ):
				"""simple docstring"""




				def __init__(     self  :Dict						,    __lowercase  :List[str]						,    __lowercase  :Dict			):
											__lowerCamelCase      :      Optional[Any]         =list(__lowercase			)
											__lowerCamelCase      :      Any         =begin_index




				def __call__(     self  :Optional[int]						,    __lowercase  :int						,    __lowercase  :List[Any]						,    __lowercase  :int			):
											__lowerCamelCase      :      int         =1 - jnp.bool_(cur_len - self.begin_index			)

											__lowerCamelCase      :      Optional[int]         =jnp.where(__lowercase						,    scores.at[:, self.begin_suppress_tokens].set(-float('''inf'''			)			)						,    __lowercase			)

											return scores

class 		SCREAMING_SNAKE_CASE_							(    snake_case__   ):
				"""simple docstring"""




				def __init__(     self  :Any						,    __lowercase  :list			):
											__lowerCamelCase      :      List[str]         =list(__lowercase			)




				def __call__(     self  :Dict						,    __lowercase  :jnp.ndarray						,    __lowercase  :jnp.ndarray						,    __lowercase  :int			):
											__lowerCamelCase      :      int         =scores.at[..., self.suppress_tokens].set(-float('''inf'''			)			)

											return scores

class 		SCREAMING_SNAKE_CASE_							(    snake_case__   ):
				"""simple docstring"""




				def __init__(     self  :int						,    __lowercase  :Optional[Any]			):
											__lowerCamelCase      :      Tuple         =dict(__lowercase			)
											# Converts the dictionary of format {index: token} containing the tokens to be forced to an array, where the
											# index of the array corresponds to the index of the token to be forced, for XLA compatibility.
											# Indexes without forced tokens will have a negative value.
											__lowerCamelCase      :      List[str]         =jnp.ones((max(force_token_map.keys()			) + 1)						,    dtype=jnp.intaa			) * -1
											for index, token in force_token_map.items():
																		if token is not None:
																									__lowerCamelCase      :      str         =force_token_array.at[index].set(__lowercase			)
											__lowerCamelCase      :      Tuple         =jnp.intaa(__lowercase			)




				def __call__(     self  :Tuple						,    __lowercase  :jnp.ndarray						,    __lowercase  :jnp.ndarray						,    __lowercase  :int			):
											def _force_token(__lowercase  :Any			):
																		__lowerCamelCase      :      Any         =scores.shape[0]
																		__lowerCamelCase      :      List[str]         =self.force_token_array[generation_idx]

																		__lowerCamelCase      :      Optional[Any]         =jnp.ones_like(__lowercase						,    dtype=scores.dtype			) * -float('''inf'''			)
																		__lowerCamelCase      :      str         =jnp.zeros((batch_size, 1)						,    dtype=scores.dtype			)
																		__lowerCamelCase      :      Optional[int]         =lax.dynamic_update_slice(__lowercase						,    __lowercase						,    (0, current_token)			)
																		return new_scores

											__lowerCamelCase      :      Optional[Any]         =lax.cond(
											    cur_len >= self.force_token_array.shape[0]						,    lambda: scores						,    lambda: lax.cond(
											        self.force_token_array[cur_len] >= 0						,    lambda: _force_token(__lowercase			)						,    lambda: scores						,    )						,    )
											return scores

class 		SCREAMING_SNAKE_CASE_							(    snake_case__   ):
				"""simple docstring"""




				def __init__(     self  :List[str]						,    __lowercase  :str						,    __lowercase  :Tuple						,    __lowercase  :Dict			):
											__lowerCamelCase      :      List[Any]         =generate_config.eos_token_id
											__lowerCamelCase      :      str         =generate_config.no_timestamps_token_id
											__lowerCamelCase      :      List[str]         =generate_config.no_timestamps_token_id + 1

											__lowerCamelCase      :      Optional[Any]         =decoder_input_length + 1

											if generate_config.is_multilingual:
																		# room for language token and task token
																		self.begin_index += 2
											if hasattr(__lowercase						,    '''max_initial_timestamp_index'''			):
																		__lowerCamelCase      :      str         =generate_config.max_initial_timestamp_index
											else:
																		__lowerCamelCase      :      List[Any]         =model_config.vocab_size
											if self.max_initial_timestamp_index is None:
																		__lowerCamelCase      :      List[str]         =model_config.vocab_size




				def __call__(     self  :Tuple						,    __lowercase  :Dict						,    __lowercase  :int						,    __lowercase  :Optional[Any]			):
											# suppress <|notimestamps|> which is handled by without_timestamps
											__lowerCamelCase      :      Dict         =scores.at[:, self.no_timestamps_token_id].set(-float('''inf'''			)			)

											def handle_pairs(__lowercase  :List[Any]						,    __lowercase  :Dict			):
																		__lowerCamelCase      :      int         =jnp.where((cur_len - self.begin_index) >= 1						,    __lowercase						,    __lowercase			)
																		__lowerCamelCase      :      List[Any]         =jnp.where(
																		    input_ids_k[cur_len - 1] >= self.timestamp_begin						,    True and last_was_timestamp						,    __lowercase						,    )

																		__lowerCamelCase      :      Any         =jnp.where((cur_len - self.begin_index) < 2						,    __lowercase						,    __lowercase			)
																		__lowerCamelCase      :      str         =jnp.where(
																		    input_ids_k[cur_len - 2] >= self.timestamp_begin						,    __lowercase						,    __lowercase						,    )

																		return jnp.where(
																		    __lowercase						,    jnp.where(
																		        penultimate_was_timestamp > 0						,    scores_k.at[self.timestamp_begin :].set(-float('''inf'''			)			)						,    scores_k.at[: self.eos_token_id].set(-float('''inf'''			)			)						,    )						,    __lowercase						,    )

											__lowerCamelCase      :      int         =jax.vmap(__lowercase			)(__lowercase						,    __lowercase			)

											__lowerCamelCase      :      Union[str, Any]         =jnp.where(cur_len == self.begin_index						,    __lowercase						,    __lowercase			)
											__lowerCamelCase      :      Dict         =jnp.where(
											    self.max_initial_timestamp_index is not None						,    True and apply_max_initial_timestamp						,    __lowercase						,    )

											__lowerCamelCase      :      Dict         =self.timestamp_begin + self.max_initial_timestamp_index

											__lowerCamelCase      :      List[str]         =jnp.where(
											    __lowercase						,    scores.at[:, last_allowed + 1 :].set(-float('''inf'''			)			)						,    __lowercase						,    )

											# if sum of probability over timestamps is above any other token, sample timestamp
											__lowerCamelCase      :      Tuple         =jax.nn.log_softmax(__lowercase						,    axis=-1			)

											def handle_cumulative_probs(__lowercase  :int						,    __lowercase  :Tuple			):
																		__lowerCamelCase      :      int         =jax.nn.logsumexp(logprobs_k[self.timestamp_begin :]						,    axis=-1			)
																		__lowerCamelCase      :      List[Any]         =jnp.max(logprobs_k[: self.timestamp_begin]			)
																		return jnp.where(
																		    timestamp_logprob > max_text_token_logprob						,    scores_k.at[: self.timestamp_begin].set(-float('''inf'''			)			)						,    __lowercase						,    )

											__lowerCamelCase      :      Any         =jax.vmap(__lowercase			)(__lowercase						,    __lowercase			)

											return scores











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




