





'''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()
a   	=    logging.get_logger(__name__)

def     __magic_name__							(				__UpperCAmelCase				,						__UpperCAmelCase   )    ->       Tuple:

   '''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     __magic_name__							(				__UpperCAmelCase				,						__UpperCAmelCase   )    ->       Tuple:

   '''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     __magic_name__							(				__UpperCAmelCase				,						__UpperCAmelCase				,						__UpperCAmelCase   )    ->       List[Any]:

   '''simple docstring'''







   __SCREAMING_SNAKE_CASE   				=							dct.pop(__UpperCAmelCase   )
   __SCREAMING_SNAKE_CASE   				=							val

def     __magic_name__							(				__UpperCAmelCase   )    ->       List[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(__UpperCAmelCase				,						stream=__UpperCAmelCase   ).raw   ).convert("""RGB"""   )
   return im

@torch.no_grad()
def     __magic_name__							(				__UpperCAmelCase				,						__UpperCAmelCase   )    ->       int:

   '''simple docstring'''







   __SCREAMING_SNAKE_CASE   				=							ViTConfig(image_size=384				,						qkv_bias=__UpperCAmelCase   )
   __SCREAMING_SNAKE_CASE   				=							TrOCRConfig()

   # size of the architecture
   if "base" in checkpoint_url:
      __SCREAMING_SNAKE_CASE   				=							768
   elif "large" in checkpoint_url:
      # use ViT-large encoder
      __SCREAMING_SNAKE_CASE   				=							1024
      __SCREAMING_SNAKE_CASE   				=							4096
      __SCREAMING_SNAKE_CASE   				=							24
      __SCREAMING_SNAKE_CASE   				=							16
      __SCREAMING_SNAKE_CASE   				=							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:
      __SCREAMING_SNAKE_CASE   				=							False
      __SCREAMING_SNAKE_CASE   				=							"""relu"""
      __SCREAMING_SNAKE_CASE   				=							1024
      __SCREAMING_SNAKE_CASE   				=							True
      __SCREAMING_SNAKE_CASE   				=							False
      __SCREAMING_SNAKE_CASE   				=							False

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

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

   __SCREAMING_SNAKE_CASE   				=							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():
      __SCREAMING_SNAKE_CASE   				=							state_dict.pop(__UpperCAmelCase   )
      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(__UpperCAmelCase   )

   # 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(__UpperCAmelCase				,						__UpperCAmelCase   )

   __SCREAMING_SNAKE_CASE   				=							processor(images=prepare_img(__UpperCAmelCase   )				,						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=__UpperCAmelCase				,						decoder_input_ids=__UpperCAmelCase   )
   __SCREAMING_SNAKE_CASE   				=							outputs.logits

   __SCREAMING_SNAKE_CASE   				=							torch.Size([1, 1, 50265]   )
   if "trocr-base-handwritten" in checkpoint_url:
      __SCREAMING_SNAKE_CASE   				=							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:
      __SCREAMING_SNAKE_CASE   				=							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:
      __SCREAMING_SNAKE_CASE   				=							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:
      __SCREAMING_SNAKE_CASE   				=							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]				,						__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__":
    a   	=    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."
    )
    a   	=    parser.parse_args()
    convert_tr_ocr_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)









import numpy as np
import torch
import torch.nn as nn
from transformers import CLIPConfig, CLIPVisionModelWithProjection, PreTrainedModel

from ...utils import logging


snake_case__				:	int					       =					logging.get_logger(__name__)





class   SCREAMING_SNAKE_CASE_			(a__						):




       '''simple docstring'''



       _a   							=							CLIPConfig

       _a   							=							["CLIPEncoderLayer"]



       def __init__(						self				:				Tuple		,				__a				:				CLIPConfig							)					->Union[str, Any]:
          super().__init__(__a							)

          lowerCamelCase_						:   List[Any]		= CLIPVisionModelWithProjection(config.vision_config							)

          lowerCamelCase_						:   int		= nn.Linear(config.vision_config.projection_dim		,				1							)
          lowerCamelCase_						:   int		= nn.Linear(config.vision_config.projection_dim		,				1							)




       @torch.no_grad()
       def 				_lowerCAmelCase							(						self				:				Dict		,				__a				:				List[Any]		,				__a				:				Tuple		,				__a				:				int=0.5		,				__a				:				Optional[Any]=0.5							)					->Union[str, Any]:
          lowerCamelCase_						:   Dict		= self.vision_model(__a							)[0]

          lowerCamelCase_						:   str		= self.p_head(__a							)
          lowerCamelCase_						:   Union[str, Any]		= nsfw_detected.flatten()
          lowerCamelCase_						:   Tuple		= nsfw_detected > p_threshold
          lowerCamelCase_						:   Dict		= nsfw_detected.tolist()

          if any(__a							):
             logger.warning(
                 """Potential NSFW content was detected in one or more images. A black image will be returned instead."""
                 """ Try again with a different prompt and/or seed."""							)

          for idx, nsfw_detected_ in enumerate(__a							):
             if nsfw_detected_:
                lowerCamelCase_						:   int		= np.zeros(images[idx].shape							)

          lowerCamelCase_						:   Union[str, Any]		= self.w_head(__a							)
          lowerCamelCase_						:   List[str]		= watermark_detected.flatten()
          lowerCamelCase_						:   Dict		= watermark_detected > w_threshold
          lowerCamelCase_						:   Optional[int]		= watermark_detected.tolist()

          if any(__a							):
             logger.warning(
                 """Potential watermarked content was detected in one or more images. A black image will be returned instead."""
                 """ Try again with a different prompt and/or seed."""							)

          for idx, watermark_detected_ in enumerate(__a							):
             if watermark_detected_:
                lowerCamelCase_						:   Tuple		= np.zeros(images[idx].shape							)

          return images, nsfw_detected, watermark_detected





