


"""simple docstring"""





import unittest

from transformers import (
    MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
    TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
    SummarizationPipeline,
    TFPreTrainedModel,
    pipeline,
)
from transformers.testing_utils import get_gpu_count, is_pipeline_test, require_tf, require_torch, slow, torch_device
from transformers.tokenization_utils import TruncationStrategy

from .test_pipelines_common import ANY


__UpperCAmelCase    				=      -1 if torch_device == "cpu" else 0

@is_pipeline_test
class  lowercase_  ( unittest.TestCase						):
      __magic_name__     :							int          =				MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING
      __magic_name__     :							Any          =				TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING







      def 				_lowerCAmelCase       (      self :							List[str]		,  _lowercase :							int		,  _lowercase :							List[Any]		,  _lowercase :							Dict    ):
        lowerCAmelCase__    :    Union[str, Any]								=					SummarizationPipeline(model=_lowercase		,  tokenizer=_lowercase    )
        return summarizer, ["(CNN)The Palestinian Authority officially became", "Some other text"]







      def 				_lowerCAmelCase       (      self :							Dict		,  _lowercase :							Tuple		,  _lowercase :							Dict    ):
        lowerCAmelCase__    :    int								=					summarizer.model

        lowerCAmelCase__    :    int								=					summarizer("(CNN)The Palestinian Authority officially became"    )
        self.assertEqual(_lowercase		,  [{"summary_text": ANY(_lowercase    )}]    )

        lowerCAmelCase__    :    Any								=					summarizer(
            "(CNN)The Palestinian Authority officially became "		,  num_beams=2		,  min_length=2		,  max_length=5		,  )
        self.assertEqual(_lowercase		,  [{"summary_text": ANY(_lowercase    )}]    )

        # Some models (Switch Transformers, LED, T5, LongT5, etc) can handle long sequences.
        lowerCAmelCase__    :    Dict								=					[
            "SwitchTransformersConfig",
            "T5Config",
            "LongT5Config",
            "LEDConfig",
            "PegasusXConfig",
            "FSMTConfig",
            "M2M100Config",
            "ProphetNetConfig",  # positional embeddings up to a fixed maximum size (otherwise clamping the values)
        ]
        if model.config.__class__.__name__ not in model_can_handle_longer_seq:
          # Too long and exception is expected.
          # For TF models, if the weights are initialized in GPU context, we won't get expected index error from
          # the embedding layer.
          if not (
              isinstance(_lowercase		,  _lowercase    )
              and get_gpu_count() > 0
              and len(summarizer.model.trainable_weights    ) > 0
          ):
            with self.assertRaises(_lowercase    ):
              lowerCAmelCase__    :    Dict								=					summarizer("This " * 1_0_0_0    )
        lowerCAmelCase__    :    Any								=					summarizer("This " * 1_0_0_0		,  truncation=TruncationStrategy.ONLY_FIRST    )







      @require_torch
      def 				_lowerCAmelCase       (      self :							Tuple    ):
        lowerCAmelCase__    :    str								=					pipeline(task="summarization"		,  model="sshleifer/tiny-mbart"		,  framework="pt"    )
        lowerCAmelCase__    :    str								=					summarizer("This is a small test"    )
        self.assertEqual(
            _lowercase		,  [
                {
                    "summary_text": "เข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไป"
                }
            ]		,  )







      @require_tf
      def 				_lowerCAmelCase       (      self :							Optional[Any]    ):
        lowerCAmelCase__    :    List[Any]								=					pipeline(task="summarization"		,  model="sshleifer/tiny-mbart"		,  framework="tf"    )
        lowerCAmelCase__    :    Optional[int]								=					summarizer("This is a small test"    )
        self.assertEqual(
            _lowercase		,  [
                {
                    "summary_text": "เข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไปเข้าไป"
                }
            ]		,  )







      @require_torch
      @slow
      def 				_lowerCAmelCase       (      self :							Tuple    ):
        lowerCAmelCase__    :    List[str]								=					pipeline(task="summarization"		,  device=_lowercase    )
        lowerCAmelCase__    :    Dict								=					(
            " (CNN)The Palestinian Authority officially became the 123rd member of the International Criminal Court on"
            " Wednesday, a step that gives the court jurisdiction over alleged crimes in Palestinian territories. The"
            " formal accession was marked with a ceremony at The Hague, in the Netherlands, where the court is based."
            " The Palestinians signed the ICC's founding Rome Statute in January, when they also accepted its"
            " jurisdiction over alleged crimes committed \"in the occupied Palestinian territory, including East"
            " Jerusalem, since June 13, 2014.\" Later that month, the ICC opened a preliminary examination into the"
            " situation in Palestinian territories, paving the way for possible war crimes investigations against"
            " Israelis. As members of the court, Palestinians may be subject to counter-charges as well. Israel and"
            " the United States, neither of which is an ICC member, opposed the Palestinians' efforts to join the"
            " body. But Palestinian Foreign Minister Riad al-Malki, speaking at Wednesday's ceremony, said it was a"
            " move toward greater justice. \"As Palestine formally becomes a State Party to the Rome Statute today, the"
            " world is also a step closer to ending a long era of impunity and injustice,\" he said, according to an"
            " ICC news release. \"Indeed, today brings us closer to our shared goals of justice and peace.\" Judge"
            " Kuniko Ozaki, a vice president of the ICC, said acceding to the treaty was just the first step for the"
            " Palestinians. \"As the Rome Statute today enters into force for the State of Palestine, Palestine"
            " acquires all the rights as well as responsibilities that come with being a State Party to the Statute."
            " These are substantive commitments, which cannot be taken lightly,\" she said. Rights group Human Rights"
            " Watch welcomed the development. \"Governments seeking to penalize Palestine for joining the ICC should"
            " immediately end their pressure, and countries that support universal acceptance of the court's treaty"
            " should speak out to welcome its membership,\" said Balkees Jarrah, international justice counsel for the"
            " group. \"What's objectionable is the attempts to undermine international justice, not Palestine's"
            " decision to join a treaty to which over 100 countries around the world are members.\" In January, when"
            " the preliminary ICC examination was opened, Israeli Prime Minister Benjamin Netanyahu described it as an"
            " outrage, saying the court was overstepping its boundaries. The United States also said it \"strongly\""
            " disagreed with the court's decision. \"As we have said repeatedly, we do not believe that Palestine is a"
            " state and therefore we do not believe that it is eligible to join the ICC,\" the State Department said in"
            " a statement. It urged the warring sides to resolve their differences through direct negotiations. \"We"
            " will continue to oppose actions against Israel at the ICC as counterproductive to the cause of peace,\""
            " it said. But the ICC begs to differ with the definition of a state for its purposes and refers to the"
            " territories as \"Palestine.\" While a preliminary examination is not a formal investigation, it allows the"
            " court to review evidence and determine whether to investigate suspects on both sides. Prosecutor Fatou"
            " Bensouda said her office would \"conduct its analysis in full independence and impartiality.\" The war"
            " between Israel and Hamas militants in Gaza last summer left more than 2,000 people dead. The inquiry"
            " will include alleged war crimes committed since June. The International Criminal Court was set up in"
            " 2002 to prosecute genocide, crimes against humanity and war crimes. CNN's Vasco Cotovio, Kareem Khadder"
            " and Faith Karimi contributed to this report."
        )
        lowerCAmelCase__    :    List[str]								=					(
            " The Palestinian Authority becomes the 123rd member of the International Criminal Court . The move gives"
            " the court jurisdiction over alleged crimes in Palestinian territories . Israel and the United States"
            " opposed the Palestinians' efforts to join the court . Rights group Human Rights Watch welcomes the move,"
            " says governments seeking to penalize Palestine should end pressure ."
        )
        lowerCAmelCase__    :    List[Any]								=					summarizer(_lowercase    )
        self.assertEqual(result[0]["summary_text"]		,  _lowercase    )




