






from math import sqrt

def 	__lowerCamelCase      (		snake_case__		)							->  bool:






						"""simple docstring"""





						assert isinstance(snake_case__	,snake_case__		) and (
						    number >= 0
						), "'number' must been an int and positive"

						_SCREAMING_SNAKE_CASE     				=  True

						# 0 and 1 are none primes.
						if number <= 1:
												_SCREAMING_SNAKE_CASE     				=  False

						for divisor in range(2	,int(round(sqrt(snake_case__		)		)		) + 1		):
												# if 'number' divisible by 'divisor' then sets 'status'
												# of false and break up the loop.
												if number % divisor == 0:
																		_SCREAMING_SNAKE_CASE     				=  False
																		break

    # precondition
						assert isinstance(snake_case__	,snake_case__		), "'status' must been from type bool"

						return status

def 	__lowerCamelCase      (		snake_case__		)							->  Optional[Any]:






						"""simple docstring"""





						assert isinstance(snake_case__	,snake_case__		) and (n > 2), "'N' must been an int and > 2"

						# beginList: contains all natural numbers from 2 up to N
						_SCREAMING_SNAKE_CASE     				=  list(range(2	,n + 1		)		)

						_SCREAMING_SNAKE_CASE     				=  []  # this list will be returns.

						# actual sieve of erathostenes
						for i in range(len(snake_case__		)		):
												for j in range(i + 1	,len(snake_case__		)		):
																		if (begin_list[i] != 0) and (begin_list[j] % begin_list[i] == 0):
																								_SCREAMING_SNAKE_CASE     				=  0

    # filters actual prime numbers.
						_SCREAMING_SNAKE_CASE     				=  [x for x in begin_list if x != 0]

						# precondition
						assert isinstance(snake_case__	,snake_case__		), "'ans' must been from type list"

						return ans

def 	__lowerCamelCase      (		snake_case__		)							->  Any:






						"""simple docstring"""





						assert isinstance(snake_case__	,snake_case__		) and (n > 2), "'N' must been an int and > 2"

						_SCREAMING_SNAKE_CASE     				=  []

						# iterates over all numbers between 2 up to N+1
						# if a number is prime then appends to list 'ans'
						for number in range(2	,n + 1		):
												if is_prime(snake_case__		):
																		ans.append(snake_case__		)

    # precondition
						assert isinstance(snake_case__	,snake_case__		), "'ans' must been from type list"

						return ans

def 	__lowerCamelCase      (		snake_case__		)							->  Optional[Any]:






						"""simple docstring"""





						assert isinstance(snake_case__	,snake_case__		) and number >= 0, "'number' must been an int and >= 0"

						_SCREAMING_SNAKE_CASE     				=  []  # this list will be returns of the function.

						# potential prime number factors.

						_SCREAMING_SNAKE_CASE     				=  2

						_SCREAMING_SNAKE_CASE     				=  number

						if number == 0 or number == 1:
												ans.append(snake_case__		)

						# if 'number' not prime then builds the prime factorization of 'number'
						elif not is_prime(snake_case__		):
												while quotient != 1:
																		if is_prime(snake_case__		) and (quotient % factor == 0):
																								ans.append(snake_case__		)
																								quotient /= factor
																		else:
																								factor += 1

						else:
												ans.append(snake_case__		)

						# precondition
						assert isinstance(snake_case__	,snake_case__		), "'ans' must been from type list"

						return ans

def 	__lowerCamelCase      (		snake_case__		)							->  Tuple:






						"""simple docstring"""





						assert isinstance(snake_case__	,snake_case__		) and (
						    number >= 0
						), "'number' bust been an int and >= 0"

						_SCREAMING_SNAKE_CASE     				=  0

						# prime factorization of 'number'
						_SCREAMING_SNAKE_CASE     				=  prime_factorization(snake_case__		)

						_SCREAMING_SNAKE_CASE     				=  max(snake_case__		)

						# precondition
						assert isinstance(snake_case__	,snake_case__		), "'ans' must been from type int"

						return ans

