

import inspect
import unittest

from transformers import SegformerConfig, is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch

       from transformers import (
           MODEL_MAPPING,
           SegformerForImageClassification,
           SegformerForSemanticSegmentation,
           SegformerModel,
       )
       from transformers.models.segformer.modeling_segformer import SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
       from PIL import Image

       from transformers import SegformerImageProcessor
class        A_ (   __snake_case    ):


    '''simple docstring'''






    def   lowerCAmelCase_    (self	)  ->   Dict:
        __UpperCAmelCase					       =      self.config_class(**self.inputs_dict	)
        self.parent.assertTrue(hasattr(__lowerCamelCase  ,   '''hidden_sizes'''	)	)
        self.parent.assertTrue(hasattr(__lowerCamelCase  ,   '''num_attention_heads'''	)	)
        self.parent.assertTrue(hasattr(__lowerCamelCase  ,   '''num_encoder_blocks'''	)	)






class        A_ :


    '''simple docstring'''






    def __init__(self  ,   lowercase__  ,   lowercase__=13  ,   lowercase__=64  ,   lowercase__=3  ,   lowercase__=4  ,   lowercase__=[2, 2, 2, 2]  ,   lowercase__=[8, 4, 2, 1]  ,   lowercase__=[16, 32, 64, 128]  ,   lowercase__=[1, 4, 8, 16]  ,   lowercase__=[1, 2, 4, 8]  ,   lowercase__=True  ,   lowercase__=True  ,   lowercase__="gelu"  ,   lowercase__=0.1  ,   lowercase__=0.1  ,   lowercase__=0.02  ,   lowercase__=3  ,   lowercase__=None  ,   )  ->   Any:
        __UpperCAmelCase					       =      parent
        __UpperCAmelCase					       =      batch_size
        __UpperCAmelCase					       =      image_size
        __UpperCAmelCase					       =      num_channels
        __UpperCAmelCase					       =      num_encoder_blocks
        __UpperCAmelCase					       =      sr_ratios
        __UpperCAmelCase					       =      depths
        __UpperCAmelCase					       =      hidden_sizes
        __UpperCAmelCase					       =      downsampling_rates
        __UpperCAmelCase					       =      num_attention_heads
        __UpperCAmelCase					       =      is_training
        __UpperCAmelCase					       =      use_labels
        __UpperCAmelCase					       =      hidden_act
        __UpperCAmelCase					       =      hidden_dropout_prob
        __UpperCAmelCase					       =      attention_probs_dropout_prob
        __UpperCAmelCase					       =      initializer_range
        __UpperCAmelCase					       =      num_labels
        __UpperCAmelCase					       =      scope


    def   lowerCAmelCase_    (self	)  ->   Union[str, Any]:
        __UpperCAmelCase					       =      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

        __UpperCAmelCase					       =      None
        if self.use_labels:
            __UpperCAmelCase					       =      ids_tensor([self.batch_size, self.image_size, self.image_size]  ,   self.num_labels	)

        __UpperCAmelCase					       =      self.get_config()
        return config, pixel_values, labels


    def   lowerCAmelCase_    (self	)  ->   List[Any]:
        return SegformerConfig(
            image_size=self.image_size  ,   num_channels=self.num_channels  ,   num_encoder_blocks=self.num_encoder_blocks  ,   depths=self.depths  ,   hidden_sizes=self.hidden_sizes  ,   num_attention_heads=self.num_attention_heads  ,   hidden_act=self.hidden_act  ,   hidden_dropout_prob=self.hidden_dropout_prob  ,   attention_probs_dropout_prob=self.attention_probs_dropout_prob  ,   initializer_range=self.initializer_range  ,   )


    def   lowerCAmelCase_    (self  ,   lowercase__  ,   lowercase__  ,   lowercase__	)  ->   int:
        __UpperCAmelCase					       =      SegformerModel(config=__lowerCamelCase	)
        model.to(__lowerCamelCase	)
        model.eval()
        __UpperCAmelCase					       =      model(__lowerCamelCase	)
        __UpperCAmelCase					       =      __UpperCAmelCase					       =      self.image_size // (self.downsampling_rates[-1] * 2)
        self.parent.assertEqual(
            result.last_hidden_state.shape  ,   (self.batch_size, self.hidden_sizes[-1], expected_height, expected_width)	)


    def   lowerCAmelCase_    (self  ,   lowercase__  ,   lowercase__  ,   lowercase__	)  ->   Optional[Any]:
        __UpperCAmelCase					       =      self.num_labels
        __UpperCAmelCase					       =      SegformerForSemanticSegmentation(__lowerCamelCase	)
        model.to(__lowerCamelCase	)
        model.eval()
        __UpperCAmelCase					       =      model(__lowerCamelCase	)
        self.parent.assertEqual(
            result.logits.shape  ,   (self.batch_size, self.num_labels, self.image_size // 4, self.image_size // 4)	)
        __UpperCAmelCase					       =      model(__lowerCamelCase  ,   labels=__lowerCamelCase	)
        self.parent.assertEqual(
            result.logits.shape  ,   (self.batch_size, self.num_labels, self.image_size // 4, self.image_size // 4)	)
        self.parent.assertGreater(result.loss  ,   0.0	)


    def   lowerCAmelCase_    (self  ,   lowercase__  ,   lowercase__  ,   lowercase__	)  ->   List[Any]:
        __UpperCAmelCase					       =      1
        __UpperCAmelCase					       =      SegformerForSemanticSegmentation(config=__lowerCamelCase	)
        model.to(__lowerCamelCase	)
        model.eval()
        __UpperCAmelCase					       =      torch.randint(0  ,   1  ,   (self.batch_size, self.image_size, self.image_size)	).to(__lowerCamelCase	)
        __UpperCAmelCase					       =      model(__lowerCamelCase  ,   labels=__lowerCamelCase	)
        self.parent.assertGreater(result.loss  ,   0.0	)






    def   lowerCAmelCase_    (self	)  ->   str:
        __UpperCAmelCase					       =      self.prepare_config_and_inputs()
        __UpperCAmelCase       ,  __UpperCAmelCase       ,  __UpperCAmelCase					       =      config_and_inputs
        __UpperCAmelCase					       =      {'''pixel_values''': pixel_values}
        return config, inputs_dict






@require_torch
class        A_ (   __snake_case			,  __snake_case			,  unittest.TestCase    ):


    '''simple docstring'''






    a__  				=							(
        (
            SegformerModel,
            SegformerForSemanticSegmentation,
            SegformerForImageClassification,
        )
        if is_torch_available()
        else ()
    )
    a__  				=							(
        {
            "feature-extraction": SegformerModel,
            "image-classification": SegformerForImageClassification,
            "image-segmentation": SegformerForSemanticSegmentation,
        }
        if is_torch_available()
        else {}
    )

    a__  				=							True
    a__  				=							False
    a__  				=							False
    a__  				=							False


    def   lowerCAmelCase_    (self	)  ->   Any:
        __UpperCAmelCase					       =      SegformerModelTester(self	)
        __UpperCAmelCase					       =      SegformerConfigTester(self  ,   config_class=__lowerCamelCase	)


    def   lowerCAmelCase_    (self	)  ->   Optional[int]:
        self.config_tester.run_common_tests()


    def   lowerCAmelCase_    (self	)  ->   Any:
        __UpperCAmelCase					       =      self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__lowerCamelCase	)


    def   lowerCAmelCase_    (self	)  ->   Tuple:
        __UpperCAmelCase					       =      self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_binary_image_segmentation(*__lowerCamelCase	)


    def   lowerCAmelCase_    (self	)  ->   Union[str, Any]:
        __UpperCAmelCase					       =      self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_segmentation(*__lowerCamelCase	)


    @unittest.skip('''SegFormer does not use inputs_embeds'''	)
    def   lowerCAmelCase_    (self	)  ->   Dict:
        pass


    @unittest.skip('''SegFormer does not have get_input_embeddings method and get_output_embeddings methods'''	)
    def   lowerCAmelCase_    (self	)  ->   Union[str, Any]:
        pass


    def   lowerCAmelCase_    (self	)  ->   Optional[int]:
        __UpperCAmelCase       ,  __UpperCAmelCase					       =      self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            __UpperCAmelCase					       =      model_class(__lowerCamelCase	)
            __UpperCAmelCase					       =      inspect.signature(model.forward	)
            # signature.parameters is an OrderedDict => so arg_names order is deterministic
            __UpperCAmelCase					       =      [*signature.parameters.keys()]

            __UpperCAmelCase					       =      ['''pixel_values''']
            self.assertListEqual(arg_names[:1]  ,   __lowerCamelCase	)


    def   lowerCAmelCase_    (self	)  ->   Union[str, Any]:
        __UpperCAmelCase       ,  __UpperCAmelCase					       =      self.model_tester.prepare_config_and_inputs_for_common()
        __UpperCAmelCase					       =      True

        for model_class in self.all_model_classes:
            __UpperCAmelCase					       =      True
            __UpperCAmelCase					       =      False
            __UpperCAmelCase					       =      True
            __UpperCAmelCase					       =      model_class(__lowerCamelCase	)
            model.to(__lowerCamelCase	)
            model.eval()
            with torch.no_grad():
                __UpperCAmelCase					       =      model(**self._prepare_for_class(__lowerCamelCase  ,   __lowerCamelCase	)	)
            __UpperCAmelCase					       =      outputs.attentions

            __UpperCAmelCase					       =      sum(self.model_tester.depths	)
            self.assertEqual(len(__lowerCamelCase	)  ,   __lowerCamelCase	)

            # check that output_attentions also work using config
            del inputs_dict["output_attentions"]
            __UpperCAmelCase					       =      True
            __UpperCAmelCase					       =      model_class(__lowerCamelCase	)
            model.to(__lowerCamelCase	)
            model.eval()
            with torch.no_grad():
                __UpperCAmelCase					       =      model(**self._prepare_for_class(__lowerCamelCase  ,   __lowerCamelCase	)	)
            __UpperCAmelCase					       =      outputs.attentions

            self.assertEqual(len(__lowerCamelCase	)  ,   __lowerCamelCase	)

            # verify the first attentions (first block, first layer)
            __UpperCAmelCase					       =      (self.model_tester.image_size // 4) ** 2
            __UpperCAmelCase					       =      (self.model_tester.image_size // (4 * self.model_tester.sr_ratios[0])) ** 2
            self.assertListEqual(
                list(attentions[0].shape[-3:]	)  ,   [self.model_tester.num_attention_heads[0], expected_seq_len, expected_reduced_seq_len]  ,   )

            # verify the last attentions (last block, last layer)
            __UpperCAmelCase					       =      (self.model_tester.image_size // 32) ** 2
            __UpperCAmelCase					       =      (self.model_tester.image_size // (32 * self.model_tester.sr_ratios[-1])) ** 2
            self.assertListEqual(
                list(attentions[-1].shape[-3:]	)  ,   [self.model_tester.num_attention_heads[-1], expected_seq_len, expected_reduced_seq_len]  ,   )
            __UpperCAmelCase					       =      len(__lowerCamelCase	)

            # Check attention is always last and order is fine
            __UpperCAmelCase					       =      True
            __UpperCAmelCase					       =      True
            __UpperCAmelCase					       =      model_class(__lowerCamelCase	)
            model.to(__lowerCamelCase	)
            model.eval()
            with torch.no_grad():
                __UpperCAmelCase					       =      model(**self._prepare_for_class(__lowerCamelCase  ,   __lowerCamelCase	)	)

            self.assertEqual(out_len + 1  ,   len(__lowerCamelCase	)	)

            __UpperCAmelCase					       =      outputs.attentions

            self.assertEqual(len(__lowerCamelCase	)  ,   __lowerCamelCase	)
            # verify the first attentions (first block, first layer)
            __UpperCAmelCase					       =      (self.model_tester.image_size // 4) ** 2
            __UpperCAmelCase					       =      (self.model_tester.image_size // (4 * self.model_tester.sr_ratios[0])) ** 2
            self.assertListEqual(
                list(self_attentions[0].shape[-3:]	)  ,   [self.model_tester.num_attention_heads[0], expected_seq_len, expected_reduced_seq_len]  ,   )


    def   lowerCAmelCase_    (self	)  ->   List[str]:







        def check_hidden_states_output(lowercase__  ,   lowercase__  ,   lowercase__	):
            __UpperCAmelCase					       =      model_class(__lowerCamelCase	)
            model.to(__lowerCamelCase	)
            model.eval()

            with torch.no_grad():
                __UpperCAmelCase					       =      model(**self._prepare_for_class(__lowerCamelCase  ,   __lowerCamelCase	)	)

            __UpperCAmelCase					       =      outputs.hidden_states

            __UpperCAmelCase					       =      self.model_tester.num_encoder_blocks
            self.assertEqual(len(__lowerCamelCase	)  ,   __lowerCamelCase	)

            # verify the first hidden states (first block)
            self.assertListEqual(
                list(hidden_states[0].shape[-3:]	)  ,   [
                    self.model_tester.hidden_sizes[0],
                    self.model_tester.image_size // 4,
                    self.model_tester.image_size // 4,
                ]  ,   )

        __UpperCAmelCase       ,  __UpperCAmelCase					       =      self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            __UpperCAmelCase					       =      True
            check_hidden_states_output(__lowerCamelCase  ,   __lowerCamelCase  ,   __lowerCamelCase	)

            # check that output_hidden_states also work using config
            del inputs_dict["output_hidden_states"]
            __UpperCAmelCase					       =      True

            check_hidden_states_output(__lowerCamelCase  ,   __lowerCamelCase  ,   __lowerCamelCase	)


    def   lowerCAmelCase_    (self	)  ->   str:







        if not self.model_tester.is_training:
            return

        __UpperCAmelCase       ,  __UpperCAmelCase					       =      self.model_tester.prepare_config_and_inputs_for_common()
        __UpperCAmelCase					       =      True

        for model_class in self.all_model_classes:
            if model_class in get_values(__lowerCamelCase	):
                continue

            __UpperCAmelCase					       =      model_class(__lowerCamelCase	)
            model.to(__lowerCamelCase	)
            model.train()
            __UpperCAmelCase					       =      self._prepare_for_class(__lowerCamelCase  ,   __lowerCamelCase  ,   return_labels=__lowerCamelCase	)
            __UpperCAmelCase					       =      model(**__lowerCamelCase	).loss
            loss.backward()


    @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''	)
    def   lowerCAmelCase_    (self	)  ->   List[str]:
        pass






    @slow
    def   lowerCAmelCase_    (self	)  ->   str:







        for model_name in SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            __UpperCAmelCase					       =      SegformerModel.from_pretrained(__lowerCamelCase	)
            self.assertIsNotNone(__lowerCamelCase	)






def    __a							( )   ->  Any:






    '''simple docstring'''
    __UpperCAmelCase					       =      Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''   )
    return image
@require_torch
class        A_ (   unittest.TestCase    ):


    '''simple docstring'''






    @slow
    def   lowerCAmelCase_    (self	)  ->   Tuple:
        __UpperCAmelCase					       =      SegformerImageProcessor(
            image_scale=(512, 512)  ,   keep_ratio=__lowerCamelCase  ,   align=__lowerCamelCase  ,   do_random_crop=__lowerCamelCase	)
        __UpperCAmelCase					       =      SegformerForSemanticSegmentation.from_pretrained('''nvidia/segformer-b0-finetuned-ade-512-512'''	).to(
            __lowerCamelCase	)

        __UpperCAmelCase					       =      prepare_img()
        __UpperCAmelCase					       =      image_processor(images=__lowerCamelCase  ,   return_tensors='''pt'''	)
        __UpperCAmelCase					       =      encoded_inputs.pixel_values.to(__lowerCamelCase	)

        with torch.no_grad():
            __UpperCAmelCase					       =      model(__lowerCamelCase	)

        __UpperCAmelCase					       =      torch.Size((1, model.config.num_labels, 128, 128)	)
        self.assertEqual(outputs.logits.shape  ,   __lowerCamelCase	)

        __UpperCAmelCase					       =      torch.tensor(
            [
                [[-4.6310, -5.5232, -6.2356], [-5.1921, -6.1444, -6.5996], [-5.4424, -6.2790, -6.7574]],
                [[-12.1391, -13.3122, -13.9554], [-12.8732, -13.9352, -14.3563], [-12.9438, -13.8226, -14.2513]],
                [[-12.5134, -13.4686, -14.4915], [-12.8669, -14.4343, -14.7758], [-13.2523, -14.5819, -15.0694]],
            ]	).to(__lowerCamelCase	)
        self.assertTrue(torch.allclose(outputs.logits[0, :3, :3, :3]  ,   __lowerCamelCase  ,   atol=1E-4	)	)


    @slow
    def   lowerCAmelCase_    (self	)  ->   Union[str, Any]:
        __UpperCAmelCase					       =      SegformerImageProcessor(
            image_scale=(512, 512)  ,   keep_ratio=__lowerCamelCase  ,   align=__lowerCamelCase  ,   do_random_crop=__lowerCamelCase	)
        __UpperCAmelCase					       =      SegformerForSemanticSegmentation.from_pretrained(
            '''nvidia/segformer-b1-finetuned-cityscapes-1024-1024'''	).to(__lowerCamelCase	)

        __UpperCAmelCase					       =      prepare_img()
        __UpperCAmelCase					       =      image_processor(images=__lowerCamelCase  ,   return_tensors='''pt'''	)
        __UpperCAmelCase					       =      encoded_inputs.pixel_values.to(__lowerCamelCase	)

        with torch.no_grad():
            __UpperCAmelCase					       =      model(__lowerCamelCase	)

        __UpperCAmelCase					       =      torch.Size((1, model.config.num_labels, 128, 128)	)
        self.assertEqual(outputs.logits.shape  ,   __lowerCamelCase	)

        __UpperCAmelCase					       =      torch.tensor(
            [
                [[-13.5748, -13.9111, -12.6500], [-14.3500, -15.3683, -14.2328], [-14.7532, -16.0424, -15.6087]],
                [[-17.1651, -15.8725, -12.9653], [-17.2580, -17.3718, -14.8223], [-16.6058, -16.8783, -16.7452]],
                [[-3.6456, -3.0209, -1.4203], [-3.0797, -3.1959, -2.0000], [-1.8757, -1.9217, -1.6997]],
            ]	).to(__lowerCamelCase	)
        self.assertTrue(torch.allclose(outputs.logits[0, :3, :3, :3]  ,   __lowerCamelCase  ,   atol=1E-1	)	)






    @slow
    def   lowerCAmelCase_    (self	)  ->   str:
        __UpperCAmelCase					       =      SegformerImageProcessor(
            image_scale=(512, 512)  ,   keep_ratio=__lowerCamelCase  ,   align=__lowerCamelCase  ,   do_random_crop=__lowerCamelCase	)
        __UpperCAmelCase					       =      SegformerForSemanticSegmentation.from_pretrained('''nvidia/segformer-b0-finetuned-ade-512-512'''	).to(
            __lowerCamelCase	)

        __UpperCAmelCase					       =      prepare_img()
        __UpperCAmelCase					       =      image_processor(images=__lowerCamelCase  ,   return_tensors='''pt'''	)
        __UpperCAmelCase					       =      encoded_inputs.pixel_values.to(__lowerCamelCase	)

        with torch.no_grad():
            __UpperCAmelCase					       =      model(__lowerCamelCase	)

        __UpperCAmelCase					       =      outputs.logits.detach().cpu()

        __UpperCAmelCase					       =      image_processor.post_process_semantic_segmentation(outputs=__lowerCamelCase  ,   target_sizes=[(500, 300)]	)
        __UpperCAmelCase					       =      torch.Size((500, 300)	)
        self.assertEqual(segmentation[0].shape  ,   __lowerCamelCase	)

        __UpperCAmelCase					       =      image_processor.post_process_semantic_segmentation(outputs=__lowerCamelCase	)
        __UpperCAmelCase					       =      torch.Size((128, 128)	)
        self.assertEqual(segmentation[0].shape  ,   __lowerCamelCase	)








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

      UpperCAmelCase      					=  in_proj_weight[
          : encoder_config.hidden_size, :
      ]
      UpperCAmelCase      					=  in_proj_weight[
          encoder_config.hidden_size : encoder_config.hidden_size * 2, :
      ]
      UpperCAmelCase      					=  in_proj_weight[
          -encoder_config.hidden_size :, :
      ]
def 	_UpperCamelCase					(    lowerCAmelCase_	,    lowerCAmelCase_	,    lowerCAmelCase_ )    ->Optional[Any]:
   UpperCAmelCase      					=  dct.pop(lowerCAmelCase_ )
   UpperCAmelCase      					=  val
def 	_UpperCamelCase					(    lowerCAmelCase_ )    ->Any:
   if "handwritten" in checkpoint_url:
      UpperCAmelCase      					=  """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      					=  """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      					=  Image.open(requests.get(lowerCAmelCase_	,    stream=lowerCAmelCase_ ).raw ).convert("""RGB""" )
   return im
@torch.no_grad()
def 	_UpperCamelCase					(    lowerCAmelCase_	,    lowerCAmelCase_ )    ->Optional[int]:
   UpperCAmelCase      					=  ViTConfig(image_size=3_8_4	,    qkv_bias=lowerCAmelCase_ )
   UpperCAmelCase      					=  TrOCRConfig()

   # size of the architecture
   if "base" in checkpoint_url:
      UpperCAmelCase      					=  7_6_8
   elif "large" in checkpoint_url:
      # use ViT-large encoder
      UpperCAmelCase      					=  1_0_2_4
      UpperCAmelCase      					=  4_0_9_6
      UpperCAmelCase      					=  2_4
      UpperCAmelCase      					=  1_6
      UpperCAmelCase      					=  1_0_2_4
   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      					=  False
      UpperCAmelCase      					=  """relu"""
      UpperCAmelCase      					=  1_0_2_4
      UpperCAmelCase      					=  True
      UpperCAmelCase      					=  False
      UpperCAmelCase      					=  False

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

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

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

    # load state dict
   model.load_state_dict(lowerCAmelCase_ )

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

   UpperCAmelCase      					=  processor(images=prepare_img(lowerCAmelCase_ )	,    return_tensors="""pt""" ).pixel_values

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

   UpperCAmelCase      					=  torch.Size([1, 1, 5_0_2_6_5] )
   if "trocr-base-handwritten" in checkpoint_url:
      UpperCAmelCase      					=  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      					=  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      					=  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      					=  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, :1_0]	,    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__":
 __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)



