



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 		UpperCamelCase     (       snake_case__	:   Optional[int]	,					snake_case__	:   Optional[Any]     )     ->					Optional[int]:
   UpperCamelCase   :  List[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 		UpperCamelCase     (       snake_case__	:   List[Any]	,					snake_case__	:   List[Any]     )     ->					Dict:
   for i in range(encoder_config.num_hidden_layers     ):
      # queries, keys and values (only weights, no biases)
      UpperCamelCase   :  Any						     =     state_dict.pop(F"""encoder.deit.blocks.{i}.attn.qkv.weight"""     )

      UpperCamelCase   :  Tuple						     =     in_proj_weight[
          : encoder_config.hidden_size, :
      ]
      UpperCamelCase   :  str						     =     in_proj_weight[
          encoder_config.hidden_size : encoder_config.hidden_size * 2, :
      ]
      UpperCamelCase   :  Dict						     =     in_proj_weight[
          -encoder_config.hidden_size :, :
      ]
def 		UpperCamelCase     (       snake_case__	:   Optional[int]	,					snake_case__	:   int	,					snake_case__	:   List[Any]     )     ->					Tuple:
   UpperCamelCase   :  List[str]						     =     dct.pop(snake_case__     )
   UpperCamelCase   :  Optional[int]						     =     val
def 		UpperCamelCase     (       snake_case__	:   int     )     ->					Any:
   if "handwritten" in checkpoint_url:
      UpperCamelCase   :  Optional[int]						     =     '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:
      UpperCamelCase   :  Dict						     =     'https://www.researchgate.net/profile/Dinh-Sang/publication/338099565/figure/fig8/AS:840413229350922@1577381536857/An-receipt-example-in-the-SROIE-2019-dataset_Q640.jpg'
   UpperCamelCase   :  List[str]						     =     Image.open(requests.get(snake_case__	,					stream=snake_case__     ).raw     ).convert('RGB'     )
   return im
@torch.no_grad()
def 		UpperCamelCase     (       snake_case__	:   Optional[int]	,					snake_case__	:   Any     )     ->					Tuple:
   UpperCamelCase   :  Optional[int]						     =     ViTConfig(image_size=384	,					qkv_bias=snake_case__     )
   UpperCamelCase   :  str						     =     TrOCRConfig()

   # size of the architecture
   if "base" in checkpoint_url:
      UpperCamelCase   :  Optional[int]						     =     768
   elif "large" in checkpoint_url:
      # use ViT-large encoder
      UpperCamelCase   :  Union[str, Any]						     =     1024
      UpperCamelCase   :  int						     =     4096
      UpperCamelCase   :  Optional[Any]						     =     24
      UpperCamelCase   :  Union[str, Any]						     =     16
      UpperCamelCase   :  Dict						     =     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:
      UpperCamelCase   :  Any						     =     False
      UpperCamelCase   :  List[Any]						     =     'relu'
      UpperCamelCase   :  Tuple						     =     1024
      UpperCamelCase   :  Optional[int]						     =     True
      UpperCamelCase   :  Union[str, Any]						     =     False
      UpperCamelCase   :  List[Any]						     =     False

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

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

   UpperCamelCase   :  int						     =     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():
      UpperCamelCase   :  Dict						     =     state_dict.pop(snake_case__     )
      if key.startswith('decoder'     ) and "output_projection" not in key:
         UpperCamelCase   :  Tuple						     =     val
      else:
         UpperCamelCase   :  str						     =     val

    # load state dict
   model.load_state_dict(snake_case__     )

   # Check outputs on an image
   UpperCamelCase   :  Optional[Any]						     =     ViTImageProcessor(size=encoder_config.image_size     )
   UpperCamelCase   :  List[str]						     =     RobertaTokenizer.from_pretrained('roberta-large'     )
   UpperCamelCase   :  Optional[Any]						     =     TrOCRProcessor(snake_case__	,					snake_case__     )

   UpperCamelCase   :  int						     =     processor(images=prepare_img(snake_case__     )	,					return_tensors='pt'     ).pixel_values

   # verify logits
   UpperCamelCase   :  Any						     =     torch.tensor([[model.config.decoder.decoder_start_token_id]]     )
   UpperCamelCase   :  int						     =     model(pixel_values=snake_case__	,					decoder_input_ids=snake_case__     )
   UpperCamelCase   :  Dict						     =     outputs.logits

   UpperCamelCase   :  Union[str, Any]						     =     torch.Size([1, 1, 50265]     )
   if "trocr-base-handwritten" in checkpoint_url:
      UpperCamelCase   :  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:
      UpperCamelCase   :  List[str]						     =     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:
      UpperCamelCase   :  Union[str, Any]						     =     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:
      UpperCamelCase   :  Tuple						     =     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]	,					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)








import re
def 		UpperCamelCase     (       snake_case__	:   str     )     ->					str:

   if len(re.findall('[ATCG]'	,					snake_case__     )     ) != len(snake_case__     ):
      raise ValueError('Invalid Strand'     )

   return dna.translate(dna.maketrans('ATCG'	,					'TAGC'     )     )


if __name__ == "__main__":
    import doctest

    doctest.testmod()



