



'''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_  :							Any             =    logging.get_logger(__name__)
def        UpperCAmelCase (    A	:					Optional[Any]       ,		A	:					Dict=False					):
							SCREAMING_SNAKE_CASE       :   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"
														SCREAMING_SNAKE_CASE       :   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 (    A	:					Union[str, Any]       ,		A	:					Tuple       ,		A	:					List[str]=False					):
							for i in range(config.num_hidden_layers					):
														if base_model:
																					SCREAMING_SNAKE_CASE       :   List[str]         =      ''''''
														else:
																					SCREAMING_SNAKE_CASE       :   int         =      '''vit.'''
														# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
														SCREAMING_SNAKE_CASE       :   Dict         =      state_dict.pop(F"""blocks.{i}.attn.qkv.weight"""					)
														SCREAMING_SNAKE_CASE       :   List[str]         =      state_dict.pop(F"""blocks.{i}.attn.qkv.bias"""					)
														# next, add query, keys and values (in that order) to the state dict
														SCREAMING_SNAKE_CASE       :   Optional[int]         =      in_proj_weight[
														    : config.hidden_size, :
														]
														SCREAMING_SNAKE_CASE       :   int         =      in_proj_bias[: config.hidden_size]
														SCREAMING_SNAKE_CASE       :   Dict         =      in_proj_weight[
														    config.hidden_size : config.hidden_size * 2, :
														]
														SCREAMING_SNAKE_CASE       :   List[str]         =      in_proj_bias[
														    config.hidden_size : config.hidden_size * 2
														]
														SCREAMING_SNAKE_CASE       :   List[Any]         =      in_proj_weight[
														    -config.hidden_size :, :
														]
														SCREAMING_SNAKE_CASE       :   Dict         =      in_proj_bias[-config.hidden_size :]
def        UpperCAmelCase (    A	:					Tuple					):
							SCREAMING_SNAKE_CASE       :   int         =      ['''head.weight''', '''head.bias''']
							for k in ignore_keys:
														state_dict.pop(A       ,		A					)
def        UpperCAmelCase (    A	:					str       ,		A	:					Any       ,		A	:					Tuple					):
							SCREAMING_SNAKE_CASE       :   Any         =      dct.pop(A					)
							SCREAMING_SNAKE_CASE       :   str         =      val
def        UpperCAmelCase (    ):
							SCREAMING_SNAKE_CASE       :   List[Any]         =      '''http://images.cocodataset.org/val2017/000000039769.jpg'''
							SCREAMING_SNAKE_CASE       :   str         =      Image.open(requests.get(A       ,		stream=A					).raw					)
							return im
@torch.no_grad()
def        UpperCAmelCase (    A	:					Tuple       ,		A	:					List[Any]       ,		A	:					Dict=False					):
							SCREAMING_SNAKE_CASE       :   Any         =      BitConfig(
							    global_padding='''same'''       ,		layer_type='''bottleneck'''       ,		depths=(3, 4, 9)       ,		out_features=['''stage3''']       ,		embedding_dynamic_padding=A       ,		)
							SCREAMING_SNAKE_CASE       :   Optional[int]         =      ViTHybridConfig(backbone_config=A       ,		image_size=384       ,		num_labels=1000					)
							SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      False

							# load original model from timm
							SCREAMING_SNAKE_CASE       :   Optional[int]         =      timm.create_model(A       ,		pretrained=A					)
							timm_model.eval()

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

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

							# load HuggingFace model
							if vit_name[-5:] == "in21k":
														SCREAMING_SNAKE_CASE       :   List[str]         =      ViTHybridModel(A					).eval()
							else:
														SCREAMING_SNAKE_CASE       :   str         =      ViTHybridForImageClassification(A					).eval()
							model.load_state_dict(A					)

							# create image processor
							SCREAMING_SNAKE_CASE       :   str         =      create_transform(**resolve_data_config({}       ,		model=A					)					)
							SCREAMING_SNAKE_CASE       :   Tuple         =      transform.transforms

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

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

							SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      prepare_img()
							SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      transform(A					).unsqueeze(0					)
							SCREAMING_SNAKE_CASE       :   str         =      processor(A       ,		return_tensors='''pt'''					).pixel_values

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

							# verify logits
							with torch.no_grad():
														SCREAMING_SNAKE_CASE       :   str         =      model(A					)
														SCREAMING_SNAKE_CASE       :   str         =      outputs.logits

							print('''Predicted class:'''       ,		logits.argmax(-1					).item()					)
							if base_model:
														SCREAMING_SNAKE_CASE       :   int         =      timm_model.forward_features(A					)
														assert timm_pooled_output.shape == outputs.pooler_output.shape
														assert torch.allclose(A       ,		outputs.pooler_output       ,		atol=1e-3					)
							else:
														SCREAMING_SNAKE_CASE       :   List[str]         =      timm_model(A					)
														assert timm_logits.shape == outputs.logits.shape
														assert torch.allclose(A       ,		outputs.logits       ,		atol=1e-3					)
							print('''Looks ok!'''					)

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

							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_  :							List[Any]             =    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_  :							Union[str, Any]             =    parser.parse_args()
							convert_vit_checkpoint(args.vit_name, args.pytorch_dump_folder_path, args.push_to_hub)









