






'''simple docstring'''




import argparse
import json
from pathlib import Path

import requests
import timm
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from timm.data import resolve_data_config
from timm.data.transforms_factory import create_transform

from transformers import (
    BitConfig,
    ViTHybridConfig,
    ViTHybridForImageClassification,
    ViTHybridImageProcessor,
    ViTHybridModel,
)
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


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







def    __UpperCAmelCase    (	lowerCamelCase_  ,   lowerCamelCase_=False)				->  List[Any]:
		UpperCamelCase__					:       Any       =       []

		# fmt: off
		# stem:
		rename_keys.append(('cls_token', 'vit.embeddings.cls_token'))
		rename_keys.append(('pos_embed', 'vit.embeddings.position_embeddings'))

		rename_keys.append(('patch_embed.proj.weight', 'vit.embeddings.patch_embeddings.projection.weight'))
		rename_keys.append(('patch_embed.proj.bias', 'vit.embeddings.patch_embeddings.projection.bias'))

		# backbone
		rename_keys.append(('patch_embed.backbone.stem.conv.weight', 'vit.embeddings.patch_embeddings.backbone.bit.embedder.convolution.weight'))
		rename_keys.append(('patch_embed.backbone.stem.norm.weight', 'vit.embeddings.patch_embeddings.backbone.bit.embedder.norm.weight'))
		rename_keys.append(('patch_embed.backbone.stem.norm.bias', 'vit.embeddings.patch_embeddings.backbone.bit.embedder.norm.bias'))

		for stage_idx in range(len(config.backbone_config.depths)):
				for layer_idx in range(config.backbone_config.depths[stage_idx]):
						rename_keys.append((f'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.conv1.weight', f'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.conv1.weight'))
						rename_keys.append((f'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm1.weight', f'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm1.weight'))
						rename_keys.append((f'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm1.bias', f'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm1.bias'))
						rename_keys.append((f'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.conv2.weight', f'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.conv2.weight'))
						rename_keys.append((f'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm2.weight', f'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm2.weight'))
						rename_keys.append((f'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm2.bias', f'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm2.bias'))
						rename_keys.append((f'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.conv3.weight', f'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.conv3.weight'))
						rename_keys.append((f'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm3.weight', f'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm3.weight'))
						rename_keys.append((f'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm3.bias', f'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm3.bias'))

				rename_keys.append((f'patch_embed.backbone.stages.{stage_idx}.blocks.0.downsample.conv.weight', f'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.0.downsample.conv.weight'))
				rename_keys.append((f'patch_embed.backbone.stages.{stage_idx}.blocks.0.downsample.norm.weight', f'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.0.downsample.norm.weight'))
				rename_keys.append((f'patch_embed.backbone.stages.{stage_idx}.blocks.0.downsample.norm.bias', f'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.0.downsample.norm.bias'))

		# transformer encoder
		for i in range(config.num_hidden_layers):
				# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
				rename_keys.append((f'blocks.{i}.norm1.weight', f'vit.encoder.layer.{i}.layernorm_before.weight'))
				rename_keys.append((f'blocks.{i}.norm1.bias', f'vit.encoder.layer.{i}.layernorm_before.bias'))
				rename_keys.append((f'blocks.{i}.attn.proj.weight', f'vit.encoder.layer.{i}.attention.output.dense.weight'))
				rename_keys.append((f'blocks.{i}.attn.proj.bias', f'vit.encoder.layer.{i}.attention.output.dense.bias'))
				rename_keys.append((f'blocks.{i}.norm2.weight', f'vit.encoder.layer.{i}.layernorm_after.weight'))
				rename_keys.append((f'blocks.{i}.norm2.bias', f'vit.encoder.layer.{i}.layernorm_after.bias'))
				rename_keys.append((f'blocks.{i}.mlp.fc1.weight', f'vit.encoder.layer.{i}.intermediate.dense.weight'))
				rename_keys.append((f'blocks.{i}.mlp.fc1.bias', f'vit.encoder.layer.{i}.intermediate.dense.bias'))
				rename_keys.append((f'blocks.{i}.mlp.fc2.weight', f'vit.encoder.layer.{i}.output.dense.weight'))
				rename_keys.append((f'blocks.{i}.mlp.fc2.bias', f'vit.encoder.layer.{i}.output.dense.bias'))

		if base_model:
				# layernorm + pooler
				rename_keys.extend(
				    [
				        ('norm.weight', 'layernorm.weight'),
				        ('norm.bias', 'layernorm.bias'),
				        ('pre_logits.fc.weight', 'pooler.dense.weight'),
				        ('pre_logits.fc.bias', 'pooler.dense.bias'),
				    ])

				# if just the base model, we should remove "vit" from all keys that start with "vit"
				UpperCamelCase__					:       Union[str, Any]       =       [(pair[0], pair[1][4:]) if pair[1].startswith('vit') else pair for pair in rename_keys]
		else:
				# layernorm + classification head
				rename_keys.extend(
				    [
				        ('norm.weight', 'vit.layernorm.weight'),
				        ('norm.bias', 'vit.layernorm.bias'),
				        ('head.weight', 'classifier.weight'),
				        ('head.bias', 'classifier.bias'),
				    ])
		# fmt: on

		return rename_keys







def    __UpperCAmelCase    (	lowerCamelCase_  ,   lowerCamelCase_  ,   lowerCamelCase_=False)				->  Tuple:


		for i in range(config.num_hidden_layers):
				if base_model:
						UpperCamelCase__					:       Union[str, Any]       =       ''
				else:
						UpperCamelCase__					:       Optional[int]       =       'vit.'
				# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
				UpperCamelCase__					:       List[Any]       =       state_dict.pop(f'blocks.{i}.attn.qkv.weight')
				UpperCamelCase__					:       Optional[Any]       =       state_dict.pop(f'blocks.{i}.attn.qkv.bias')
				# next, add query, keys and values (in that order) to the state dict
				UpperCamelCase__					:       List[str]       =       in_proj_weight[
				    : config.hidden_size, :
				]
				UpperCamelCase__					:       str       =       in_proj_bias[: config.hidden_size]
				UpperCamelCase__					:       Optional[int]       =       in_proj_weight[
				    config.hidden_size : config.hidden_size * 2, :
				]
				UpperCamelCase__					:       Tuple       =       in_proj_bias[
				    config.hidden_size : config.hidden_size * 2
				]
				UpperCamelCase__					:       Union[str, Any]       =       in_proj_weight[
				    -config.hidden_size :, :
				]
				UpperCamelCase__					:       Dict       =       in_proj_bias[-config.hidden_size :]







def    __UpperCAmelCase    (	lowerCamelCase_)				->  str:
		UpperCamelCase__					:       Tuple       =       ['head.weight', 'head.bias']
		for k in ignore_keys:
				state_dict.pop(UpperCamelCase__  ,   UpperCamelCase__)







def    __UpperCAmelCase    (	lowerCamelCase_  ,   lowerCamelCase_  ,   lowerCamelCase_)				->  List[Any]:
		UpperCamelCase__					:       int       =       dct.pop(UpperCamelCase__)
		UpperCamelCase__					:       Union[str, Any]       =       val







def    __UpperCAmelCase    (	)				->  Optional[Any]:
		UpperCamelCase__					:       Any       =       'http://images.cocodataset.org/val2017/000000039769.jpg'
		UpperCamelCase__					:       Optional[int]       =       Image.open(requests.get(UpperCamelCase__  ,   stream=UpperCamelCase__).raw)
		return im







@torch.no_grad()
def    __UpperCAmelCase    (	lowerCamelCase_  ,   lowerCamelCase_  ,   lowerCamelCase_=False)				->  Any:
		UpperCamelCase__					:       Dict       =       BitConfig(
		    global_padding='same'  ,   layer_type='bottleneck'  ,   depths=(3, 4, 9)  ,   out_features=['stage3']  ,   embedding_dynamic_padding=UpperCamelCase__  ,   )
		UpperCamelCase__					:       int       =       ViTHybridConfig(backbone_config=UpperCamelCase__  ,   image_size=384  ,   num_labels=1_000)
		UpperCamelCase__					:       List[str]       =       False

		# load original model from timm
		UpperCamelCase__					:       List[str]       =       timm.create_model(UpperCamelCase__  ,   pretrained=UpperCamelCase__)
		timm_model.eval()

		# load state_dict of original model, remove and rename some keys
		UpperCamelCase__					:       Any       =       timm_model.state_dict()
		if base_model:
				remove_classification_head_(UpperCamelCase__)
		UpperCamelCase__					:       Optional[Any]       =       create_rename_keys(UpperCamelCase__  ,   UpperCamelCase__)
		for src, dest in rename_keys:
				rename_key(UpperCamelCase__  ,   UpperCamelCase__  ,   UpperCamelCase__)
		read_in_q_k_v(UpperCamelCase__  ,   UpperCamelCase__  ,   UpperCamelCase__)

		UpperCamelCase__					:       Union[str, Any]       =       'huggingface/label-files'
		UpperCamelCase__					:       Tuple       =       'imagenet-1k-id2label.json'
		UpperCamelCase__					:       Optional[int]       =       json.load(open(hf_hub_download(UpperCamelCase__  ,   UpperCamelCase__  ,   repo_type='dataset')  ,   'r'))
		UpperCamelCase__					:       List[Any]       =       {int(UpperCamelCase__): v for k, v in idalabel.items()}
		UpperCamelCase__					:       Any       =       idalabel
		UpperCamelCase__					:       Tuple       =       {v: k for k, v in idalabel.items()}

		# load HuggingFace model
		if vit_name[-5:] == "in21k":
				UpperCamelCase__					:       Union[str, Any]       =       ViTHybridModel(UpperCamelCase__).eval()
		else:
				UpperCamelCase__					:       int       =       ViTHybridForImageClassification(UpperCamelCase__).eval()
		model.load_state_dict(UpperCamelCase__)

		# create image processor
		UpperCamelCase__					:       Any       =       create_transform(**resolve_data_config({}  ,   model=UpperCamelCase__))
		UpperCamelCase__					:       Any       =       transform.transforms

		UpperCamelCase__					:       Any       =       {
		    'bilinear': PILImageResampling.BILINEAR,
		    'bicubic': PILImageResampling.BICUBIC,
		    'nearest': PILImageResampling.NEAREST,
		}

		UpperCamelCase__					:       Any       =       ViTHybridImageProcessor(
		    do_resize=UpperCamelCase__  ,   size={'shortest_edge': timm_transforms[0].size}  ,   resample=pillow_resamplings[timm_transforms[0].interpolation.value]  ,   do_center_crop=UpperCamelCase__  ,   crop_size={'height': timm_transforms[1].size[0], 'width': timm_transforms[1].size[1]}  ,   do_normalize=UpperCamelCase__  ,   image_mean=timm_transforms[-1].mean.tolist()  ,   image_std=timm_transforms[-1].std.tolist()  ,   )

		UpperCamelCase__					:       Optional[Any]       =       prepare_img()
		UpperCamelCase__					:       Tuple       =       transform(UpperCamelCase__).unsqueeze(0)
		UpperCamelCase__					:       int       =       processor(UpperCamelCase__  ,   return_tensors='pt').pixel_values

		# verify pixel values
		assert torch.allclose(UpperCamelCase__  ,   UpperCamelCase__)

		# verify logits
		with torch.no_grad():
				UpperCamelCase__					:       Dict       =       model(UpperCamelCase__)
				UpperCamelCase__					:       str       =       outputs.logits

		print('Predicted class:'  ,   logits.argmax(-1).item())
		if base_model:
				UpperCamelCase__					:       Dict       =       timm_model.forward_features(UpperCamelCase__)
				assert timm_pooled_output.shape == outputs.pooler_output.shape
				assert torch.allclose(UpperCamelCase__  ,   outputs.pooler_output  ,   atol=1e-3)
		else:
				UpperCamelCase__					:       Tuple       =       timm_model(UpperCamelCase__)
				assert timm_logits.shape == outputs.logits.shape
				assert torch.allclose(UpperCamelCase__  ,   outputs.logits  ,   atol=1e-3)
		print('Looks ok!')

		if pytorch_dump_folder_path is not None:
				Path(UpperCamelCase__).mkdir(exist_ok=UpperCamelCase__)
				print(f'Saving model {vit_name} to {pytorch_dump_folder_path}')
				model.save_pretrained(UpperCamelCase__)
				print(f'Saving processor to {pytorch_dump_folder_path}')
				processor.save_pretrained(UpperCamelCase__)

		if push_to_hub:
				print(f'Pushing model and processor to the hub {vit_name}')
				model.push_to_hub(f'ybelkada/{vit_name}')
				processor.push_to_hub(f'ybelkada/{vit_name}')


if __name__ == "__main__":
		lowerCAmelCase__ 							=    argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    '--vit_name',
		    default='vit_base_r50_s16_384',
		    type=str,
		    help='Name of the hybrid ViT timm model you\'d like to convert.',
		)
		parser.add_argument(
		    '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
		)
		parser.add_argument(
		    '--push_to_hub', action='store_true', help='Whether to upload the model to the HuggingFace hub.'
		)

		lowerCAmelCase__ 							=    parser.parse_args()
		convert_vit_checkpoint(args.vit_name, args.pytorch_dump_folder_path, args.push_to_hub)










def      lowerCamelCase_						(			)	->			List[str]:





				"""simple docstring"""


				__lowerCamelCase									=       0
				for i in range(1					,	1001 ):
								total += i**i
				return str(UpperCamelCase__ )[-10:]


if __name__ == "__main__":
					print(solution())


