


'''simple docstring'''
# Copyright (c) 2021-, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

####################################################################################################

#
# Note: If when running this conversion script you're getting an exception:
#     ModuleNotFoundError: No module named 'megatron.model.enums'
# you need to tell python where to find the clone of Megatron-LM, e.g.:
#
# cd /tmp
# git clone https://github.com/NVIDIA/Megatron-LM
# PYTHONPATH=/tmp/Megatron-LM python src/transformers/models/megatron_gpt2/convert_megatron_gpt2_checkpoint.py ...
#
# if you already have it cloned elsewhere, simply adjust the path to the existing path
#
# If the training was done using a Megatron-LM fork, e.g.,
# https://github.com/microsoft/Megatron-DeepSpeed/ then chances are that you need to have that one
# in your path, i.e., /path/to/Megatron-DeepSpeed/
#

import argparse
import os
import re
import zipfile

import torch

from transformers import AutoTokenizer, GPTaConfig


def 			_lowercase       (				UpperCamelCase__       :					Union[str, Any],				UpperCamelCase__       :					Dict,				UpperCamelCase__       :					Optional[Any]=0			):
 # Format the message.
 if name is None:
  __A  :		Dict       =			None
 else:
  __A  :		List[Any]       =			'.' * max(0,				spaces - 2			) + '# {:' + str(50 - spaces			) + 's}'
  __A  :		List[Any]       =			fmt.format(UpperCamelCase__			)

 # Print and recurse (if needed).
 if isinstance(UpperCamelCase__,				UpperCamelCase__			):
  if msg is not None:
   print(UpperCamelCase__			)
  for k in val.keys():
   recursive_print(UpperCamelCase__,				val[k],				spaces + 2			)
 elif isinstance(UpperCamelCase__,				torch.Tensor			):
  print(UpperCamelCase__,				':',				val.size()			)
 else:
  print(UpperCamelCase__,				':',				UpperCamelCase__			)


def 			_lowercase       (				UpperCamelCase__       :					Optional[Any],				UpperCamelCase__       :					Any,				UpperCamelCase__       :					List[Any],				UpperCamelCase__       :					List[str],				UpperCamelCase__       :					Optional[Any]			):
 # Permutes layout of param tensor to [num_splits * num_heads * hidden_size, :]
 # for compatibility with later versions of NVIDIA Megatron-LM.
 # The inverse operation is performed inside Megatron-LM to read checkpoints:
 # https://github.com/NVIDIA/Megatron-LM/blob/v2.4/megatron/checkpointing.py#L209
 # If param is the weight tensor of the self-attention block, the returned tensor
 # will have to be transposed one more time to be read by HuggingFace GPT2.
 __A  :		Optional[int]       =			param.size()
 if checkpoint_version == 1.0:
  # version 1.0 stores [num_heads * hidden_size * num_splits, :]
  __A  :		List[Any]       =			(num_heads, hidden_size, num_splits) + input_shape[1:]
  __A  :		Union[str, Any]       =			param.view(*UpperCamelCase__			)
  __A  :		int       =			param.transpose(0,				2			)
  __A  :		Dict       =			param.transpose(1,				2			).contiguous()
 elif checkpoint_version >= 2.0:
  # other versions store [num_heads * num_splits * hidden_size, :]
  __A  :		Tuple       =			(num_heads, num_splits, hidden_size) + input_shape[1:]
  __A  :		List[Any]       =			param.view(*UpperCamelCase__			)
  __A  :		List[str]       =			param.transpose(0,				1			).contiguous()
 __A  :		Dict       =			param.view(*UpperCamelCase__			)
 return param


