





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 lowerCAmelCase__(__snake_case       ,__snake_case=False							)				->      Any:







	'''simple docstring'''







	lowerCamelCase__								=		[]

	# 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"
		lowerCamelCase__								=		[(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 lowerCAmelCase__(__snake_case       ,__snake_case       ,__snake_case=False							)				->      str:







	'''simple docstring'''







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





def lowerCAmelCase__(__snake_case							)				->      Union[str, Any]:







	'''simple docstring'''







	lowerCamelCase__								=		['''head.weight''', '''head.bias''']
	for k in ignore_keys:
		state_dict.pop(__snake_case       ,__snake_case							)





def lowerCAmelCase__(__snake_case       ,__snake_case       ,__snake_case							)				->      Any:







	'''simple docstring'''







	lowerCamelCase__								=		dct.pop(__snake_case							)
	lowerCamelCase__								=		val





def lowerCAmelCase__()				->      Tuple:







	'''simple docstring'''







	lowerCamelCase__								=		'''http://images.cocodataset.org/val2017/000000039769.jpg'''
	lowerCamelCase__								=		Image.open(requests.get(__snake_case       ,stream=__snake_case							).raw							)
	return im





@torch.no_grad()
def lowerCAmelCase__(__snake_case       ,__snake_case       ,__snake_case=False							)				->      str:







	'''simple docstring'''







	lowerCamelCase__								=		BitConfig(
	    global_padding='''same'''       ,layer_type='''bottleneck'''       ,depths=(3, 4, 9)       ,out_features=['''stage3''']       ,embedding_dynamic_padding=__snake_case       ,)
	lowerCamelCase__								=		ViTHybridConfig(backbone_config=__snake_case       ,image_size=384       ,num_labels=1000							)
	lowerCamelCase__								=		False

	# load original model from timm
	lowerCamelCase__								=		timm.create_model(__snake_case       ,pretrained=__snake_case							)
	timm_model.eval()

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

	lowerCamelCase__								=		'''huggingface/label-files'''
	lowerCamelCase__								=		'''imagenet-1k-id2label.json'''
	lowerCamelCase__								=		json.load(open(hf_hub_download(__snake_case       ,__snake_case       ,repo_type='''dataset'''							)       ,'''r'''							)							)
	lowerCamelCase__								=		{int(__snake_case							): v for k, v in idalabel.items()}
	lowerCamelCase__								=		idalabel
	lowerCamelCase__								=		{v: k for k, v in idalabel.items()}

	# load HuggingFace model
	if vit_name[-5:] == "in21k":
		lowerCamelCase__								=		ViTHybridModel(__snake_case							).eval()
	else:
		lowerCamelCase__								=		ViTHybridForImageClassification(__snake_case							).eval()
	model.load_state_dict(__snake_case							)

	# create image processor
	lowerCamelCase__								=		create_transform(**resolve_data_config({}       ,model=__snake_case							)							)
	lowerCamelCase__								=		transform.transforms

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

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

	lowerCamelCase__								=		prepare_img()
	lowerCamelCase__								=		transform(__snake_case							).unsqueeze(0							)
	lowerCamelCase__								=		processor(__snake_case       ,return_tensors='''pt'''							).pixel_values

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

	# verify logits
	with torch.no_grad():
		lowerCamelCase__								=		model(__snake_case							)
		lowerCamelCase__								=		outputs.logits

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

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

	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)










import pytest

from datasets.parallel import ParallelBackendConfig, parallel_backend
from datasets.utils.py_utils import map_nested

from .utils import require_dill_gt_0_3_2, require_joblibspark, require_not_windows





def lowerCAmelCase__(__snake_case							)				->      int:  # picklable for multiprocessing







	'''simple docstring'''







	return i + 1





@require_dill_gt_0_3_2
@require_joblibspark
@require_not_windows
def lowerCAmelCase__()				->      Any:







	'''simple docstring'''







	with parallel_backend('''spark'''							):
		assert ParallelBackendConfig.backend_name == "spark"

	lowerCamelCase__								=		[1, 2, 3]
	with pytest.raises(__snake_case							):
		with parallel_backend('''unsupported backend'''							):
			map_nested(__snake_case       ,__snake_case       ,num_proc=2							)

	with pytest.raises(__snake_case							):
		with parallel_backend('''unsupported backend'''							):
			map_nested(__snake_case       ,__snake_case       ,num_proc=-1							)





@require_dill_gt_0_3_2
@require_joblibspark
@require_not_windows
@pytest.mark.parametrize('''num_proc'''       ,[2, -1]							)
def lowerCAmelCase__(__snake_case							)				->      Tuple:







	'''simple docstring'''







	lowerCamelCase__								=		[1, 2]
	lowerCamelCase__								=		{'''a''': 1, '''b''': 2}
	lowerCamelCase__								=		{'''a''': [1, 2], '''b''': [3, 4]}
	lowerCamelCase__								=		{'''a''': {'''1''': 1}, '''b''': 2}
	lowerCamelCase__								=		{'''a''': 1, '''b''': 2, '''c''': 3, '''d''': 4}
	lowerCamelCase__								=		[2, 3]
	lowerCamelCase__								=		{'''a''': 2, '''b''': 3}
	lowerCamelCase__								=		{'''a''': [2, 3], '''b''': [4, 5]}
	lowerCamelCase__								=		{'''a''': {'''1''': 2}, '''b''': 3}
	lowerCamelCase__								=		{'''a''': 2, '''b''': 3, '''c''': 4, '''d''': 5}

	with parallel_backend('''spark'''							):
		assert map_nested(__snake_case       ,__snake_case       ,num_proc=__snake_case							) == expected_map_nested_sa
		assert map_nested(__snake_case       ,__snake_case       ,num_proc=__snake_case							) == expected_map_nested_sa
		assert map_nested(__snake_case       ,__snake_case       ,num_proc=__snake_case							) == expected_map_nested_sa
		assert map_nested(__snake_case       ,__snake_case       ,num_proc=__snake_case							) == expected_map_nested_sa
		assert map_nested(__snake_case       ,__snake_case       ,num_proc=__snake_case							) == expected_map_nested_sa



