


'''simple docstring'''



# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pathlib import Path

import torch

from ...utils import is_npu_available, is_xpu_available
from .config_args import ClusterConfig, default_json_config_file
from .config_utils import SubcommandHelpFormatter


__A		='Create a default config file for Accelerate with only a few flags set.'







def       _UpperCamelCase (						UpperCamelCase__="no"	, UpperCamelCase__ = default_json_config_file	, UpperCamelCase__ = False   ):
			UpperCAmelCase__					:   Any       				=       Path(UpperCamelCase__   )
			path.parent.mkdir(parents=UpperCamelCase__	, exist_ok=UpperCamelCase__   )
			if path.exists():
						print(
						    f'''Configuration already exists at {save_location}, will not override. Run `accelerate config` manually or pass a different `save_location`.'''   )
						return False
			UpperCAmelCase__					:   Tuple       				=       mixed_precision.lower()
			if mixed_precision not in ["no", "fp16", "bf16", "fp8"]:
						raise ValueError(
						    f'''`mixed_precision` should be one of \'no\', \'fp16\', \'bf16\', or \'fp8\'. Received {mixed_precision}'''   )
			UpperCAmelCase__					:   Optional[int]       				=       {
			    """compute_environment""": """LOCAL_MACHINE""",
			    """mixed_precision""": mixed_precision,
			}
			if torch.cuda.is_available():
						UpperCAmelCase__					:   str       				=       torch.cuda.device_count()
						UpperCAmelCase__					:   int       				=       num_gpus
						UpperCAmelCase__					:   Union[str, Any]       				=       False
						if num_gpus > 1:
									UpperCAmelCase__					:   Union[str, Any]       				=       """MULTI_GPU"""
						else:
									UpperCAmelCase__					:   Dict       				=       """NO"""
			elif is_xpu_available() and use_xpu:
						UpperCAmelCase__					:   Any       				=       torch.xpu.device_count()
						UpperCAmelCase__					:   Dict       				=       num_xpus
						UpperCAmelCase__					:   Dict       				=       False
						if num_xpus > 1:
									UpperCAmelCase__					:   int       				=       """MULTI_XPU"""
						else:
									UpperCAmelCase__					:   Any       				=       """NO"""
			elif is_npu_available():
						UpperCAmelCase__					:   Dict       				=       torch.npu.device_count()
						UpperCAmelCase__					:   str       				=       num_npus
						UpperCAmelCase__					:   Tuple       				=       False
						if num_npus > 1:
									UpperCAmelCase__					:   Union[str, Any]       				=       """MULTI_NPU"""
						else:
									UpperCAmelCase__					:   Optional[Any]       				=       """NO"""
			else:
						UpperCAmelCase__					:   int       				=       0
						UpperCAmelCase__					:   Optional[int]       				=       True
						UpperCAmelCase__					:   List[str]       				=       1
						UpperCAmelCase__					:   Union[str, Any]       				=       """NO"""
			UpperCAmelCase__					:   Dict       				=       ClusterConfig(**UpperCamelCase__   )
			config.to_json_file(UpperCamelCase__   )
			return path







def       _UpperCamelCase (						UpperCamelCase__	, UpperCamelCase__   ):
			UpperCAmelCase__					:   Optional[int]       				=       parser.add_parser("""default"""	, parents=UpperCamelCase__	, help=UpperCamelCase__	, formatter_class=UpperCamelCase__   )
			parser.add_argument(
			    """--config_file"""	, default=UpperCamelCase__	, help=(
			        """The path to use to store the config file. Will default to a file named default_config.yaml in the cache """
			        """location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have """
			        """such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed """
			        """with 'huggingface'."""
			    )	, dest="""save_location"""	, )

			parser.add_argument(
			    """--mixed_precision"""	, choices=["""no""", """fp16""", """bf16"""]	, type=UpperCamelCase__	, help="""Whether or not to use mixed precision training. """
			    """Choose between FP16 and BF16 (bfloat16) training. """
			    """BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later."""	, default="""no"""	, )
			parser.set_defaults(func=UpperCamelCase__   )
			return parser