"""simple docstring"""





import argparse
import json
from pathlib import Path

import requests
import torch
from datasets import load_dataset
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    BeitConfig,
    BeitForImageClassification,
    BeitForMaskedImageModeling,
    BeitForSemanticSegmentation,
    BeitImageProcessor,
)
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


logging.set_verbosity_info()
__UpperCAmelCase    				=      logging.get_logger(__name__)





def     lowercase__	(  lowerCamelCase     :  Optional[int]				,	lowerCamelCase     :  Union[str, Any]=False				,	lowerCamelCase     :  Optional[int]=False				)					->     Union[str, Any]:
  lowerCAmelCase__    :    str								=					"backbone." if is_semantic else ""

  lowerCAmelCase__    :    Optional[Any]								=					[]
  for i in range(config.num_hidden_layers				):
    # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
    rename_keys.append((F"{prefix}blocks.{i}.norm1.weight", F"beit.encoder.layer.{i}.layernorm_before.weight")				)
    rename_keys.append((F"{prefix}blocks.{i}.norm1.bias", F"beit.encoder.layer.{i}.layernorm_before.bias")				)
    rename_keys.append(
        (F"{prefix}blocks.{i}.attn.proj.weight", F"beit.encoder.layer.{i}.attention.output.dense.weight")				)
    rename_keys.append(
        (F"{prefix}blocks.{i}.attn.proj.bias", F"beit.encoder.layer.{i}.attention.output.dense.bias")				)
    rename_keys.append((F"{prefix}blocks.{i}.norm2.weight", F"beit.encoder.layer.{i}.layernorm_after.weight")				)
    rename_keys.append((F"{prefix}blocks.{i}.norm2.bias", F"beit.encoder.layer.{i}.layernorm_after.bias")				)
    rename_keys.append((F"{prefix}blocks.{i}.mlp.fc1.weight", F"beit.encoder.layer.{i}.intermediate.dense.weight")				)
    rename_keys.append((F"{prefix}blocks.{i}.mlp.fc1.bias", F"beit.encoder.layer.{i}.intermediate.dense.bias")				)
    rename_keys.append((F"{prefix}blocks.{i}.mlp.fc2.weight", F"beit.encoder.layer.{i}.output.dense.weight")				)
    rename_keys.append((F"{prefix}blocks.{i}.mlp.fc2.bias", F"beit.encoder.layer.{i}.output.dense.bias")				)

  # projection layer + position embeddings
  rename_keys.extend(
      [
          (F"{prefix}cls_token", "beit.embeddings.cls_token"),
          (F"{prefix}patch_embed.proj.weight", "beit.embeddings.patch_embeddings.projection.weight"),
          (F"{prefix}patch_embed.proj.bias", "beit.embeddings.patch_embeddings.projection.bias"),
      ]				)

  if has_lm_head:
    # mask token + shared relative position bias + layernorm
    rename_keys.extend(
        [
            ("mask_token", "beit.embeddings.mask_token"),
            (
                "rel_pos_bias.relative_position_bias_table",
                "beit.encoder.relative_position_bias.relative_position_bias_table",
            ),
            (
                "rel_pos_bias.relative_position_index",
                "beit.encoder.relative_position_bias.relative_position_index",
            ),
            ("norm.weight", "layernorm.weight"),
            ("norm.bias", "layernorm.bias"),
        ]				)
  elif is_semantic:
    # semantic segmentation classification heads
    rename_keys.extend(
        [
            ("decode_head.conv_seg.weight", "decode_head.classifier.weight"),
            ("decode_head.conv_seg.bias", "decode_head.classifier.bias"),
            ("auxiliary_head.conv_seg.weight", "auxiliary_head.classifier.weight"),
            ("auxiliary_head.conv_seg.bias", "auxiliary_head.classifier.bias"),
        ]				)
  else:
    # layernorm + classification head
    rename_keys.extend(
        [
            ("fc_norm.weight", "beit.pooler.layernorm.weight"),
            ("fc_norm.bias", "beit.pooler.layernorm.bias"),
            ("head.weight", "classifier.weight"),
            ("head.bias", "classifier.bias"),
        ]				)

  return rename_keys