'''simple docstring'''



import unittest
from pathlib import Path
from tempfile import TemporaryDirectory

from transformers import AutoConfig, TFAutoModel, is_tensorflow_text_available, is_tf_available
from transformers.models.bert.tokenization_bert import BertTokenizer
from transformers.testing_utils import require_tensorflow_text, require_tf, slow


if is_tf_available():
							import tensorflow as tf

if is_tensorflow_text_available():
							from transformers.models.bert import TFBertTokenizer


lowerCAmelCase_  :							List[str]             =    ['bert-base-uncased', 'bert-base-cased']
lowerCAmelCase_  :							int             =    'hf-internal-testing/tiny-bert-tf-only'

if is_tf_available():
							class 						lowerCamelCase_				(       tf.keras.Model      ):







													def __init__(	self :   Optional[int]					,     lowerCAmelCase__ :   Optional[Any]					):

																				"""simple docstring"""
																				super().__init__()
																				SCREAMING_SNAKE_CASE       :   List[str]         =      tokenizer
																				SCREAMING_SNAKE_CASE       :   str         =      AutoConfig.from_pretrained(lowerCAmelCase__					)
																				SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      TFAutoModel.from_config(lowerCAmelCase__					)







													def        __lowercase     (	self :   Tuple					,     lowerCAmelCase__ :   Tuple					):

																				"""simple docstring"""
																				SCREAMING_SNAKE_CASE       :   Optional[Any]         =      self.tokenizer(lowerCAmelCase__					)
																				SCREAMING_SNAKE_CASE       :   str         =      self.bert(**lowerCAmelCase__					)
																				return out["pooler_output"]
