

import argparse
import json
from pathlib import Path

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

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


logging.set_verbosity_info()
lowerCAmelCase__       = logging.get_logger(__name__)
def      _lowerCAmelCase(	__A   ,       __A=False   ,       __A=False			):
  UpperCAmelCase					      =  "backbone." if is_semantic else ""

  UpperCAmelCase					      =  []
  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"),
          (F"{prefix}pos_embed", "beit.embeddings.position_embeddings"),
      ]			)

  if has_lm_head:
    # mask token + layernorm
    rename_keys.extend(
        [
            ("mask_token", "beit.embeddings.mask_token"),
            ("norm.weight", "layernorm.weight"),
            ("norm.bias", "layernorm.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      _lowerCAmelCase(	__A   ,       __A   ,       __A=False   ,       __A=False			):




  for i in range(config.num_hidden_layers			):
    UpperCAmelCase					      =  "backbone." if is_semantic else ""
    # queries, keys and values
    UpperCAmelCase					      =  state_dict.pop(F"{prefix}blocks.{i}.attn.qkv.weight"			)
    UpperCAmelCase					      =  state_dict.pop(F"{prefix}blocks.{i}.attn.q_bias"			)
    UpperCAmelCase					      =  state_dict.pop(F"{prefix}blocks.{i}.attn.v_bias"			)

    UpperCAmelCase					      =  in_proj_weight[
        : config.hidden_size, :
    ]
    UpperCAmelCase					      =  q_bias
    UpperCAmelCase					      =  in_proj_weight[
        config.hidden_size : config.hidden_size * 2, :
    ]
    UpperCAmelCase					      =  in_proj_weight[
        -config.hidden_size :, :
    ]
    UpperCAmelCase					      =  v_bias

    # gamma_1 and gamma_2
    # we call them lambda because otherwise they are renamed when using .from_pretrained
    UpperCAmelCase					      =  state_dict.pop(F"{prefix}blocks.{i}.gamma_1"			)
    UpperCAmelCase					      =  state_dict.pop(F"{prefix}blocks.{i}.gamma_2"			)

    UpperCAmelCase					      =  gamma_a
    UpperCAmelCase					      =  gamma_a
def      _lowerCAmelCase(	__A   ,       __A   ,       __A			):
  UpperCAmelCase					      =  dct.pop(__A			)
  UpperCAmelCase					      =  val
def      _lowerCAmelCase(	):
  UpperCAmelCase					      =  "http://images.cocodataset.org/val2017/000000039769.jpg"
  UpperCAmelCase					      =  Image.open(requests.get(__A   ,       stream=__A			).raw			)
  return im
@torch.no_grad()
def      _lowerCAmelCase(	__A   ,       __A   ,       __A=False			):
  UpperCAmelCase					      =  False if "rvlcdip" in checkpoint_url else True
  UpperCAmelCase					      =  BeitConfig(use_absolute_position_embeddings=__A   ,       use_mask_token=__A			)

  # size of the architecture
  if "large" in checkpoint_url or "dit-l" in checkpoint_url:
    UpperCAmelCase					      =  1024
    UpperCAmelCase					      =  4096
    UpperCAmelCase					      =  24
    UpperCAmelCase					      =  16

  # labels
  if "rvlcdip" in checkpoint_url:
    UpperCAmelCase					      =  16
    UpperCAmelCase					      =  "huggingface/label-files"
    UpperCAmelCase					      =  "rvlcdip-id2label.json"
    UpperCAmelCase					      =  json.load(open(hf_hub_download(__A   ,       __A   ,       repo_type="dataset"			)   ,       "r"			)			)
    UpperCAmelCase					      =  {int(__A			): v for k, v in idalabel.items()}
    UpperCAmelCase					      =  idalabel
    UpperCAmelCase					      =  {v: k for k, v in idalabel.items()}

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

  UpperCAmelCase					      =  create_rename_keys(__A   ,       has_lm_head=__A			)
  for src, dest in rename_keys:
    rename_key(__A   ,       __A   ,       __A			)
  read_in_q_k_v(__A   ,       __A   ,       has_lm_head=__A			)

  # load HuggingFace model
  UpperCAmelCase					      =  BeitForMaskedImageModeling(__A			) if has_lm_head else BeitForImageClassification(__A			)
  model.eval()
  model.load_state_dict(__A			)

  # Check outputs on an image
  UpperCAmelCase					      =  BeitImageProcessor(
      size=config.image_size   ,       resample=PILImageResampling.BILINEAR   ,       do_center_crop=__A			)
  UpperCAmelCase					      =  prepare_img()

  UpperCAmelCase					      =  image_processor(images=__A   ,       return_tensors="pt"			)
  UpperCAmelCase					      =  encoding["pixel_values"]

  UpperCAmelCase					      =  model(__A			)
  UpperCAmelCase					      =  outputs.logits

  # verify logits
  UpperCAmelCase					      =  [1, 16] if "rvlcdip" in checkpoint_url else [1, 196, 8192]
  assert logits.shape == torch.Size(__A			), "Shape of logits not as expected"

  Path(__A			).mkdir(exist_ok=__A			)
  print(F"Saving model to {pytorch_dump_folder_path}"			)
  model.save_pretrained(__A			)
  print(F"Saving image processor to {pytorch_dump_folder_path}"			)
  image_processor.save_pretrained(__A			)

  if push_to_hub:
    if has_lm_head:
      UpperCAmelCase					      =  "dit-base" if "base" in checkpoint_url else "dit-large"
    else:
      UpperCAmelCase					      =  "dit-base-finetuned-rvlcdip" if "dit-b" in checkpoint_url else "dit-large-finetuned-rvlcdip"
    image_processor.push_to_hub(
        repo_path_or_name=Path(__A   ,       __A			)   ,       organization="nielsr"   ,       commit_message="Add image processor"   ,       use_temp_dir=__A   ,       )
    model.push_to_hub(
        repo_path_or_name=Path(__A   ,       __A			)   ,       organization="nielsr"   ,       commit_message="Add model"   ,       use_temp_dir=__A   ,       )


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

    parser.add_argument(
        "--checkpoint_url",
        default="https://layoutlm.blob.core.windows.net/dit/dit-pts/dit-base-224-p16-500k-62d53a.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."
    )
    parser.add_argument(
        "--push_to_hub",
        action="store_true",
    )
    lowerCAmelCase__       = parser.parse_args()
    convert_dit_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub)






import sys
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import AutoFeatureExtractor, WavaVecaFeatureExtractor
from transformers.testing_utils import TOKEN, USER, get_tests_dir, is_staging_test


sys.path.append(str(Path(__file__).parent.parent / "utils"))

from test_module.custom_feature_extraction import CustomFeatureExtractor  # noqa E402


lowerCAmelCase__       = get_tests_dir("fixtures")






class 						__magic_name__    (       unittest.TestCase    ):


 def _UpperCamelCase	(						self		: Tuple    )    ->      Union[str, Any]:
   # A mock response for an HTTP head request to emulate server down
   UpperCAmelCase					      =  mock.Mock()
   UpperCAmelCase					      =  5_0_0
   UpperCAmelCase					      =  {}
   UpperCAmelCase					      =  HTTPError
   UpperCAmelCase					      =  {}

   # Download this model to make sure it's in the cache.
   UpperCAmelCase					      =  WavaVecaFeatureExtractor.from_pretrained("hf-internal-testing/tiny-random-wav2vec2"    )
   # Under the mock environment we get a 500 error when trying to reach the model.
   with mock.patch("requests.Session.request"			,						return_value=lowerCAmelCase__    ) as mock_head:
     UpperCAmelCase					      =  WavaVecaFeatureExtractor.from_pretrained("hf-internal-testing/tiny-random-wav2vec2"    )
     # This check we did call the fake head request
     mock_head.assert_called()


 def _UpperCamelCase	(						self		: List[Any]    )    ->      Dict:
   # This test is for deprecated behavior and can be removed in v5
   UpperCAmelCase					      =  WavaVecaFeatureExtractor.from_pretrained(
       "https://huggingface.co/hf-internal-testing/tiny-random-wav2vec2/resolve/main/preprocessor_config.json"    )






@is_staging_test
class 						__magic_name__    (       unittest.TestCase    ):


 @classmethod
 def _UpperCamelCase	(						cls		: List[str]    )    ->      List[Any]:
   UpperCAmelCase					      =  TOKEN
   HfFolder.save_token(lowerCAmelCase__    )


 @classmethod
 def _UpperCamelCase	(						cls		: Optional[int]    )    ->      Union[str, Any]:
   try:
     delete_repo(token=cls._token			,						repo_id="test-feature-extractor"    )
   except HTTPError:
     pass

   try:
     delete_repo(token=cls._token			,						repo_id="valid_org/test-feature-extractor-org"    )
   except HTTPError:
     pass

   try:
     delete_repo(token=cls._token			,						repo_id="test-dynamic-feature-extractor"    )
   except HTTPError:
     pass


 def _UpperCamelCase	(						self		: Any    )    ->      Any:
   UpperCAmelCase					      =  WavaVecaFeatureExtractor.from_pretrained(lowerCAmelCase__    )
   feature_extractor.push_to_hub("test-feature-extractor"			,						use_auth_token=self._token    )

   UpperCAmelCase					      =  WavaVecaFeatureExtractor.from_pretrained(f"{USER}/test-feature-extractor"    )
   for k, v in feature_extractor.__dict__.items():
     self.assertEqual(lowerCAmelCase__			,						getattr(lowerCAmelCase__			,						lowerCAmelCase__    )    )

   # Reset repo
   delete_repo(token=self._token			,						repo_id="test-feature-extractor"    )

   # Push to hub via save_pretrained
   with tempfile.TemporaryDirectory() as tmp_dir:
     feature_extractor.save_pretrained(
         lowerCAmelCase__			,						repo_id="test-feature-extractor"			,						push_to_hub=lowerCAmelCase__			,						use_auth_token=self._token    )

   UpperCAmelCase					      =  WavaVecaFeatureExtractor.from_pretrained(f"{USER}/test-feature-extractor"    )
   for k, v in feature_extractor.__dict__.items():
     self.assertEqual(lowerCAmelCase__			,						getattr(lowerCAmelCase__			,						lowerCAmelCase__    )    )


 def _UpperCamelCase	(						self		: List[Any]    )    ->      Tuple:
   UpperCAmelCase					      =  WavaVecaFeatureExtractor.from_pretrained(lowerCAmelCase__    )
   feature_extractor.push_to_hub("valid_org/test-feature-extractor"			,						use_auth_token=self._token    )

   UpperCAmelCase					      =  WavaVecaFeatureExtractor.from_pretrained("valid_org/test-feature-extractor"    )
   for k, v in feature_extractor.__dict__.items():
     self.assertEqual(lowerCAmelCase__			,						getattr(lowerCAmelCase__			,						lowerCAmelCase__    )    )

   # Reset repo
   delete_repo(token=self._token			,						repo_id="valid_org/test-feature-extractor"    )

   # Push to hub via save_pretrained
   with tempfile.TemporaryDirectory() as tmp_dir:
     feature_extractor.save_pretrained(
         lowerCAmelCase__			,						repo_id="valid_org/test-feature-extractor-org"			,						push_to_hub=lowerCAmelCase__			,						use_auth_token=self._token    )

   UpperCAmelCase					      =  WavaVecaFeatureExtractor.from_pretrained("valid_org/test-feature-extractor-org"    )
   for k, v in feature_extractor.__dict__.items():
     self.assertEqual(lowerCAmelCase__			,						getattr(lowerCAmelCase__			,						lowerCAmelCase__    )    )


 def _UpperCamelCase	(						self		: Dict    )    ->      List[str]:
   CustomFeatureExtractor.register_for_auto_class()
   UpperCAmelCase					      =  CustomFeatureExtractor.from_pretrained(lowerCAmelCase__    )

   feature_extractor.push_to_hub("test-dynamic-feature-extractor"			,						use_auth_token=self._token    )

   # This has added the proper auto_map field to the config
   self.assertDictEqual(
       feature_extractor.auto_map			,						{"AutoFeatureExtractor": "custom_feature_extraction.CustomFeatureExtractor"}			,						)

   UpperCAmelCase					      =  AutoFeatureExtractor.from_pretrained(
       f"{USER}/test-dynamic-feature-extractor"			,						trust_remote_code=lowerCAmelCase__    )
   # Can't make an isinstance check because the new_feature_extractor is from the CustomFeatureExtractor class of a dynamic module
   self.assertEqual(new_feature_extractor.__class__.__name__			,						"CustomFeatureExtractor"    )