def 	__lowerCamelCase      (		snake_case__		)							->  Optional[Any]:






						"""simple docstring"""





						assert isinstance(snake_case__	,snake_case__		) and (
						    number >= 0
						), "'number' bust been an int and >= 0"

						_SCREAMING_SNAKE_CASE     				=  0

						# prime factorization of 'number'
						_SCREAMING_SNAKE_CASE     				=  prime_factorization(snake_case__		)

						_SCREAMING_SNAKE_CASE     				=  min(snake_case__		)

						# precondition
						assert isinstance(snake_case__	,snake_case__		), "'ans' must been from type int"

						return ans

def 	__lowerCamelCase      (		snake_case__		)							->  str:






						"""simple docstring"""





						assert isinstance(snake_case__	,snake_case__		), "'number' must been an int"
						assert isinstance(number % 2 == 0	,snake_case__		), "compare bust been from type bool"

						return number % 2 == 0

def 	__lowerCamelCase      (		snake_case__		)							->  str:






						"""simple docstring"""





						assert isinstance(snake_case__	,snake_case__		), "'number' must been an int"
						assert isinstance(number % 2 != 0	,snake_case__		), "compare bust been from type bool"

						return number % 2 != 0

def 	__lowerCamelCase      (		snake_case__		)							->  str:






						"""simple docstring"""





						assert (
						    isinstance(snake_case__	,snake_case__		) and (number > 2) and is_even(snake_case__		)
						), "'number' must been an int, even and > 2"

						_SCREAMING_SNAKE_CASE     				=  []  # this list will returned

						# creates a list of prime numbers between 2 up to 'number'
						_SCREAMING_SNAKE_CASE     				=  get_prime_numbers(snake_case__		)
						_SCREAMING_SNAKE_CASE     				=  len(snake_case__		)

						# run variable for while-loops.
						_SCREAMING_SNAKE_CASE     				=  0
						_SCREAMING_SNAKE_CASE     				=  None

						# exit variable. for break up the loops
						_SCREAMING_SNAKE_CASE     				=  True

						while i < len_pn and loop:
												_SCREAMING_SNAKE_CASE     				=  i + 1

												while j < len_pn and loop:
																		if prime_numbers[i] + prime_numbers[j] == number:
																								_SCREAMING_SNAKE_CASE     				=  False
																								ans.append(prime_numbers[i]		)
																								ans.append(prime_numbers[j]		)

																		j += 1

												i += 1

						# precondition
						assert (
						    isinstance(snake_case__	,snake_case__		)
						    and (len(snake_case__		) == 2)
						    and (ans[0] + ans[1] == number)
						    and is_prime(ans[0]		)
						    and is_prime(ans[1]		)
						), "'ans' must contains two primes. And sum of elements must been eq 'number'"

						return ans

def 	__lowerCamelCase      (		snake_case__	,snake_case__		)							->  Optional[int]:






						"""simple docstring"""





						assert (
						    isinstance(snake_case__	,snake_case__		)
						    and isinstance(snake_case__	,snake_case__		)
						    and (numbera >= 0)
						    and (numbera >= 0)
						), "'number1' and 'number2' must been positive integer."

						_SCREAMING_SNAKE_CASE     				=  0

						while numbera != 0:
												_SCREAMING_SNAKE_CASE     				=  numbera % numbera
												_SCREAMING_SNAKE_CASE     				=  numbera
												_SCREAMING_SNAKE_CASE     				=  rest

						# precondition
						assert isinstance(snake_case__	,snake_case__		) and (
						    numbera >= 0
						), "'number' must been from type int and positive"

						return numbera