def     lowercase__	(  lowerCamelCase     :  List[str]				,	lowerCamelCase     :  Any				,	lowerCamelCase     :  Dict=False				,	lowerCamelCase     :  Optional[int]=False				)					->     Tuple:
  for i in range(config.num_hidden_layers				):
    lowerCAmelCase__    :    Dict								=					"backbone." if is_semantic else ""
    # queries, keys and values
    lowerCAmelCase__    :    str								=					state_dict.pop(F"{prefix}blocks.{i}.attn.qkv.weight"				)
    lowerCAmelCase__    :    str								=					state_dict.pop(F"{prefix}blocks.{i}.attn.q_bias"				)
    lowerCAmelCase__    :    Optional[Any]								=					state_dict.pop(F"{prefix}blocks.{i}.attn.v_bias"				)

    lowerCAmelCase__    :    Union[str, Any]								=					in_proj_weight[
        : config.hidden_size, :
    ]
    lowerCAmelCase__    :    Dict								=					q_bias
    lowerCAmelCase__    :    str								=					in_proj_weight[
        config.hidden_size : config.hidden_size * 2, :
    ]
    lowerCAmelCase__    :    Optional[int]								=					in_proj_weight[
        -config.hidden_size :, :
    ]
    lowerCAmelCase__    :    Optional[Any]								=					v_bias

    # gamma_1 and gamma_2
    # we call them lambda because otherwise they are renamed when using .from_pretrained
    lowerCAmelCase__    :    Any								=					state_dict.pop(F"{prefix}blocks.{i}.gamma_1"				)
    lowerCAmelCase__    :    Optional[int]								=					state_dict.pop(F"{prefix}blocks.{i}.gamma_2"				)

    lowerCAmelCase__    :    Tuple								=					gamma_a
    lowerCAmelCase__    :    Any								=					gamma_a

    # relative_position bias table + index
    if not has_lm_head:
      # each layer has its own relative position bias
      lowerCAmelCase__    :    int								=					state_dict.pop(F"{prefix}blocks.{i}.attn.relative_position_bias_table"				)
      lowerCAmelCase__    :    Optional[int]								=					state_dict.pop(F"{prefix}blocks.{i}.attn.relative_position_index"				)

      lowerCAmelCase__    :    Any								=					table
      lowerCAmelCase__    :    int								=					index