def 			_lowercase       (				UpperCamelCase__       :					Optional[int],				UpperCamelCase__       :					int,				UpperCamelCase__       :					Union[str, Any]			):
 # The converted output model.
 __A  :		Any       =			{}

 # old versions did not store training args
 __A  :		Optional[Any]       =			input_state_dict.get('args',				UpperCamelCase__			)
 if ds_args is not None:
  # do not make the user write a config file when the exact dimensions/sizes are already in the checkpoint
  # from pprint import pprint
  # pprint(vars(ds_args))

  __A  :		str       =			ds_args.padded_vocab_size
  __A  :		int       =			ds_args.max_position_embeddings
  __A  :		int       =			ds_args.hidden_size
  __A  :		Dict       =			ds_args.num_layers
  __A  :		Union[str, Any]       =			ds_args.num_attention_heads
  __A  :		str       =			ds_args.ffn_hidden_size
  # pprint(config)

 # The number of heads.
 __A  :		List[str]       =			config.n_head
 # The hidden_size per head.
 __A  :		List[str]       =			config.n_embd // config.n_head
 # Megatron-LM checkpoint version
 if "checkpoint_version" in input_state_dict.keys():
  __A  :		Union[str, Any]       =			input_state_dict['checkpoint_version']
 else:
  __A  :		Dict       =			0.0

 # The model.
 __A  :		List[str]       =			input_state_dict['model']
 # The language model.
 __A  :		Dict       =			model['language_model']
 # The embeddings.
 __A  :		Tuple       =			lm['embedding']

 # The word embeddings.
 __A  :		int       =			embeddings['word_embeddings']['weight']
 # Truncate the embedding table to vocab_size rows.
 __A  :		List[str]       =			word_embeddings[: config.vocab_size, :]
 __A  :		Optional[int]       =			word_embeddings

 # The position embeddings.
 __A  :		Dict       =			embeddings['position_embeddings']['weight']
 # Read the causal mask dimension (seqlen). [max_sequence_length, hidden_size]
 __A  :		Optional[Any]       =			pos_embeddings.size(0			)
 if n_positions != config.n_positions:
  raise ValueError(
      f"""pos_embeddings.max_sequence_length={n_positions} and config.n_positions={config.n_positions} don't match"""			)
 # Store the position embeddings.
 __A  :		int       =			pos_embeddings

 # The transformer.
 __A  :		str       =			lm['transformer'] if 'transformer' in lm.keys() else lm['encoder']

 # The regex to extract layer names.
 __A  :		str       =			re.compile(r'layers\.(\d+)\.([a-z0-9_.]+)\.([a-z]+)'			)

 # The simple map of names for "automated" rules.
 __A  :		Optional[Any]       =			{
     'attention.dense': '.attn.c_proj.',
     'self_attention.dense': '.attn.c_proj.',
     'mlp.dense_h_to_4h': '.mlp.c_fc.',
     'mlp.dense_4h_to_h': '.mlp.c_proj.',
 }

 # Extract the layers.
 for key, val in transformer.items():
  # Match the name.
  __A  :		int       =			layer_re.match(UpperCamelCase__			)

  # Stop if that's not a layer
  if m is None:
   break

  # The index of the layer.
  __A  :		List[Any]       =			int(m.group(1			)			)
  # The name of the operation.
  __A  :		Optional[int]       =			m.group(2			)
  # Is it a weight or a bias?
  __A  :		Union[str, Any]       =			m.group(3			)

  # The name of the layer.
  __A  :		str       =			f"""transformer.h.{layer_idx}"""

  # For layernorm(s), simply store the layer norm.
  if op_name.endswith('layernorm'			):
   __A  :		str       =			'ln_1' if op_name.startswith('input'			) else 'ln_2'
   __A  :		List[Any]       =			val

  # Transpose the QKV matrix.
  elif (
      op_name == "attention.query_key_value" or op_name == "self_attention.query_key_value"
  ) and weight_or_bias == "weight":
   # Insert a tensor of 1x1xDxD bias.
   __A  :		str       =			torch.tril(torch.ones((n_positions, n_positions),				dtype=torch.floataa			)			).view(
       1,				1,				UpperCamelCase__,				UpperCamelCase__			)
   __A  :		Union[str, Any]       =			causal_mask

   # Insert a "dummy" tensor for masked_bias.
   __A  :		Union[str, Any]       =			torch.tensor(-1E4,				dtype=torch.floataa			)
   __A  :		Any       =			masked_bias

   __A  :		str       =			fix_query_key_value_ordering(UpperCamelCase__,				UpperCamelCase__,				3,				UpperCamelCase__,				UpperCamelCase__			)
   # Megatron stores (3*D) x D but transformers-GPT2 expects D x 3*D.
   __A  :		str       =			out_val.transpose(0,				1			).contiguous()
   # Store.
   __A  :		str       =			out_val

  # Transpose the bias.
  elif (
      op_name == "attention.query_key_value" or op_name == "self_attention.query_key_value"
  ) and weight_or_bias == "bias":
   __A  :		Dict       =			fix_query_key_value_ordering(UpperCamelCase__,				UpperCamelCase__,				3,				UpperCamelCase__,				UpperCamelCase__			)
   # Store. No change of shape.
   __A  :		List[Any]       =			out_val

  # Transpose the weights.
  elif weight_or_bias == "weight":
   __A  :		List[Any]       =			megatron_to_transformers[op_name]
   __A  :		List[Any]       =			val.transpose(0,				1			)

  # Copy the bias.
  elif weight_or_bias == "bias":
   __A  :		Tuple       =			megatron_to_transformers[op_name]
   __A  :		Union[str, Any]       =			val

    # DEBUG.
 assert config.n_layer == layer_idx + 1

 # The final layernorm.
 __A  :		List[Any]       =			transformer['final_layernorm.weight']
 __A  :		Any       =			transformer['final_layernorm.bias']

 # For LM head, transformers' wants the matrix to weight embeddings.
 __A  :		Optional[Any]       =			word_embeddings

 # It should be done!
 return output_state_dict