def 	__lowerCamelCase      (		snake_case__	,snake_case__		)							->  Union[str, Any]:






						"""simple docstring"""





						assert (
						    isinstance(snake_case__	,snake_case__		)
						    and isinstance(snake_case__	,snake_case__		)
						    and (numbera >= 1)
						    and (numbera >= 1)
						), "'number1' and 'number2' must been positive integer."

						_SCREAMING_SNAKE_CASE     				=  1  # actual answer that will be return.

						# for kgV (x,1)
						if numbera > 1 and numbera > 1:
												# builds the prime factorization of 'number1' and 'number2'
												_SCREAMING_SNAKE_CASE     				=  prime_factorization(snake_case__		)
												_SCREAMING_SNAKE_CASE     				=  prime_factorization(snake_case__		)

						elif numbera == 1 or numbera == 1:
												_SCREAMING_SNAKE_CASE     				=  []
												_SCREAMING_SNAKE_CASE     				=  []
												_SCREAMING_SNAKE_CASE     				=  max(snake_case__	,snake_case__		)

						_SCREAMING_SNAKE_CASE     				=  0
						_SCREAMING_SNAKE_CASE     				=  0

						_SCREAMING_SNAKE_CASE     				=  []  # captured numbers int both 'primeFac1' and 'primeFac2'

						# iterates through primeFac1
						for n in prime_fac_a:
												if n not in done:
																		if n in prime_fac_a:
																								_SCREAMING_SNAKE_CASE     				=  prime_fac_a.count(snake_case__		)
																								_SCREAMING_SNAKE_CASE     				=  prime_fac_a.count(snake_case__		)

																								for _ in range(max(snake_case__	,snake_case__		)		):
																														ans *= n

																		else:
																								_SCREAMING_SNAKE_CASE     				=  prime_fac_a.count(snake_case__		)

																								for _ in range(snake_case__		):
																														ans *= n

																		done.append(snake_case__		)

    # iterates through primeFac2
						for n in prime_fac_a:
												if n not in done:
																		_SCREAMING_SNAKE_CASE     				=  prime_fac_a.count(snake_case__		)

																		for _ in range(snake_case__		):
																								ans *= n

																		done.append(snake_case__		)

    # precondition
						assert isinstance(snake_case__	,snake_case__		) and (
						    ans >= 0
						), "'ans' must been from type int and positive"

						return ans

def 	__lowerCamelCase      (		snake_case__		)							->  Any:






						"""simple docstring"""





						assert isinstance(snake_case__	,snake_case__		) and (n >= 0), "'number' must been a positive int"

						_SCREAMING_SNAKE_CASE     				=  0
						_SCREAMING_SNAKE_CASE     				=  2  # this variable holds the answer

						while index < n:
												index += 1

												ans += 1  # counts to the next number

												# if ans not prime then
												# runs to the next prime number.
												while not is_prime(snake_case__		):
																		ans += 1

    # precondition
						assert isinstance(snake_case__	,snake_case__		) and is_prime(
						    snake_case__		), "'ans' must been a prime number and from type int"

						return ans

def 	__lowerCamelCase      (		snake_case__	,snake_case__		)							->  Any:






						"""simple docstring"""





						assert (
						    is_prime(snake_case__		) and is_prime(snake_case__		) and (p_number_a < p_number_a)
						), "The arguments must been prime numbers and 'pNumber1' < 'pNumber2'"

						_SCREAMING_SNAKE_CASE     				=  p_number_a + 1  # jump to the next number

						_SCREAMING_SNAKE_CASE     				=  []  # this list will be returns.

						# if number is not prime then
						# fetch the next prime number.
						while not is_prime(snake_case__		):
												number += 1

						while number < p_number_a:
												ans.append(snake_case__		)

												number += 1

												# fetch the next prime number.
												while not is_prime(snake_case__		):
																		number += 1

    # precondition
						assert (
						    isinstance(snake_case__	,snake_case__		)
						    and ans[0] != p_number_a
						    and ans[len(snake_case__		) - 1] != p_number_a
						), "'ans' must been a list without the arguments"

						# 'ans' contains not 'pNumber1' and 'pNumber2' !
						return ans

def 	__lowerCamelCase      (		snake_case__		)							->  List[Any]:






						"""simple docstring"""





						assert isinstance(snake_case__	,snake_case__		) and (n >= 1), "'n' must been int and >= 1"

						_SCREAMING_SNAKE_CASE     				=  []  # will be returned.

						for divisor in range(1	,n + 1		):
												if n % divisor == 0:
																		ans.append(snake_case__		)

    # precondition
						assert ans[0] == 1 and ans[len(snake_case__		) - 1] == n, "Error in function getDivisiors(...)"

						return ans

