




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()
lowercase__				:      Optional[int]	   = logging.get_logger(__name__)

def  SCREAMING_SNAKE_CASE							(       __UpperCamelCase		, __UpperCamelCase)      ->					Tuple:
     a     =				[]
     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  SCREAMING_SNAKE_CASE							(       __UpperCamelCase		, __UpperCamelCase)      ->					List[str]:
     for i in range(encoder_config.num_hidden_layers):
          # queries, keys and values (only weights, no biases)
          a     =				state_dict.pop(f'''encoder.deit.blocks.{i}.attn.qkv.weight''')

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

def  SCREAMING_SNAKE_CASE							(       __UpperCamelCase		, __UpperCamelCase		, __UpperCamelCase)      ->					Union[str, Any]:
     a     =				dct.pop(__UpperCamelCase)
     a     =				val

def  SCREAMING_SNAKE_CASE							(       __UpperCamelCase)      ->					List[str]:
     if "handwritten" in checkpoint_url:
          a     =				"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:
          a     =				"https://www.researchgate.net/profile/Dinh-Sang/publication/338099565/figure/fig8/AS:840413229350922@1577381536857/An-receipt-example-in-the-SROIE-2019-dataset_Q640.jpg"
     a     =				Image.open(requests.get(__UpperCamelCase		, stream=__UpperCamelCase).raw).convert("RGB")
     return im

@torch.no_grad()
def  SCREAMING_SNAKE_CASE							(       __UpperCamelCase		, __UpperCamelCase)      ->					Dict:
     a     =				ViTConfig(image_size=3_84		, qkv_bias=__UpperCamelCase)
     a     =				TrOCRConfig()

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

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

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

     a     =				create_rename_keys(__UpperCamelCase		, __UpperCamelCase)
     for src, dest in rename_keys:
          rename_key(__UpperCamelCase		, __UpperCamelCase		, __UpperCamelCase)
     read_in_q_k_v(__UpperCamelCase		, __UpperCamelCase)

     # 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():
          a     =				state_dict.pop(__UpperCamelCase)
          if key.startswith("decoder") and "output_projection" not in key:
               a     =				val
          else:
               a     =				val

    # load state dict
     model.load_state_dict(__UpperCamelCase)

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

     a     =				processor(images=prepare_img(__UpperCamelCase)		, return_tensors="pt").pixel_values

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

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

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


if __name__ == "__main__":
   lowercase__				:      Optional[int]	   = 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."
   )
   lowercase__				:      str	   = parser.parse_args()
   convert_tr_ocr_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)







import math

def  SCREAMING_SNAKE_CASE							(       __UpperCamelCase)      ->					bool:

     if 1 < number < 4:
          # 2 and 3 are primes
          return True
     elif number < 2 or number % 2 == 0 or number % 3 == 0:
          # Negatives, 0, 1, all even numbers, all multiples of 3 are not primes
          return False

     # All primes number are in format of 6k +/- 1
     for i in range(5		, int(math.sqrt(__UpperCamelCase) + 1)		, 6):
          if number % i == 0 or number % (i + 2) == 0:
               return False
     return True

def  SCREAMING_SNAKE_CASE							(       __UpperCamelCase = 0.1)      ->					int:
     a     =				3
     a     =				3

     while primes / (2 * j - 1) >= ratio:
          for i in range(j * j + j + 1		, (j + 2) * (j + 2)		, j + 1):
               primes += is_prime(__UpperCamelCase)
          j += 2
     return j


if __name__ == "__main__":
   import doctest

   doctest.testmod()