def       _UpperCamelCase (						UpperCamelCase__   ):
			UpperCAmelCase__					:   Any       				=       write_basic_config(args.mixed_precision	, args.save_location   )
			if config_file:
						print(f'''accelerate configuration saved at {config_file}'''   )



'''simple docstring'''



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()
__A		=logging.get_logger(__name__)







def       _UpperCamelCase (						UpperCamelCase__	, UpperCamelCase__=False	, UpperCamelCase__=False   ):
			UpperCAmelCase__					:   str       				=       """backbone.""" if is_semantic else """"""

			UpperCAmelCase__					:   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"""),
			        (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       _UpperCamelCase (						UpperCamelCase__	, UpperCamelCase__	, UpperCamelCase__=False	, UpperCamelCase__=False   ):
			for i in range(config.num_hidden_layers   ):
						UpperCAmelCase__					:   Optional[Any]       				=       """backbone.""" if is_semantic else """"""
						# queries, keys and values
						UpperCAmelCase__					:   Any       				=       state_dict.pop(f'''{prefix}blocks.{i}.attn.qkv.weight'''   )
						UpperCAmelCase__					:   List[str]       				=       state_dict.pop(f'''{prefix}blocks.{i}.attn.q_bias'''   )
						UpperCAmelCase__					:   int       				=       state_dict.pop(f'''{prefix}blocks.{i}.attn.v_bias'''   )

						UpperCAmelCase__					:   Optional[Any]       				=       in_proj_weight[
						    : config.hidden_size, :
						]
						UpperCAmelCase__					:   Any       				=       q_bias
						UpperCAmelCase__					:   Dict       				=       in_proj_weight[
						    config.hidden_size : config.hidden_size * 2, :
						]
						UpperCAmelCase__					:   Any       				=       in_proj_weight[
						    -config.hidden_size :, :
						]
						UpperCAmelCase__					:   Any       				=       v_bias

						# gamma_1 and gamma_2
						# we call them lambda because otherwise they are renamed when using .from_pretrained
						UpperCAmelCase__					:   Dict       				=       state_dict.pop(f'''{prefix}blocks.{i}.gamma_1'''   )
						UpperCAmelCase__					:   Dict       				=       state_dict.pop(f'''{prefix}blocks.{i}.gamma_2'''   )

						UpperCAmelCase__					:   Union[str, Any]       				=       gamma_a
						UpperCAmelCase__					:   str       				=       gamma_a







def       _UpperCamelCase (						UpperCamelCase__	, UpperCamelCase__	, UpperCamelCase__   ):
			UpperCAmelCase__					:   int       				=       dct.pop(UpperCamelCase__   )
			UpperCAmelCase__					:   Optional[Any]       				=       val







def       _UpperCamelCase (						):
			UpperCAmelCase__					:   Any       				=       """http://images.cocodataset.org/val2017/000000039769.jpg"""
			UpperCAmelCase__					:   Dict       				=       Image.open(requests.get(UpperCamelCase__	, stream=UpperCamelCase__   ).raw   )
			return im