def 	__lowerCamelCase      (		snake_case__		)							->  int:






						"""simple docstring"""





						assert isinstance(snake_case__	,snake_case__		) and (
						    number > 1
						), "'number' must been an int and >= 1"

						_SCREAMING_SNAKE_CASE     				=  get_divisors(snake_case__		)

						# precondition
						assert (
						    isinstance(snake_case__	,snake_case__		)
						    and (divisors[0] == 1)
						    and (divisors[len(snake_case__		) - 1] == number)
						), "Error in help-function getDivisiors(...)"

						# summed all divisors up to 'number' (exclusive), hence [:-1]
						return sum(divisors[:-1]		) == number

def 	__lowerCamelCase      (		snake_case__	,snake_case__		)							->  Tuple:






						"""simple docstring"""





						assert (
						    isinstance(snake_case__	,snake_case__		)
						    and isinstance(snake_case__	,snake_case__		)
						    and (denominator != 0)
						), "The arguments must been from type int and 'denominator' != 0"

						# build the greatest common divisor of numerator and denominator.
						_SCREAMING_SNAKE_CASE     				=  gcd(abs(snake_case__		)	,abs(snake_case__		)		)

						# precondition
						assert (
						    isinstance(snake_case__	,snake_case__		)
						    and (numerator % gcd_of_fraction == 0)
						    and (denominator % gcd_of_fraction == 0)
						), "Error in function gcd(...,...)"

						return (numerator // gcd_of_fraction, denominator // gcd_of_fraction)

def 	__lowerCamelCase      (		snake_case__		)							->  Optional[int]:






						"""simple docstring"""





						assert isinstance(snake_case__	,snake_case__		) and (n >= 0), "'n' must been a int and >= 0"

						_SCREAMING_SNAKE_CASE     				=  1  # this will be return.

						for factor in range(1	,n + 1		):
												ans *= factor

						return ans




def 	__lowerCamelCase      (		snake_case__		)							->  Any:






						"""simple docstring"""





						assert isinstance(snake_case__	,snake_case__		) and (n >= 0), "'n' must been an int and >= 0"

						_SCREAMING_SNAKE_CASE     				=  0
						_SCREAMING_SNAKE_CASE     				=  1
						_SCREAMING_SNAKE_CASE     				=  1  # this will be return

						for _ in range(n - 1		):
												_SCREAMING_SNAKE_CASE     				=  ans
												ans += fiba
												_SCREAMING_SNAKE_CASE     				=  tmp

						return ans









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      (		snake_case__	,snake_case__		)							->  Any:






						"""simple docstring"""





						_SCREAMING_SNAKE_CASE     				=  []
						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      (		snake_case__	,snake_case__		)							->  Dict:






						"""simple docstring"""





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

												_SCREAMING_SNAKE_CASE     				=  in_proj_weight[
												    : encoder_config.hidden_size, :
												]
												_SCREAMING_SNAKE_CASE     				=  in_proj_weight[
												    encoder_config.hidden_size : encoder_config.hidden_size * 2, :
												]
												_SCREAMING_SNAKE_CASE     				=  in_proj_weight[
												    -encoder_config.hidden_size :, :
												]

def 	__lowerCamelCase      (		snake_case__	,snake_case__	,snake_case__		)							->  Union[str, Any]:






						"""simple docstring"""





						_SCREAMING_SNAKE_CASE     				=  dct.pop(snake_case__		)
						_SCREAMING_SNAKE_CASE     				=  val

def 	__lowerCamelCase      (		snake_case__		)							->  Union[str, Any]:






						"""simple docstring"""





						if "handwritten" in checkpoint_url:
												_SCREAMING_SNAKE_CASE     				=  """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     				=  """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     				=  Image.open(requests.get(snake_case__	,stream=snake_case__		).raw		).convert("""RGB"""		)
						return im

@torch.no_grad()
def 	__lowerCamelCase      (		snake_case__	,snake_case__		)							->  Tuple:






						"""simple docstring"""





						_SCREAMING_SNAKE_CASE     				=  ViTConfig(image_size=3_84	,qkv_bias=snake_case__		)
						_SCREAMING_SNAKE_CASE     				=  TrOCRConfig()

						# size of the architecture
						if "base" in checkpoint_url:
												_SCREAMING_SNAKE_CASE     				=  7_68
						elif "large" in checkpoint_url:
												# use ViT-large encoder
												_SCREAMING_SNAKE_CASE     				=  10_24
												_SCREAMING_SNAKE_CASE     				=  40_96
												_SCREAMING_SNAKE_CASE     				=  24
												_SCREAMING_SNAKE_CASE     				=  16
												_SCREAMING_SNAKE_CASE     				=  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     				=  False
												_SCREAMING_SNAKE_CASE     				=  """relu"""
												_SCREAMING_SNAKE_CASE     				=  10_24
												_SCREAMING_SNAKE_CASE     				=  True
												_SCREAMING_SNAKE_CASE     				=  False
												_SCREAMING_SNAKE_CASE     				=  False

						# load HuggingFace model
						_SCREAMING_SNAKE_CASE     				=  ViTModel(snake_case__	,add_pooling_layer=snake_case__		)
						_SCREAMING_SNAKE_CASE     				=  TrOCRForCausalLM(snake_case__		)
						_SCREAMING_SNAKE_CASE     				=  VisionEncoderDecoderModel(encoder=snake_case__	,decoder=snake_case__		)
						model.eval()

						# load state_dict of original model, rename some keys
						_SCREAMING_SNAKE_CASE     				=  torch.hub.load_state_dict_from_url(snake_case__	,map_location="""cpu"""	,check_hash=snake_case__		)["""model"""]

						_SCREAMING_SNAKE_CASE     				=  create_rename_keys(snake_case__	,snake_case__		)
						for src, dest in rename_keys:
												rename_key(snake_case__	,snake_case__	,snake_case__		)
						read_in_q_k_v(snake_case__	,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     				=  state_dict.pop(snake_case__		)
												if key.startswith("""decoder"""		) and "output_projection" not in key:
																		_SCREAMING_SNAKE_CASE     				=  val
												else:
																		_SCREAMING_SNAKE_CASE     				=  val

    # load state dict
						model.load_state_dict(snake_case__		)

						# Check outputs on an image
						_SCREAMING_SNAKE_CASE     				=  ViTImageProcessor(size=encoder_config.image_size		)
						_SCREAMING_SNAKE_CASE     				=  RobertaTokenizer.from_pretrained("""roberta-large"""		)
						_SCREAMING_SNAKE_CASE     				=  TrOCRProcessor(snake_case__	,snake_case__		)

						_SCREAMING_SNAKE_CASE     				=  processor(images=prepare_img(snake_case__		)	,return_tensors="""pt"""		).pixel_values

						# verify logits
						_SCREAMING_SNAKE_CASE     				=  torch.tensor([[model.config.decoder.decoder_start_token_id]]		)
						_SCREAMING_SNAKE_CASE     				=  model(pixel_values=snake_case__	,decoder_input_ids=snake_case__		)
						_SCREAMING_SNAKE_CASE     				=  outputs.logits

						_SCREAMING_SNAKE_CASE     				=  torch.Size([1, 1, 5_02_65]		)
						if "trocr-base-handwritten" in checkpoint_url:
												_SCREAMING_SNAKE_CASE     				=  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:
												_SCREAMING_SNAKE_CASE     				=  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:
												_SCREAMING_SNAKE_CASE     				=  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:
												_SCREAMING_SNAKE_CASE     				=  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]	,snake_case__	,atol=1e-3		), "First elements of logits not as expected"

						Path(snake_case__		).mkdir(exist_ok=snake_case__		)
						print(F'Saving model to {pytorch_dump_folder_path}'		)
						model.save_pretrained(snake_case__		)
						print(F'Saving processor to {pytorch_dump_folder_path}'		)
						processor.save_pretrained(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)