def 			_lowercase       (				):
 # Create the argument parser.
 __A  :		List[str]       =			argparse.ArgumentParser()
 parser.add_argument('--print-checkpoint-structure',				action='store_true'			)
 parser.add_argument(
     'path_to_checkpoint',				type=UpperCamelCase__,				help='Path to the checkpoint file (.zip archive or direct .pt file)',				)
 parser.add_argument(
     '--config_file',				default='',				type=UpperCamelCase__,				help='An optional config json file describing the pre-trained model.',				)
 __A  :		Optional[Any]       =			parser.parse_args()

 # Extract the basename.
 __A  :		List[Any]       =			os.path.dirname(args.path_to_checkpoint			)

 # Load the model.
 # the .zip is very optional, let's keep it for backward compatibility
 print(f"""Extracting PyTorch state dictionary from {args.path_to_checkpoint}"""			)
 if args.path_to_checkpoint.endswith('.zip'			):
  with zipfile.ZipFile(args.path_to_checkpoint,				'r'			) as checkpoint:
   with checkpoint.open('release/mp_rank_00/model_optim_rng.pt'			) as pytorch_dict:
    __A  :		Union[str, Any]       =			torch.load(UpperCamelCase__,				map_location='cpu'			)
 else:
  __A  :		int       =			torch.load(args.path_to_checkpoint,				map_location='cpu'			)

 __A  :		List[Any]       =			input_state_dict.get('args',				UpperCamelCase__			)

 # Read the config, or default to the model released by NVIDIA.
 if args.config_file == "":
  if ds_args is not None:
   if ds_args.bias_gelu_fusion:
    __A  :		List[str]       =			'gelu_fast'
   elif ds_args.openai_gelu:
    __A  :		List[Any]       =			'gelu_new'
   else:
    __A  :		int       =			'gelu'
  else:
   # in the very early days this used to be "gelu_new"
   __A  :		Optional[int]       =			'gelu_new'

  # Spell out all parameters in case the defaults change.
  __A  :		str       =			GPTaConfig(
      vocab_size=50257,				n_positions=1024,				n_embd=1024,				n_layer=24,				n_head=16,				n_inner=4096,				activation_function=UpperCamelCase__,				resid_pdrop=0.1,				embd_pdrop=0.1,				attn_pdrop=0.1,				layer_norm_epsilon=1E-5,				initializer_range=0.02,				summary_type='cls_index',				summary_use_proj=UpperCamelCase__,				summary_activation=UpperCamelCase__,				summary_proj_to_labels=UpperCamelCase__,				summary_first_dropout=0.1,				scale_attn_weights=UpperCamelCase__,				use_cache=UpperCamelCase__,				bos_token_id=50256,				eos_token_id=50256,				)
 else:
  __A  :		str       =			GPTaConfig.from_json_file(args.config_file			)

 __A  :		Any       =			['GPT2LMHeadModel']

 # Convert.
 print('Converting'			)
 __A  :		Union[str, Any]       =			convert_megatron_checkpoint(UpperCamelCase__,				UpperCamelCase__,				UpperCamelCase__			)

 # Print the structure of converted state dict.
 if args.print_checkpoint_structure:
  recursive_print(UpperCamelCase__,				UpperCamelCase__			)

 # Add tokenizer class info to config
 # see https://github.com/huggingface/transformers/issues/13906)
 if ds_args is not None:
  __A  :		Dict       =			ds_args.tokenizer_type
  if tokenizer_type == "GPT2BPETokenizer":
   __A  :		Any       =			'gpt2'
  elif tokenizer_type == "PretrainedFromHF":
   __A  :		List[Any]       =			ds_args.tokenizer_name_or_path
  else:
   raise ValueError(f"""Unrecognized tokenizer_type {tokenizer_type}"""			)
 else:
  __A  :		Tuple       =			'gpt2'

 __A  :		Union[str, Any]       =			AutoTokenizer.from_pretrained(UpperCamelCase__			)
 __A  :		Optional[Any]       =			type(UpperCamelCase__			).__name__
 __A  :		int       =			tokenizer_class

 # Store the config to file.
 print('Saving config'			)
 config.save_pretrained(UpperCamelCase__			)

 # Save tokenizer based on args
 print(f"""Adding {tokenizer_class} tokenizer files"""			)
 tokenizer.save_pretrained(UpperCamelCase__			)

 # Store the state_dict to file.
 __A  :		str       =			os.path.join(UpperCamelCase__,				'pytorch_model.bin'			)
 print(f"""Saving checkpoint to \"{output_checkpoint_file}\""""			)
 torch.save(UpperCamelCase__,				UpperCamelCase__			)


####################################################################################################

if __name__ == "__main__":
 main()

####################################################################################################










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


def 			_lowercase       (				UpperCamelCase__       :					Optional[Any],				UpperCamelCase__       :					Tuple=False			):
 __A  :		List[str]       =			[]

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


def 			_lowercase       (				UpperCamelCase__       :					Optional[Any]			):
 __A  :		Any       =			['head.weight', 'head.bias']
 for k in ignore_keys:
  state_dict.pop(UpperCamelCase__,				UpperCamelCase__			)


def 			_lowercase       (				UpperCamelCase__       :					Any,				UpperCamelCase__       :					Any,				UpperCamelCase__       :					Union[str, Any]			):
 __A  :		Dict       =			dct.pop(UpperCamelCase__			)
 __A  :		Optional[Any]       =			val


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


@torch.no_grad()
def 			_lowercase       (				UpperCamelCase__       :					Tuple,				UpperCamelCase__       :					int,				UpperCamelCase__       :					Any=False			):
 __A  :		Optional[Any]       =			BitConfig(
     global_padding='same',				layer_type='bottleneck',				depths=(3, 4, 9),				out_features=['stage3'],				embedding_dynamic_padding=UpperCamelCase__,				)
 __A  :		str       =			ViTHybridConfig(backbone_config=UpperCamelCase__,				image_size=384,				num_labels=1000			)
 __A  :		Union[str, Any]       =			False

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

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

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

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

 # create image processor
 __A  :		Tuple       =			create_transform(**resolve_data_config({},				model=UpperCamelCase__			)			)
 __A  :		Union[str, Any]       =			transform.transforms

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

 __A  :		List[str]       =			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(),				)

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

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

 # verify logits
 with torch.no_grad():
  __A  :		Union[str, Any]       =			model(UpperCamelCase__			)
  __A  :		Dict       =			outputs.logits

 print('Predicted class:',				logits.argmax(-1			).item()			)
 if base_model:
  __A  :		Tuple       =			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:
  __A  :		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__":
 UpperCAmelCase_ :     Optional[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.'
 )

 UpperCAmelCase_ :     Union[str, Any]								=      parser.parse_args()
 convert_vit_checkpoint(args.vit_name, args.pytorch_dump_folder_path, args.push_to_hub)






