

'''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()
__lowerCamelCase		       =    logging.get_logger(__name__)
def       UpperCAmelCase__    (      UpperCAmelCase__,       UpperCAmelCase__				)			->			List[str]:
       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       UpperCAmelCase__    (      UpperCAmelCase__,       UpperCAmelCase__				)			->			Union[str, Any]:
       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       UpperCAmelCase__    (      UpperCAmelCase__,       UpperCAmelCase__,       UpperCAmelCase__				)			->			List[Any]:
       A_						=      dct.pop(UpperCAmelCase__				)
       A_						=      val
def       UpperCAmelCase__    (      UpperCAmelCase__				)			->			List[Any]:
       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       UpperCAmelCase__    (      UpperCAmelCase__,       UpperCAmelCase__				)			->			List[Any]:
       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__":
      __lowerCamelCase		       =    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.'''
      )
      __lowerCamelCase		       =    parser.parse_args()
      convert_tr_ocr_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)



'''simple docstring'''







__lowerCamelCase		       =    range(2, 20 + 1)
__lowerCamelCase		       =    [10**k for k in range(ks[-1] + 1)]
__lowerCamelCase		       =    {}
def       UpperCAmelCase__    (      UpperCAmelCase__,       UpperCAmelCase__,       UpperCAmelCase__,       UpperCAmelCase__				)			->			Tuple:
       A_						=      sum(a_i[j] for j in range(UpperCAmelCase__,       len(UpperCAmelCase__				)				)				)
       A_						=      sum(a_i[j] * base[j] for j in range(min(len(UpperCAmelCase__				),       UpperCAmelCase__				)				)				)

       A_					,			A_						=      0, 0
       A_						=      n - i

       A_						=      memo.get(UpperCAmelCase__				)

       if sub_memo is not None:
              A_						=      sub_memo.get(UpperCAmelCase__				)

              if jumps is not None and len(UpperCAmelCase__				) > 0:
                     # find and make the largest jump without going over
                     A_						=      -1
                     for _k in range(len(UpperCAmelCase__				) - 1,       -1,       -1				):
                            if jumps[_k][2] <= k and jumps[_k][1] <= max_dn:
                                   A_						=      _k
                                   break

                     if max_jump >= 0:
                            A_					,			A_					,			A_						=      jumps[max_jump]
                            # since the difference between jumps is cached, add c
                            A_						=      diff + c
                            for j in range(min(UpperCAmelCase__,       len(UpperCAmelCase__				)				)				):
                                   A_					,			A_						=      divmod(UpperCAmelCase__,       10				)
                            if new_c > 0:
                                   add(UpperCAmelCase__,       UpperCAmelCase__,       UpperCAmelCase__				)

              else:
                     A_						=      []
       else:
              A_						=      {c: []}
              A_						=      sub_memo

       if dn >= max_dn or c + diff >= base[k]:
              return diff, dn

       if k > ks[0]:
              while True:
                     # keep doing smaller jumps
                     A_					,			A_						=      next_term(UpperCAmelCase__,       k - 1,       i + dn,       UpperCAmelCase__				)
                     diff += _diff
                     dn += terms_jumped

                     if dn >= max_dn or c + diff >= base[k]:
                            break
       else:
              # would be too small a jump, just compute sequential terms instead
              A_					,			A_						=      compute(UpperCAmelCase__,       UpperCAmelCase__,       i + dn,       UpperCAmelCase__				)
              diff += _diff
              dn += terms_jumped

       A_						=      sub_memo[c]

       # keep jumps sorted by # of terms skipped
       A_						=      0
       while j < len(UpperCAmelCase__				):
              if jumps[j][1] > dn:
                     break
              j += 1

       # cache the jump for this value digitsum(b) and c
       sub_memo[c].insert(UpperCAmelCase__,       (diff, dn, k)				)
       return (diff, dn)
def       UpperCAmelCase__    (      UpperCAmelCase__,       UpperCAmelCase__,       UpperCAmelCase__,       UpperCAmelCase__				)			->			int:
       if i >= n:
              return 0, i
       if k > len(UpperCAmelCase__				):
              a_i.extend([0 for _ in range(k - len(UpperCAmelCase__				)				)]				)

       # note: a_i -> b * 10^k + c
       # ds_b -> digitsum(b)
       # ds_c -> digitsum(c)
       A_						=      i
       A_					,			A_					,			A_						=      0, 0, 0
       for j in range(len(UpperCAmelCase__				)				):
              if j >= k:
                     ds_b += a_i[j]
              else:
                     ds_c += a_i[j]

       while i < n:
              i += 1
              A_						=      ds_c + ds_b
              diff += addend
              A_						=      0
              for j in range(UpperCAmelCase__				):
                     A_						=      a_i[j] + addend
                     A_					,			A_						=      divmod(UpperCAmelCase__,       10				)

                     ds_c += a_i[j]

              if addend > 0:
                     break

       if addend > 0:
              add(UpperCAmelCase__,       UpperCAmelCase__,       UpperCAmelCase__				)
       return diff, i - start_i
def       UpperCAmelCase__    (      UpperCAmelCase__,       UpperCAmelCase__,       UpperCAmelCase__				)			->			str:
       for j in range(UpperCAmelCase__,       len(UpperCAmelCase__				)				):
              A_						=      digits[j] + addend
              if s >= 10:
                     A_					,			A_						=      divmod(UpperCAmelCase__,       10				)
                     A_						=      addend // 10 + quotient
              else:
                     A_						=      s
                     A_						=      addend // 10

              if addend == 0:
                     break

       while addend > 0:
              A_					,			A_						=      divmod(UpperCAmelCase__,       10				)
              digits.append(UpperCAmelCase__				)
def       UpperCAmelCase__    (      UpperCAmelCase__ = 10**15				)			->			int:
       A_						=      [1]
       A_						=      1
       A_						=      0
       while True:
              A_					,			A_						=      next_term(UpperCAmelCase__,       20,       i + dn,       UpperCAmelCase__				)
              dn += terms_jumped
              if dn == n - i:
                     break

       A_						=      0
       for j in range(len(UpperCAmelCase__				)				):
              a_n += digits[j] * 10**j
       return a_n


if __name__ == "__main__":
      print(f"""{solution() = }""")
