






"""simple docstring"""





import datasets

from .evaluate import (
    apply_no_ans_threshold,
    find_all_best_thresh,
    get_raw_scores,
    make_eval_dict,
    make_qid_to_has_ans,
    merge_eval,
)


__UpperCAmelCase    				=      "\\n@inproceedings{Rajpurkar2016SQuAD10,\n  title={SQuAD: 100, 000+ Questions for Machine Comprehension of Text},\n  author={Pranav Rajpurkar and Jian Zhang and Konstantin Lopyrev and Percy Liang},\n  booktitle={EMNLP},\n  year={2016}\n}\n"

__UpperCAmelCase    				=      "\nThis metric wrap the official scoring script for version 2 of the Stanford Question\nAnswering Dataset (SQuAD).\n\nStanford Question Answering Dataset (SQuAD) is a reading comprehension dataset, consisting of questions posed by\ncrowdworkers on a set of Wikipedia articles, where the answer to every question is a segment of text, or span,\nfrom the corresponding reading passage, or the question might be unanswerable.\n\nSQuAD2.0 combines the 100,000 questions in SQuAD1.1 with over 50,000 unanswerable questions\nwritten adversarially by crowdworkers to look similar to answerable ones.\nTo do well on SQuAD2.0, systems must not only answer questions when possible, but also\ndetermine when no answer is supported by the paragraph and abstain from answering.\n"

__UpperCAmelCase    				=      "\nComputes SQuAD v2 scores (F1 and EM).\nArgs:\n    predictions: List of triple for question-answers to score with the following elements:\n        - the question-answer 'id' field as given in the references (see below)\n        - the text of the answer\n        - the probability that the question has no answer\n    references: List of question-answers dictionaries with the following key-values:\n            - 'id': id of the question-answer pair (see above),\n            - 'answers': a list of Dict {'text': text of the answer as a string}\n    no_answer_threshold: float\n        Probability threshold to decide that a question has no answer.\nReturns:\n    'exact': Exact match (the normalized answer exactly match the gold answer)\n    'f1': The F-score of predicted tokens versus the gold answer\n    'total': Number of score considered\n    'HasAns_exact': Exact match (the normalized answer exactly match the gold answer)\n    'HasAns_f1': The F-score of predicted tokens versus the gold answer\n    'HasAns_total': Number of score considered\n    'NoAns_exact': Exact match (the normalized answer exactly match the gold answer)\n    'NoAns_f1': The F-score of predicted tokens versus the gold answer\n    'NoAns_total': Number of score considered\n    'best_exact': Best exact match (with varying threshold)\n    'best_exact_thresh': No-answer probability threshold associated to the best exact match\n    'best_f1': Best F1 (with varying threshold)\n    'best_f1_thresh': No-answer probability threshold associated to the best F1\nExamples:\n\n    >>> predictions = [{'prediction_text': '1976', 'id': '56e10a3be3433e1400422b22', 'no_answer_probability': 0.}]\n    >>> references = [{'answers': {'answer_start': [97], 'text': ['1976']}, 'id': '56e10a3be3433e1400422b22'}]\n    >>> squad_v2_metric = datasets.load_metric(\"squad_v2\")\n    >>> results = squad_v2_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {'exact': 100.0, 'f1': 100.0, 'total': 1, 'HasAns_exact': 100.0, 'HasAns_f1': 100.0, 'HasAns_total': 1, 'best_exact': 100.0, 'best_exact_thresh': 0.0, 'best_f1': 100.0, 'best_f1_thresh': 0.0}\n"

@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION							,							_KWARGS_DESCRIPTION						)
class  lowercase_  ( datasets.Metric						):







      def 				_lowerCAmelCase       (      self :							Optional[Any]    ):
        return datasets.MetricInfo(
            description=_DESCRIPTION		,  citation=_CITATION		,  inputs_description=_KWARGS_DESCRIPTION		,  features=datasets.Features(
                {
                    "predictions": {
                        "id": datasets.Value("string"    ),
                        "prediction_text": datasets.Value("string"    ),
                        "no_answer_probability": datasets.Value("float32"    ),
                    },
                    "references": {
                        "id": datasets.Value("string"    ),
                        "answers": datasets.features.Sequence(
                            {"text": datasets.Value("string"    ), "answer_start": datasets.Value("int32"    )}    ),
                    },
                }    )		,  codebase_urls=["https://rajpurkar.github.io/SQuAD-explorer/"]		,  reference_urls=["https://rajpurkar.github.io/SQuAD-explorer/"]		,  )







      def 				_lowerCAmelCase       (      self :							Any		,  _lowercase :							List[str]		,  _lowercase :							Union[str, Any]		,  _lowercase :							Optional[Any]=1.0    ):
        lowerCAmelCase__    :    str								=					{p["id"]: p["no_answer_probability"] for p in predictions}
        lowerCAmelCase__    :    Optional[int]								=					[{"paragraphs": [{"qas": references}]}]
        lowerCAmelCase__    :    Union[str, Any]								=					{p["id"]: p["prediction_text"] for p in predictions}

        lowerCAmelCase__    :    Union[str, Any]								=					make_qid_to_has_ans(_lowercase    )  # maps qid to True/False
        lowerCAmelCase__    :    str								=					[k for k, v in qid_to_has_ans.items() if v]
        lowerCAmelCase__    :    Union[str, Any]								=					[k for k, v in qid_to_has_ans.items() if not v]

        lowerCAmelCase__    :    List[Any]								=					get_raw_scores(_lowercase		,  _lowercase    )
        lowerCAmelCase__    :    List[Any]								=					apply_no_ans_threshold(_lowercase		,  _lowercase		,  _lowercase		,  _lowercase    )
        lowerCAmelCase__    :    Union[str, Any]								=					apply_no_ans_threshold(_lowercase		,  _lowercase		,  _lowercase		,  _lowercase    )
        lowerCAmelCase__    :    Tuple								=					make_eval_dict(_lowercase		,  _lowercase    )

        if has_ans_qids:
          lowerCAmelCase__    :    List[str]								=					make_eval_dict(_lowercase		,  _lowercase		,  qid_list=_lowercase    )
          merge_eval(_lowercase		,  _lowercase		,  "HasAns"    )
        if no_ans_qids:
          lowerCAmelCase__    :    Optional[int]								=					make_eval_dict(_lowercase		,  _lowercase		,  qid_list=_lowercase    )
          merge_eval(_lowercase		,  _lowercase		,  "NoAns"    )
        find_all_best_thresh(_lowercase		,  _lowercase		,  _lowercase		,  _lowercase		,  _lowercase		,  _lowercase    )
        return dict(_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)
