



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







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





			"""simple docstring"""







			snake_case:    Optional[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"
						snake_case:    Dict      							=[(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 a_     (    __UpperCAmelCase				,							__UpperCAmelCase				,							__UpperCAmelCase=False					)				->					Dict:





			"""simple docstring"""







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







def a_     (    __UpperCAmelCase					)				->					Any:





			"""simple docstring"""







			snake_case:    Dict      							=['head.weight', 'head.bias']
			for k in ignore_keys:
						state_dict.pop(__UpperCAmelCase				,							__UpperCAmelCase					)







def a_     (    __UpperCAmelCase				,							__UpperCAmelCase				,							__UpperCAmelCase					)				->					Any:





			"""simple docstring"""







			snake_case:    Optional[int]      							=dct.pop(__UpperCAmelCase					)
			snake_case:    List[Any]      							=val







def a_     (    )				->					Tuple:





			"""simple docstring"""







			snake_case:    List[str]      							='http://images.cocodataset.org/val2017/000000039769.jpg'
			snake_case:    Any      							=Image.open(requests.get(__UpperCAmelCase				,							stream=__UpperCAmelCase					).raw					)
			return im







@torch.no_grad()
def a_     (    __UpperCAmelCase				,							__UpperCAmelCase				,							__UpperCAmelCase=False					)				->					Optional[Any]:





			"""simple docstring"""







			snake_case:    Union[str, Any]      							=BitConfig(
			    global_padding='same'				,							layer_type='bottleneck'				,							depths=(3, 4, 9)				,							out_features=['stage3']				,							embedding_dynamic_padding=__UpperCAmelCase				,							)
			snake_case:    List[Any]      							=ViTHybridConfig(backbone_config=__UpperCAmelCase				,							image_size=3_84				,							num_labels=10_00					)
			snake_case:    Union[str, Any]      							=False

			# load original model from timm
			snake_case:    Dict      							=timm.create_model(__UpperCAmelCase				,							pretrained=__UpperCAmelCase					)
			timm_model.eval()

			# load state_dict of original model, remove and rename some keys
			snake_case:    Union[str, Any]      							=timm_model.state_dict()
			if base_model:
						remove_classification_head_(__UpperCAmelCase					)
			snake_case:    List[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					)

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

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

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

			snake_case:    Union[str, Any]      							={
			    'bilinear': PILImageResampling.BILINEAR,
			    'bicubic': PILImageResampling.BICUBIC,
			    'nearest': PILImageResampling.NEAREST,
			}

			snake_case:    List[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()				,							)

			snake_case:    Dict      							=prepare_img()
			snake_case:    List[Any]      							=transform(__UpperCAmelCase					).unsqueeze(0					)
			snake_case:    int      							=processor(__UpperCAmelCase				,							return_tensors='pt'					).pixel_values

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

			# verify logits
			with torch.no_grad():
						snake_case:    Tuple      							=model(__UpperCAmelCase					)
						snake_case:    List[str]      							=outputs.logits

			print('Predicted class:'				,							logits.argmax(-1					).item()					)
			if base_model:
						snake_case:    Union[str, Any]      							=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:
						snake_case:    Union[str, Any]      							=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__":
	a							   =				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.'
	)

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









'''simple docstring'''

from __future__ import annotations







def a_     (    __UpperCAmelCase					)				->					list[int]:





			"""simple docstring"""







			snake_case:    Tuple      							=[True] * limit
			snake_case:    Optional[int]      							=False
			snake_case:    Union[str, Any]      							=False
			snake_case:    List[Any]      							=True

			for i in range(3				,							int(limit**0.5 + 1					)				,							2					):
						snake_case:    str      							=i * 2
						while index < limit:
									snake_case:    List[Any]      							=False
									snake_case:    List[str]      							=index + i

			snake_case:    Union[str, Any]      							=[2]

			for i in range(3				,							__UpperCAmelCase				,							2					):
						if is_prime[i]:
									primes.append(__UpperCAmelCase					)

			return primes







def a_     (    __UpperCAmelCase = 1_00_00_00					)				->					int:





			"""simple docstring"""







			snake_case:    str      							=prime_sieve(__UpperCAmelCase					)
			snake_case:    str      							=0
			snake_case:    str      							=0

			for i in range(len(__UpperCAmelCase					)					):
						for j in range(i + length				,							len(__UpperCAmelCase					)					):
									snake_case:    Tuple      							=sum(primes[i:j]					)
									if sol >= ceiling:
												break

									if sol in primes:
												snake_case:    List[str]      							=j - i
												snake_case:    Optional[int]      							=sol

			return largest


if __name__ == "__main__":
	print(F"""{solution() = }""")




