


'''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 					a__       (					lowerCAmelCase__  ,					lowerCAmelCase__						)    ->	str:
    UpperCAmelCase__			:		int													=				[]
    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 					a__       (					lowerCAmelCase__  ,					lowerCAmelCase__						)    ->	List[Any]:
    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__			:		str													=				in_proj_weight[
            : encoder_config.hidden_size, :
        ]
        UpperCAmelCase__			:		int													=				in_proj_weight[
            encoder_config.hidden_size : encoder_config.hidden_size * 2, :
        ]
        UpperCAmelCase__			:		Tuple													=				in_proj_weight[
            -encoder_config.hidden_size :, :
        ]





def 					a__       (					lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__						)    ->	Tuple:
    UpperCAmelCase__			:		Optional[Any]													=				dct.pop(lowerCAmelCase__						)
    UpperCAmelCase__			:		Optional[int]													=				val





def 					a__       (					lowerCAmelCase__						)    ->	Union[str, Any]:
    if "handwritten" in checkpoint_url:
        UpperCAmelCase__			:		Dict													=				'''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__			:		int													=				'''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__			:		Union[str, Any]													=				Image.open(requests.get(lowerCAmelCase__  ,					stream=lowerCAmelCase__						).raw						).convert('''RGB'''						)
    return im





@torch.no_grad()
def 					a__       (					lowerCAmelCase__  ,					lowerCAmelCase__						)    ->	List[str]:
    UpperCAmelCase__			:		Tuple													=				ViTConfig(image_size=3_84  ,					qkv_bias=lowerCAmelCase__						)
    UpperCAmelCase__			:		int													=				TrOCRConfig()

    # size of the architecture
    if "base" in checkpoint_url:
        UpperCAmelCase__			:		int													=				7_68
    elif "large" in checkpoint_url:
        # use ViT-large encoder
        UpperCAmelCase__			:		str													=				10_24
        UpperCAmelCase__			:		Any													=				40_96
        UpperCAmelCase__			:		Optional[Any]													=				24
        UpperCAmelCase__			:		int													=				16
        UpperCAmelCase__			:		Union[str, Any]													=				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:
        UpperCAmelCase__			:		Any													=				False
        UpperCAmelCase__			:		Dict													=				'''relu'''
        UpperCAmelCase__			:		Union[str, Any]													=				10_24
        UpperCAmelCase__			:		int													=				True
        UpperCAmelCase__			:		List[Any]													=				False
        UpperCAmelCase__			:		Optional[int]													=				False

    # load HuggingFace model
    UpperCAmelCase__			:		Tuple													=				ViTModel(lowerCAmelCase__  ,					add_pooling_layer=lowerCAmelCase__						)
    UpperCAmelCase__			:		Union[str, Any]													=				TrOCRForCausalLM(lowerCAmelCase__						)
    UpperCAmelCase__			:		Tuple													=				VisionEncoderDecoderModel(encoder=lowerCAmelCase__  ,					decoder=lowerCAmelCase__						)
    model.eval()

    # load state_dict of original model, rename some keys
    UpperCAmelCase__			:		int													=				torch.hub.load_state_dict_from_url(lowerCAmelCase__  ,					map_location='''cpu'''  ,					check_hash=lowerCAmelCase__						)['''model''']

    UpperCAmelCase__			:		Tuple													=				create_rename_keys(lowerCAmelCase__  ,					lowerCAmelCase__						)
    for src, dest in rename_keys:
        rename_key(lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__						)
    read_in_q_k_v(lowerCAmelCase__  ,					lowerCAmelCase__						)

    # 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__			:		Optional[int]													=				state_dict.pop(lowerCAmelCase__						)
        if key.startswith('''decoder'''						) and "output_projection" not in key:
            UpperCAmelCase__			:		Tuple													=				val
        else:
            UpperCAmelCase__			:		Union[str, Any]													=				val

    # load state dict
    model.load_state_dict(lowerCAmelCase__						)

    # Check outputs on an image
    UpperCAmelCase__			:		Optional[Any]													=				ViTImageProcessor(size=encoder_config.image_size						)
    UpperCAmelCase__			:		Tuple													=				RobertaTokenizer.from_pretrained('''roberta-large'''						)
    UpperCAmelCase__			:		int													=				TrOCRProcessor(lowerCAmelCase__  ,					lowerCAmelCase__						)

    UpperCAmelCase__			:		str													=				processor(images=prepare_img(lowerCAmelCase__						)  ,					return_tensors='''pt'''						).pixel_values

    # verify logits
    UpperCAmelCase__			:		Optional[Any]													=				torch.tensor([[model.config.decoder.decoder_start_token_id]]						)
    UpperCAmelCase__			:		Union[str, Any]													=				model(pixel_values=lowerCAmelCase__  ,					decoder_input_ids=lowerCAmelCase__						)
    UpperCAmelCase__			:		Tuple													=				outputs.logits

    UpperCAmelCase__			:		Dict													=				torch.Size([1, 1, 5_02_65]						)
    if "trocr-base-handwritten" in checkpoint_url:
        UpperCAmelCase__			:		Any													=				torch.tensor(
            [-1.4_5_0_2, -4.6_6_8_3, -0.5_3_4_7, -2.9_2_9_1, 9.1_4_3_5, -3.0_5_7_1, 8.9_7_6_4, 1.7_5_6_0, 8.7_3_5_8, -1.5_3_1_1]						)
    elif "trocr-large-handwritten" in checkpoint_url:
        UpperCAmelCase__			:		Optional[int]													=				torch.tensor(
            [-2.6_4_3_7, -1.3_1_2_9, -2.2_5_9_6, -5.3_4_5_5, 6.3_5_3_9, 1.7_6_0_4, 5.4_9_9_1, 1.4_7_0_2, 5.6_1_1_3, 2.0_1_7_0]						)
    elif "trocr-base-printed" in checkpoint_url:
        UpperCAmelCase__			:		Tuple													=				torch.tensor(
            [-5.6_8_1_6, -5.8_3_8_8, 1.1_3_9_8, -6.9_0_3_4, 6.8_5_0_5, -2.4_3_9_3, 1.2_2_8_4, -1.0_2_3_2, -1.9_6_6_1, -3.9_2_1_0]						)
    elif "trocr-large-printed" in checkpoint_url:
        UpperCAmelCase__			:		int													=				torch.tensor(
            [-6.0_1_6_2, -7.0_9_5_9, 4.4_1_5_5, -5.1_0_6_3, 7.0_4_6_8, -3.1_6_3_1, 2.6_4_6_6, -0.3_0_8_1, -0.8_1_0_6, -1.7_5_3_5]						)

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

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


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)




'''simple docstring'''
import pprint

import requests

UpperCamelCase__    					=      '''https://zenquotes.io/api'''





def 					a__       (					)    ->	list:
    return requests.get(API_ENDPOINT_URL + '''/today'''						).json()





def 					a__       (					)    ->	list:
    return requests.get(API_ENDPOINT_URL + '''/random'''						).json()


if __name__ == "__main__":
  UpperCamelCase__    					=      random_quotes()
  pprint.pprint(response)