@torch.no_grad()
def       _UpperCamelCase (						UpperCamelCase__	, UpperCamelCase__	, UpperCamelCase__=False   ):
			UpperCAmelCase__					:   Optional[Any]       				=       False if """rvlcdip""" in checkpoint_url else True
			UpperCAmelCase__					:   Any       				=       BeitConfig(use_absolute_position_embeddings=UpperCamelCase__	, use_mask_token=UpperCamelCase__   )

			# size of the architecture
			if "large" in checkpoint_url or "dit-l" in checkpoint_url:
						UpperCAmelCase__					:   Optional[Any]       				=       1_0_2_4
						UpperCAmelCase__					:   Dict       				=       4_0_9_6
						UpperCAmelCase__					:   Any       				=       2_4
						UpperCAmelCase__					:   Tuple       				=       1_6

			# labels
			if "rvlcdip" in checkpoint_url:
						UpperCAmelCase__					:   int       				=       1_6
						UpperCAmelCase__					:   List[str]       				=       """huggingface/label-files"""
						UpperCAmelCase__					:   Optional[Any]       				=       """rvlcdip-id2label.json"""
						UpperCAmelCase__					:   Optional[int]       				=       json.load(open(hf_hub_download(UpperCamelCase__	, UpperCamelCase__	, repo_type="""dataset"""   )	, """r"""   )   )
						UpperCAmelCase__					:   Union[str, Any]       				=       {int(UpperCamelCase__   ): v for k, v in idalabel.items()}
						UpperCAmelCase__					:   Optional[Any]       				=       idalabel
						UpperCAmelCase__					:   List[Any]       				=       {v: k for k, v in idalabel.items()}

			# load state_dict of original model, remove and rename some keys
			UpperCAmelCase__					:   Optional[Any]       				=       torch.hub.load_state_dict_from_url(UpperCamelCase__	, map_location="""cpu"""   )["""model"""]

			UpperCAmelCase__					:   List[str]       				=       create_rename_keys(UpperCamelCase__	, has_lm_head=UpperCamelCase__   )
			for src, dest in rename_keys:
						rename_key(UpperCamelCase__	, UpperCamelCase__	, UpperCamelCase__   )
			read_in_q_k_v(UpperCamelCase__	, UpperCamelCase__	, has_lm_head=UpperCamelCase__   )

			# load HuggingFace model
			UpperCAmelCase__					:   str       				=       BeitForMaskedImageModeling(UpperCamelCase__   ) if has_lm_head else BeitForImageClassification(UpperCamelCase__   )
			model.eval()
			model.load_state_dict(UpperCamelCase__   )

			# Check outputs on an image
			UpperCAmelCase__					:   List[str]       				=       BeitImageProcessor(
			    size=config.image_size	, resample=PILImageResampling.BILINEAR	, do_center_crop=UpperCamelCase__   )
			UpperCAmelCase__					:   List[str]       				=       prepare_img()

			UpperCAmelCase__					:   Optional[Any]       				=       image_processor(images=UpperCamelCase__	, return_tensors="""pt"""   )
			UpperCAmelCase__					:   Optional[Any]       				=       encoding["""pixel_values"""]

			UpperCAmelCase__					:   Optional[Any]       				=       model(UpperCamelCase__   )
			UpperCAmelCase__					:   int       				=       outputs.logits

			# verify logits
			UpperCAmelCase__					:   int       				=       [1, 1_6] if """rvlcdip""" in checkpoint_url else [1, 1_9_6, 8_1_9_2]
			assert logits.shape == torch.Size(UpperCamelCase__   ), "Shape 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 image processor to {pytorch_dump_folder_path}'''   )
			image_processor.save_pretrained(UpperCamelCase__   )

			if push_to_hub:
						if has_lm_head:
									UpperCAmelCase__					:   Union[str, Any]       				=       """dit-base""" if """base""" in checkpoint_url else """dit-large"""
						else:
									UpperCAmelCase__					:   Tuple       				=       """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(UpperCamelCase__	, UpperCamelCase__   )	, organization="""nielsr"""	, commit_message="""Add image processor"""	, use_temp_dir=UpperCamelCase__	, )
						model.push_to_hub(
						    repo_path_or_name=Path(UpperCamelCase__	, UpperCamelCase__   )	, organization="""nielsr"""	, commit_message="""Add model"""	, use_temp_dir=UpperCamelCase__	, )


if __name__ == "__main__":
		__A		=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',
		)
		__A		=parser.parse_args()
		convert_dit_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub)