

import json
import os
import shutil
import sys
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import AutoConfig, BertConfig, GPTaConfig
from transformers.configuration_utils import PretrainedConfig
from transformers.testing_utils import TOKEN, USER, is_staging_test


sys.path.append(str(Path(__file__).parent.parent / 'utils'))

from test_module.custom_configuration import CustomConfig  # noqa E402


SCREAMING_SNAKE_CASE_            = {
    'return_dict': False,
    'output_hidden_states': True,
    'output_attentions': True,
    'torchscript': True,
    'torch_dtype': 'float16',
    'use_bfloat16': True,
    'tf_legacy_loss': True,
    'pruned_heads': {'a': 1},
    'tie_word_embeddings': False,
    'is_decoder': True,
    'cross_attention_hidden_size': 128,
    'add_cross_attention': True,
    'tie_encoder_decoder': True,
    'max_length': 50,
    'min_length': 3,
    'do_sample': True,
    'early_stopping': True,
    'num_beams': 3,
    'num_beam_groups': 3,
    'diversity_penalty': 0.5,
    'temperature': 2.0,
    'top_k': 10,
    'top_p': 0.7,
    'typical_p': 0.2,
    'repetition_penalty': 0.8,
    'length_penalty': 0.8,
    'no_repeat_ngram_size': 5,
    'encoder_no_repeat_ngram_size': 5,
    'bad_words_ids': [1, 2, 3],
    'num_return_sequences': 3,
    'chunk_size_feed_forward': 5,
    'output_scores': True,
    'return_dict_in_generate': True,
    'forced_bos_token_id': 2,
    'forced_eos_token_id': 3,
    'remove_invalid_values': True,
    'architectures': ['BertModel'],
    'finetuning_task': 'translation',
    'id2label': {0: 'label'},
    'label2id': {'label': '0'},
    'tokenizer_class': 'BertTokenizerFast',
    'prefix': 'prefix',
    'bos_token_id': 6,
    'pad_token_id': 7,
    'eos_token_id': 8,
    'sep_token_id': 9,
    'decoder_start_token_id': 10,
    'exponential_decay_length_penalty': (5, 1.01),
    'suppress_tokens': [0, 1],
    'begin_suppress_tokens': 2,
    'task_specific_params': {'translation': 'some_params'},
    'problem_type': 'regression',
}

@is_staging_test
class 			a					(		unittest.TestCase			):


						@classmethod
						def     _UpperCAmelCase			(    cls		):







										'''simple docstring'''


										_UpperCAmelCase    :						int							=      TOKEN
										HfFolder.save_token(A_		)


						@classmethod
						def     _UpperCAmelCase			(    cls		):







										'''simple docstring'''


										try:
														delete_repo(token=cls._token  ,   repo_id="test-config"		)
										except HTTPError:
														pass

										try:
														delete_repo(token=cls._token  ,   repo_id="valid_org/test-config-org"		)
										except HTTPError:
														pass

										try:
														delete_repo(token=cls._token  ,   repo_id="test-dynamic-config"		)
										except HTTPError:
														pass


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										_UpperCAmelCase    :						List[str]							=      BertConfig(
										    vocab_size=99  ,   hidden_size=32  ,   num_hidden_layers=5  ,   num_attention_heads=4  ,   intermediate_size=37		)
										config.push_to_hub("test-config"  ,   use_auth_token=self._token		)

										_UpperCAmelCase    :						Tuple							=      BertConfig.from_pretrained(f'{USER}/test-config'		)
										for k, v in config.to_dict().items():
														if k != "transformers_version":
																		self.assertEqual(A_  ,   getattr(A_  ,   A_		)		)

        # Reset repo
										delete_repo(token=self._token  ,   repo_id="test-config"		)

										# Push to hub via save_pretrained
										with tempfile.TemporaryDirectory() as tmp_dir:
														config.save_pretrained(A_  ,   repo_id="test-config"  ,   push_to_hub=A_  ,   use_auth_token=self._token		)

										_UpperCAmelCase    :						Any							=      BertConfig.from_pretrained(f'{USER}/test-config'		)
										for k, v in config.to_dict().items():
														if k != "transformers_version":
																		self.assertEqual(A_  ,   getattr(A_  ,   A_		)		)


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										_UpperCAmelCase    :						Optional[Any]							=      BertConfig(
										    vocab_size=99  ,   hidden_size=32  ,   num_hidden_layers=5  ,   num_attention_heads=4  ,   intermediate_size=37		)
										config.push_to_hub("valid_org/test-config-org"  ,   use_auth_token=self._token		)

										_UpperCAmelCase    :						Dict							=      BertConfig.from_pretrained("valid_org/test-config-org"		)
										for k, v in config.to_dict().items():
														if k != "transformers_version":
																		self.assertEqual(A_  ,   getattr(A_  ,   A_		)		)

        # Reset repo
										delete_repo(token=self._token  ,   repo_id="valid_org/test-config-org"		)

										# Push to hub via save_pretrained
										with tempfile.TemporaryDirectory() as tmp_dir:
														config.save_pretrained(
														    A_  ,   repo_id="valid_org/test-config-org"  ,   push_to_hub=A_  ,   use_auth_token=self._token		)

										_UpperCAmelCase    :						Any							=      BertConfig.from_pretrained("valid_org/test-config-org"		)
										for k, v in config.to_dict().items():
														if k != "transformers_version":
																		self.assertEqual(A_  ,   getattr(A_  ,   A_		)		)





						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										CustomConfig.register_for_auto_class()
										_UpperCAmelCase    :						Any							=      CustomConfig(attribute=42		)

										config.push_to_hub("test-dynamic-config"  ,   use_auth_token=self._token		)

										# This has added the proper auto_map field to the config
										self.assertDictEqual(config.auto_map  ,   {"AutoConfig": "custom_configuration.CustomConfig"}		)

										_UpperCAmelCase    :						Union[str, Any]							=      AutoConfig.from_pretrained(f'{USER}/test-dynamic-config'  ,   trust_remote_code=A_		)
										# Can't make an isinstance check because the new_config is from the FakeConfig class of a dynamic module
										self.assertEqual(new_config.__class__.__name__  ,   "CustomConfig"		)
										self.assertEqual(new_config.attribute  ,   42		)







class 			a					(		unittest.TestCase			):


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										_UpperCAmelCase    :						Tuple							=      GPTaConfig()

										# attempt to modify each of int/float/bool/str config records and verify they were updated
										_UpperCAmelCase    :						Dict							=      c.n_embd + 1  # int
										_UpperCAmelCase    :						Tuple							=      c.resid_pdrop + 1.0  # float
										_UpperCAmelCase    :						Optional[int]							=      not c.scale_attn_weights  # bool
										_UpperCAmelCase    :						Optional[Any]							=      c.summary_type + "foo"  # str
										c.update_from_string(
										    f'n_embd={n_embd},resid_pdrop={resid_pdrop},scale_attn_weights={scale_attn_weights},summary_type={summary_type}'		)
										self.assertEqual(A_  ,   c.n_embd  ,   "mismatch for key: n_embd"		)
										self.assertEqual(A_  ,   c.resid_pdrop  ,   "mismatch for key: resid_pdrop"		)
										self.assertEqual(A_  ,   c.scale_attn_weights  ,   "mismatch for key: scale_attn_weights"		)
										self.assertEqual(A_  ,   c.summary_type  ,   "mismatch for key: summary_type"		)


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										_UpperCAmelCase    :						int							=      PretrainedConfig()
										_UpperCAmelCase    :						int							=      [key for key in base_config.__dict__ if key not in config_common_kwargs]
										# If this part of the test fails, you have arguments to addin config_common_kwargs above.
										self.assertListEqual(
										    A_  ,   ["is_encoder_decoder", "_name_or_path", "_commit_hash", "transformers_version"]		)
										_UpperCAmelCase    :						List[Any]							=      [key for key, value in config_common_kwargs.items() if value == getattr(A_  ,   A_		)]
										if len(A_		) > 0:
														raise ValueError(
														    "The following keys are set with the default values in"
														    " `test_configuration_common.config_common_kwargs` pick another value for them:"
														    f' {", ".join(A_		)}.'		)


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										with self.assertRaises(A_		):
														# config is in subfolder, the following should not work without specifying the subfolder
														_UpperCAmelCase    :						List[str]							=      BertConfig.from_pretrained("hf-internal-testing/tiny-random-bert-subfolder"		)

										_UpperCAmelCase    :						Any							=      BertConfig.from_pretrained("hf-internal-testing/tiny-random-bert-subfolder"  ,   subfolder="bert"		)

										self.assertIsNotNone(A_		)


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										_UpperCAmelCase    :						Any							=      mock.Mock()
										_UpperCAmelCase    :						str							=      500
										_UpperCAmelCase    :						List[str]							=      {}
										_UpperCAmelCase    :						Tuple							=      HTTPError
										_UpperCAmelCase    :						Any							=      {}

										# Download this model to make sure it's in the cache.
										_UpperCAmelCase    :						Union[str, Any]							=      BertConfig.from_pretrained("hf-internal-testing/tiny-random-bert"		)

										# Under the mock environment we get a 500 error when trying to reach the model.
										with mock.patch("requests.Session.request"  ,   return_value=A_		) as mock_head:
														_UpperCAmelCase    :						Optional[int]							=      BertConfig.from_pretrained("hf-internal-testing/tiny-random-bert"		)
														# This check we did call the fake head request
														mock_head.assert_called()


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										_UpperCAmelCase    :						Tuple							=      BertConfig.from_pretrained(
										    "https://huggingface.co/hf-internal-testing/tiny-random-bert/resolve/main/config.json"		)


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										_UpperCAmelCase    :						Optional[Any]							=      AutoConfig.from_pretrained("bert-base-cased"		)
										_UpperCAmelCase    :						List[Any]							=      ["config.4.0.0.json"]

										with tempfile.TemporaryDirectory() as tmp_dir:
														configuration.save_pretrained(A_		)
														_UpperCAmelCase    :						Union[str, Any]							=      2
														json.dump(configuration.to_dict()  ,   open(os.path.join(A_  ,   "config.4.0.0.json"		)  ,   "w"		)		)

														# This should pick the new configuration file as the version of Transformers is > 4.0.0
														_UpperCAmelCase    :						List[Any]							=      AutoConfig.from_pretrained(A_		)
														self.assertEqual(new_configuration.hidden_size  ,   2		)

														# Will need to be adjusted if we reach v42 and this test is still here.
														# Should pick the old configuration file as the version of Transformers is < 4.42.0
														_UpperCAmelCase    :						int							=      ["config.42.0.0.json"]
														_UpperCAmelCase    :						str							=      768
														configuration.save_pretrained(A_		)
														shutil.move(os.path.join(A_  ,   "config.4.0.0.json"		)  ,   os.path.join(A_  ,   "config.42.0.0.json"		)		)
														_UpperCAmelCase    :						str							=      AutoConfig.from_pretrained(A_		)
														self.assertEqual(new_configuration.hidden_size  ,   768		)





						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										_UpperCAmelCase    :						Optional[int]							=      "hf-internal-testing/test-two-configs"

										import transformers as new_transformers

										_UpperCAmelCase    :						int							=      "v4.0.0"
										_UpperCAmelCase						,							_UpperCAmelCase    :						Optional[int]							=      new_transformers.models.auto.AutoConfig.from_pretrained(
										    A_  ,   return_unused_kwargs=A_		)
										self.assertEqual(new_configuration.hidden_size  ,   2		)
										# This checks `_configuration_file` ia not kept in the kwargs by mistake.
										self.assertDictEqual(A_  ,   {}		)

										# Testing an older version by monkey-patching the version in the module it's used.
										import transformers as old_transformers

										_UpperCAmelCase    :						List[str]							=      "v3.0.0"
										_UpperCAmelCase    :						Union[str, Any]							=      old_transformers.models.auto.AutoConfig.from_pretrained(A_		)
										self.assertEqual(old_configuration.hidden_size  ,   768		)





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

def 							__SCREAMING_SNAKE_CASE							(			lowerCAmelCase:					Optional[Any]					, lowerCAmelCase:					int=False	)				->    Optional[Any]:
				_UpperCAmelCase    :						Union[str, 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 							__SCREAMING_SNAKE_CASE							(			lowerCAmelCase:					Dict					, lowerCAmelCase:					Union[str, Any]					, lowerCAmelCase:					List[str]=False	)				->    int:
				for i in range(config.num_hidden_layers	):
								if base_model:
												_UpperCAmelCase    :						Optional[Any]							=      ""
								else:
												_UpperCAmelCase    :						Dict							=      "vit."
								# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
								_UpperCAmelCase    :						List[str]							=      state_dict.pop(F'blocks.{i}.attn.qkv.weight'	)
								_UpperCAmelCase    :						Tuple							=      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    :						Dict							=      in_proj_bias[: config.hidden_size]
								_UpperCAmelCase    :						Optional[Any]							=      in_proj_weight[
								    config.hidden_size : config.hidden_size * 2, :
								]
								_UpperCAmelCase    :						Optional[int]							=      in_proj_bias[
								    config.hidden_size : config.hidden_size * 2
								]
								_UpperCAmelCase    :						Optional[int]							=      in_proj_weight[
								    -config.hidden_size :, :
								]
								_UpperCAmelCase    :						int							=      in_proj_bias[-config.hidden_size :]

def 							__SCREAMING_SNAKE_CASE							(			lowerCAmelCase:					int	)				->    Optional[int]:
				_UpperCAmelCase    :						Dict							=      ["head.weight", "head.bias"]
				for k in ignore_keys:
								state_dict.pop(lowerCAmelCase					, lowerCAmelCase	)

def 							__SCREAMING_SNAKE_CASE							(			lowerCAmelCase:					Any					, lowerCAmelCase:					Optional[int]					, lowerCAmelCase:					Dict	)				->    Tuple:
				_UpperCAmelCase    :						str							=      dct.pop(lowerCAmelCase	)
				_UpperCAmelCase    :						Any							=      val

def 							__SCREAMING_SNAKE_CASE							(			)				->    List[Any]:
				_UpperCAmelCase    :						Tuple							=      "http://images.cocodataset.org/val2017/000000039769.jpg"
				_UpperCAmelCase    :						Tuple							=      Image.open(requests.get(lowerCAmelCase					, stream=lowerCAmelCase	).raw	)
				return im

@torch.no_grad()
def 							__SCREAMING_SNAKE_CASE							(			lowerCAmelCase:					Tuple					, lowerCAmelCase:					int					, lowerCAmelCase:					List[Any]=False	)				->    Any:
				_UpperCAmelCase    :						List[Any]							=      BitConfig(
				    global_padding="same"					, layer_type="bottleneck"					, depths=(3, 4, 9)					, out_features=["stage3"]					, embedding_dynamic_padding=lowerCAmelCase					, )
				_UpperCAmelCase    :						Optional[Any]							=      ViTHybridConfig(backbone_config=lowerCAmelCase					, image_size=384					, num_labels=1000	)
				_UpperCAmelCase    :						str							=      False

				# load original model from timm
				_UpperCAmelCase    :						Optional[Any]							=      timm.create_model(lowerCAmelCase					, pretrained=lowerCAmelCase	)
				timm_model.eval()

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

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

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

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

				_UpperCAmelCase    :						Tuple							=      {
				    "bilinear": PILImageResampling.BILINEAR,
				    "bicubic": PILImageResampling.BICUBIC,
				    "nearest": PILImageResampling.NEAREST,
				}

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

				_UpperCAmelCase    :						List[str]							=      prepare_img()
				_UpperCAmelCase    :						List[Any]							=      transform(lowerCAmelCase	).unsqueeze(0	)
				_UpperCAmelCase    :						Optional[Any]							=      processor(lowerCAmelCase					, return_tensors="pt"	).pixel_values

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

				# verify logits
				with torch.no_grad():
								_UpperCAmelCase    :						Union[str, Any]							=      model(lowerCAmelCase	)
								_UpperCAmelCase    :						Optional[Any]							=      outputs.logits

				print("Predicted class:"					, logits.argmax(-1	).item()	)
				if base_model:
								_UpperCAmelCase    :						List[Any]							=      timm_model.forward_features(lowerCAmelCase	)
								assert timm_pooled_output.shape == outputs.pooler_output.shape
								assert torch.allclose(lowerCAmelCase					, outputs.pooler_output					, atol=1E-3	)
				else:
								_UpperCAmelCase    :						Any							=      timm_model(lowerCAmelCase	)
								assert timm_logits.shape == outputs.logits.shape
								assert torch.allclose(lowerCAmelCase					, outputs.logits					, atol=1E-3	)
				print("Looks ok!"	)

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

				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__":
					SCREAMING_SNAKE_CASE_            = 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.'
					)

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