@require_tf
@require_tensorflow_text
class 						lowerCamelCase_				(       unittest.TestCase      ):







						def        __lowercase     (	self :   str					):

													"""simple docstring"""
													super().setUp()

													SCREAMING_SNAKE_CASE       :   int         =      [
													    BertTokenizer.from_pretrained(lowerCAmelCase__					) for checkpoint in (TOKENIZER_CHECKPOINTS * 2)
													]  # repeat for when fast_bert_tokenizer=false
													SCREAMING_SNAKE_CASE       :   Dict         =      [TFBertTokenizer.from_pretrained(lowerCAmelCase__					) for checkpoint in TOKENIZER_CHECKPOINTS] + [
													    TFBertTokenizer.from_pretrained(lowerCAmelCase__					,     use_fast_bert_tokenizer=lowerCAmelCase__					)
													    for checkpoint in TOKENIZER_CHECKPOINTS
													]
													assert len(self.tokenizers					) == len(self.tf_tokenizers					)

													SCREAMING_SNAKE_CASE       :   Any         =      [
													    '''This is a straightforward English test sentence.''',
													    '''This one has some weird characters\rto\nsee\r\nif  those\u00E9break things.''',
													    '''Now we\'re going to add some Chinese: 一 二 三 一二三''',
													    '''And some much more rare Chinese: 齉 堃 齉堃''',
													    '''Je vais aussi écrire en français pour tester les accents''',
													    '''Classical Irish also has some unusual characters, so in they go: Gaelaċ, ꝼ''',
													]
													SCREAMING_SNAKE_CASE       :   int         =      list(zip(self.test_sentences					,     self.test_sentences[::-1]					)					)







						def        __lowercase     (	self :   Any					):

													"""simple docstring"""
													for tokenizer, tf_tokenizer in zip(self.tokenizers					,     self.tf_tokenizers					):
																				for test_inputs in (self.test_sentences, self.paired_sentences):
																											SCREAMING_SNAKE_CASE       :   Tuple         =      tokenizer(lowerCAmelCase__					,     return_tensors='''tf'''					,     padding='''longest'''					)
																											SCREAMING_SNAKE_CASE       :   List[Any]         =      tf_tokenizer(lowerCAmelCase__					)

																											for key in python_outputs.keys():
																																		self.assertTrue(tf.reduce_all(python_outputs[key].shape == tf_outputs[key].shape					)					)
																																		self.assertTrue(tf.reduce_all(tf.cast(python_outputs[key]					,     tf.intaa					) == tf_outputs[key]					)					)







						@slow
						def        __lowercase     (	self :   Dict					):

													"""simple docstring"""
													for tf_tokenizer in self.tf_tokenizers:
																				SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      tf_tokenizer(self.paired_sentences					)
																				SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      tf_tokenizer(
																				    text=[sentence[0] for sentence in self.paired_sentences]					,     text_pair=[sentence[1] for sentence in self.paired_sentences]					,     )
																				for key in merged_outputs.keys():
																											self.assertTrue(tf.reduce_all(tf.cast(merged_outputs[key]					,     tf.intaa					) == separated_outputs[key]					)					)







						@slow
						def        __lowercase     (	self :   Union[str, Any]					):

													"""simple docstring"""
													for tf_tokenizer in self.tf_tokenizers:
																				SCREAMING_SNAKE_CASE       :   Any         =      tf.function(lowerCAmelCase__					)
																				for test_inputs in (self.test_sentences, self.paired_sentences):
																											SCREAMING_SNAKE_CASE       :   int         =      tf.constant(lowerCAmelCase__					)
																											SCREAMING_SNAKE_CASE       :   Any         =      compiled_tokenizer(lowerCAmelCase__					)
																											SCREAMING_SNAKE_CASE       :   List[str]         =      tf_tokenizer(lowerCAmelCase__					)

																											for key in eager_outputs.keys():
																																		self.assertTrue(tf.reduce_all(eager_outputs[key] == compiled_outputs[key]					)					)







						@slow
						def        __lowercase     (	self :   int					):

													"""simple docstring"""
													for tf_tokenizer in self.tf_tokenizers:
																				SCREAMING_SNAKE_CASE       :   int         =      ModelToSave(tokenizer=lowerCAmelCase__					)
																				SCREAMING_SNAKE_CASE       :   Tuple         =      tf.convert_to_tensor(self.test_sentences					)
																				SCREAMING_SNAKE_CASE       :   Optional[Any]         =      model(lowerCAmelCase__					)  # Build model with some sample inputs
																				with TemporaryDirectory() as tempdir:
																											SCREAMING_SNAKE_CASE       :   Optional[int]         =      Path(lowerCAmelCase__					) / '''saved.model'''
																											model.save(lowerCAmelCase__					)
																											SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      tf.keras.models.load_model(lowerCAmelCase__					)
																				SCREAMING_SNAKE_CASE       :   Tuple         =      loaded_model(lowerCAmelCase__					)
																				# We may see small differences because the loaded model is compiled, so we need an epsilon for the test
																				self.assertLessEqual(tf.reduce_max(tf.abs(out - loaded_output					)					)					,     1e-5					)