def     lowercase__	(  lowerCamelCase     :  int				,	lowerCamelCase     :  Any				,	lowerCamelCase     :  Any				)					->     Dict:
  lowerCAmelCase__    :    str								=					dct.pop(lowerCamelCase				)
  lowerCAmelCase__    :    Union[str, Any]								=					val





def     lowercase__	(  )					->     Optional[int]:
  lowerCAmelCase__    :    Union[str, Any]								=					"http://images.cocodataset.org/val2017/000000039769.jpg"
  lowerCAmelCase__    :    List[str]								=					Image.open(requests.get(lowerCamelCase				,	stream=lowerCamelCase				).raw				)
  return im





@torch.no_grad()
def     lowercase__	(  lowerCamelCase     :  Optional[int]				,	lowerCamelCase     :  Optional[int]				)					->     int:
  lowerCAmelCase__    :    str								=					BeitConfig()
  lowerCAmelCase__    :    Any								=					False
  lowerCAmelCase__    :    Optional[Any]								=					False
  lowerCAmelCase__    :    Optional[Any]								=					"huggingface/label-files"
  # set config parameters based on URL
  if checkpoint_url[-9:-4] == "pt22k":
    # masked image modeling
    lowerCAmelCase__    :    int								=					True
    lowerCAmelCase__    :    List[str]								=					True
    lowerCAmelCase__    :    Union[str, Any]								=					True
  elif checkpoint_url[-9:-4] == "ft22k":
    # intermediate fine-tuning on ImageNet-22k
    lowerCAmelCase__    :    List[Any]								=					True
    lowerCAmelCase__    :    Dict								=					2_1_8_4_1
    lowerCAmelCase__    :    str								=					"imagenet-22k-id2label.json"
    lowerCAmelCase__    :    Union[str, Any]								=					json.load(open(hf_hub_download(lowerCamelCase				,	lowerCamelCase				,	repo_type="dataset"				)				,	"r"				)				)
    lowerCAmelCase__    :    Tuple								=					{int(lowerCamelCase				): v for k, v in idalabel.items()}
    # this dataset contains 21843 labels but the model only has 21841
    # we delete the classes as mentioned in https://github.com/google-research/big_transfer/issues/18
    del idalabel[9_2_0_5]
    del idalabel[1_5_0_2_7]
    lowerCAmelCase__    :    str								=					idalabel
    lowerCAmelCase__    :    str								=					{v: k for k, v in idalabel.items()}
  elif checkpoint_url[-8:-4] == "to1k":
    # fine-tuning on ImageNet-1k
    lowerCAmelCase__    :    int								=					True
    lowerCAmelCase__    :    Optional[int]								=					1_0_0_0
    lowerCAmelCase__    :    List[Any]								=					"imagenet-1k-id2label.json"
    lowerCAmelCase__    :    List[str]								=					json.load(open(hf_hub_download(lowerCamelCase				,	lowerCamelCase				,	repo_type="dataset"				)				,	"r"				)				)
    lowerCAmelCase__    :    Optional[Any]								=					{int(lowerCamelCase				): v for k, v in idalabel.items()}
    lowerCAmelCase__    :    List[Any]								=					idalabel
    lowerCAmelCase__    :    Union[str, Any]								=					{v: k for k, v in idalabel.items()}
    if "384" in checkpoint_url:
      lowerCAmelCase__    :    str								=					3_8_4
    if "512" in checkpoint_url:
      lowerCAmelCase__    :    Union[str, Any]								=					5_1_2
  elif "ade20k" in checkpoint_url:
    # fine-tuning
    lowerCAmelCase__    :    int								=					True
    lowerCAmelCase__    :    Tuple								=					1_5_0
    lowerCAmelCase__    :    Dict								=					"ade20k-id2label.json"
    lowerCAmelCase__    :    Union[str, Any]								=					json.load(open(hf_hub_download(lowerCamelCase				,	lowerCamelCase				,	repo_type="dataset"				)				,	"r"				)				)
    lowerCAmelCase__    :    str								=					{int(lowerCamelCase				): v for k, v in idalabel.items()}
    lowerCAmelCase__    :    List[Any]								=					idalabel
    lowerCAmelCase__    :    Any								=					{v: k for k, v in idalabel.items()}
    lowerCAmelCase__    :    Any								=					6_4_0
    lowerCAmelCase__    :    Dict								=					True
  else:
    raise ValueError("Checkpoint not supported, URL should either end with 'pt22k', 'ft22k', 'to1k' or 'ade20k'"				)

  # size of the architecture
  if "base" in checkpoint_url:
    pass
  elif "large" in checkpoint_url:
    lowerCAmelCase__    :    Optional[int]								=					1_0_2_4
    lowerCAmelCase__    :    int								=					4_0_9_6
    lowerCAmelCase__    :    Optional[Any]								=					2_4
    lowerCAmelCase__    :    Optional[int]								=					1_6
    if "ade20k" in checkpoint_url:
      lowerCAmelCase__    :    Optional[int]								=					6_4_0
      lowerCAmelCase__    :    Tuple								=					[7, 1_1, 1_5, 2_3]
  else:
    raise ValueError("Should either find 'base' or 'large' in checkpoint URL"				)

  # load state_dict of original model, remove and rename some keys
  lowerCAmelCase__    :    List[str]								=					torch.hub.load_state_dict_from_url(lowerCamelCase				,	map_location="cpu"				,	check_hash=lowerCamelCase				)
  lowerCAmelCase__    :    Tuple								=					state_dict["model"] if "ade20k" not in checkpoint_url else state_dict["state_dict"]

  lowerCAmelCase__    :    Optional[Any]								=					create_rename_keys(lowerCamelCase				,	has_lm_head=lowerCamelCase				,	is_semantic=lowerCamelCase				)
  for src, dest in rename_keys:
    rename_key(lowerCamelCase				,	lowerCamelCase				,	lowerCamelCase				)
  read_in_q_k_v(lowerCamelCase				,	lowerCamelCase				,	has_lm_head=lowerCamelCase				,	is_semantic=lowerCamelCase				)
  if is_semantic:
    # add prefix to decoder keys
    for key, val in state_dict.copy().items():
      lowerCAmelCase__    :    Dict								=					state_dict.pop(lowerCamelCase				)
      if key.startswith("backbone.fpn"				):
        lowerCAmelCase__    :    Any								=					key.replace("backbone.fpn"				,	"fpn"				)
      lowerCAmelCase__    :    Any								=					val

    # load HuggingFace model
  if checkpoint_url[-9:-4] == "pt22k":
    lowerCAmelCase__    :    Union[str, Any]								=					BeitForMaskedImageModeling(lowerCamelCase				)
  elif "ade20k" in checkpoint_url:
    lowerCAmelCase__    :    List[str]								=					BeitForSemanticSegmentation(lowerCamelCase				)
  else:
    lowerCAmelCase__    :    int								=					BeitForImageClassification(lowerCamelCase				)
  model.eval()
  model.load_state_dict(lowerCamelCase				)

  # Check outputs on an image
  if is_semantic:
    lowerCAmelCase__    :    Optional[Any]								=					BeitImageProcessor(size=config.image_size				,	do_center_crop=lowerCamelCase				)
    lowerCAmelCase__    :    str								=					load_dataset("hf-internal-testing/fixtures_ade20k"				,	split="test"				)
    lowerCAmelCase__    :    Union[str, Any]								=					Image.open(ds[0]["file"]				)
  else:
    lowerCAmelCase__    :    Optional[int]								=					BeitImageProcessor(
        size=config.image_size				,	resample=PILImageResampling.BILINEAR				,	do_center_crop=lowerCamelCase				)
    lowerCAmelCase__    :    List[Any]								=					prepare_img()

  lowerCAmelCase__    :    Optional[int]								=					image_processor(images=lowerCamelCase				,	return_tensors="pt"				)
  lowerCAmelCase__    :    Union[str, Any]								=					encoding["pixel_values"]

  lowerCAmelCase__    :    Optional[Any]								=					model(lowerCamelCase				)
  lowerCAmelCase__    :    Optional[int]								=					outputs.logits

  # verify logits
  lowerCAmelCase__    :    Union[str, Any]								=					torch.Size([1, 1_0_0_0]				)
  if checkpoint_url[:-4].endswith("beit_base_patch16_224_pt22k"				):
    lowerCAmelCase__    :    Any								=					torch.Size([1, 1_9_6, 8_1_9_2]				)
  elif checkpoint_url[:-4].endswith("beit_large_patch16_224_pt22k"				):
    lowerCAmelCase__    :    Any								=					torch.Size([1, 1_9_6, 8_1_9_2]				)
  elif checkpoint_url[:-4].endswith("beit_base_patch16_224_pt22k_ft22k"				):
    lowerCAmelCase__    :    Dict								=					torch.Size([1, 2_1_8_4_1]				)
    lowerCAmelCase__    :    Optional[int]								=					torch.tensor([2.22_88, 2.46_71, 0.73_95]				)
    lowerCAmelCase__    :    Any								=					2_3_9_7
  elif checkpoint_url[:-4].endswith("beit_large_patch16_224_pt22k_ft22k"				):
    lowerCAmelCase__    :    Tuple								=					torch.Size([1, 2_1_8_4_1]				)
    lowerCAmelCase__    :    List[str]								=					torch.tensor([1.68_81, -0.27_87, 0.59_01]				)
    lowerCAmelCase__    :    int								=					2_3_9_6
  elif checkpoint_url[:-4].endswith("beit_base_patch16_224_pt22k_ft1k"				):
    lowerCAmelCase__    :    Tuple								=					torch.tensor([0.12_41, 0.07_98, -0.65_69]				)
    lowerCAmelCase__    :    Tuple								=					2_8_5
  elif checkpoint_url[:-4].endswith("beit_base_patch16_224_pt22k_ft22kto1k"				):
    lowerCAmelCase__    :    List[Any]								=					torch.tensor([-1.23_85, -1.09_87, -1.01_08]				)
    lowerCAmelCase__    :    Union[str, Any]								=					2_8_1
  elif checkpoint_url[:-4].endswith("beit_base_patch16_384_pt22k_ft22kto1k"				):
    lowerCAmelCase__    :    Optional[Any]								=					torch.tensor([-1.53_03, -0.94_84, -0.31_47]				)
    lowerCAmelCase__    :    Optional[Any]								=					7_6_1
  elif checkpoint_url[:-4].endswith("beit_large_patch16_224_pt22k_ft1k"				):
    lowerCAmelCase__    :    Union[str, Any]								=					torch.tensor([0.46_10, -0.09_28, 0.20_86]				)
    lowerCAmelCase__    :    Any								=					7_6_1
  elif checkpoint_url[:-4].endswith("beit_large_patch16_224_pt22k_ft22kto1k"				):
    lowerCAmelCase__    :    str								=					torch.tensor([-0.48_04, 0.62_57, -0.18_37]				)
    lowerCAmelCase__    :    Any								=					7_6_1
  elif checkpoint_url[:-4].endswith("beit_large_patch16_384_pt22k_ft22kto1k"				):
    lowerCAmelCase__    :    int								=					torch.tensor([[-0.51_22, 0.51_17, -0.21_13]]				)
    lowerCAmelCase__    :    Optional[int]								=					7_6_1
  elif checkpoint_url[:-4].endswith("beit_large_patch16_512_pt22k_ft22kto1k"				):
    lowerCAmelCase__    :    Union[str, Any]								=					torch.tensor([-0.30_62, 0.72_61, 0.48_52]				)
    lowerCAmelCase__    :    Optional[Any]								=					7_6_1
  elif checkpoint_url[:-4].endswith("beit_base_patch16_640_pt22k_ft22ktoade20k"				):
    lowerCAmelCase__    :    Union[str, Any]								=					(1, 1_5_0, 1_6_0, 1_6_0)
    lowerCAmelCase__    :    Union[str, Any]								=					torch.tensor(
        [
            [[-4.92_25, -2.39_54, -3.05_22], [-2.88_22, -1.00_46, -1.75_61], [-2.95_49, -1.32_28, -2.13_47]],
            [[-5.81_68, -3.41_29, -4.07_78], [-3.86_51, -2.22_14, -3.02_77], [-3.83_56, -2.46_43, -3.35_35]],
            [[-0.00_78, 3.99_52, 4.07_54], [2.98_56, 4.69_44, 5.00_35], [3.24_13, 4.78_13, 4.99_69]],
        ]				)
  elif checkpoint_url[:-4].endswith("beit_large_patch16_640_pt22k_ft22ktoade20k"				):
    lowerCAmelCase__    :    Dict								=					(1, 1_5_0, 1_6_0, 1_6_0)
    lowerCAmelCase__    :    List[Any]								=					torch.tensor(
        [
            [[-4.33_05, -2.30_49, -3.01_61], [-2.95_91, -1.53_05, -2.22_51], [-3.41_98, -1.80_04, -2.90_62]],
            [[-5.89_22, -3.74_35, -4.39_78], [-4.20_63, -2.78_72, -3.47_55], [-4.27_91, -3.18_74, -4.16_81]],
            [[0.98_95, 4.34_67, 4.76_63], [4.24_76, 5.68_30, 6.15_18], [4.55_50, 6.24_95, 6.51_54]],
        ]				)
  else:
    raise ValueError("Can't verify logits as model is not supported"				)

  assert logits.shape == expected_shape, "Shape of logits not as expected"
  if not has_lm_head:
    if is_semantic:
      assert torch.allclose(
          logits[0, :3, :3, :3]				,	lowerCamelCase				,	atol=1e-3				), "First elements of logits not as expected"
    else:
      print("Predicted class idx:"				,	logits.argmax(-1				).item()				)
      assert torch.allclose(
          logits[0, :3]				,	lowerCamelCase				,	atol=1e-3				), "First elements of logits not as expected"
      assert logits.argmax(-1				).item() == expected_class_idx, "Predicted class index 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 image processor to {pytorch_dump_folder_path}"				)
  image_processor.save_pretrained(lowerCamelCase				)


if __name__ == "__main__":
   __UpperCAmelCase    				=      argparse.ArgumentParser()

   parser.add_argument(
       "--checkpoint_url",
       default="https://conversationhub.blob.core.windows.net/beit-share-public/beit/beit_base_patch16_224_pt22k_ft22kto1k.pth",
       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_beit_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)
