


import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    SwiftFormerConfig,
    SwiftFormerForImageClassification,
    ViTImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
lowerCAmelCase			=				logging.get_logger(__name__)

lowerCAmelCase			=				torch.device('cpu')
def      _a					( ):


		"""simple docstring"""




		lowercase__				= '''http://images.cocodataset.org/val2017/000000039769.jpg'''
		lowercase__				= Image.open(requests.get(SCREAMING_SNAKE_CASE		,							stream=SCREAMING_SNAKE_CASE	).raw	)
		return im
def      _a					( SCREAMING_SNAKE_CASE	):


		"""simple docstring"""




		if swiftformer_name == "swiftformer_xs":
				return torch.tensor([-2.1703E00, 2.1107E00, -2.0811E00, 8.8685E-01, 2.4360E-01]	)

		elif swiftformer_name == "swiftformer_s":
				return torch.tensor([3.9636E-01, 2.3478E-01, -1.6963E00, -1.7381E00, -8.6337E-01]	)

		elif swiftformer_name == "swiftformer_l1":
				return torch.tensor([-4.2768E-01, -4.7429E-01, -1.0897E00, -1.0248E00, 3.5523E-02]	)

		elif swiftformer_name == "swiftformer_l3":
				return torch.tensor([-2.5330E-01, 2.4211E-01, -6.0185E-01, -8.2789E-01, -6.0446E-02]	)
def      _a					( SCREAMING_SNAKE_CASE		,							SCREAMING_SNAKE_CASE		,							SCREAMING_SNAKE_CASE	):


		"""simple docstring"""




		lowercase__				= dct.pop(SCREAMING_SNAKE_CASE	)
		lowercase__				= val
def      _a					( SCREAMING_SNAKE_CASE	):


		"""simple docstring"""




		lowercase__				= []
		for k in state_dict.keys():
				lowercase__				= k
				if ".pwconv" in k:
						lowercase__				= k_new.replace('''.pwconv'''		,							'''.point_wise_conv'''	)
				if ".dwconv" in k:
						lowercase__				= k_new.replace('''.dwconv'''		,							'''.depth_wise_conv'''	)
				if ".Proj." in k:
						lowercase__				= k_new.replace('''.Proj.'''		,							'''.proj.'''	)
				if "patch_embed" in k_new:
						lowercase__				= k_new.replace('''patch_embed'''		,							'''swiftformer.patch_embed.patch_embedding'''	)
				if "network" in k_new:
						lowercase__				= k_new.split('''.'''	)
						if ls[2].isdigit():
								lowercase__				= '''swiftformer.encoder.network.''' + ls[1] + '''.blocks.''' + ls[2] + '''.''' + '''.'''.join(ls[3:]	)
						else:
								lowercase__				= k_new.replace('''network'''		,							'''swiftformer.encoder.network'''	)
				rename_keys.append((k, k_new)	)
		return rename_keys
@torch.no_grad()
def      _a					( SCREAMING_SNAKE_CASE		,							SCREAMING_SNAKE_CASE		,							SCREAMING_SNAKE_CASE	):


		"""simple docstring"""




		lowercase__				= SwiftFormerConfig()

		# dataset (ImageNet-21k only or also fine-tuned on ImageNet 2012), patch_size and image_size
		lowercase__				= 10_00
		lowercase__				= '''huggingface/label-files'''
		lowercase__				= '''imagenet-1k-id2label.json'''
		lowercase__				= json.load(open(hf_hub_download(SCREAMING_SNAKE_CASE		,							SCREAMING_SNAKE_CASE		,							repo_type='''dataset'''	)		,							'''r'''	)	)
		lowercase__				= {int(SCREAMING_SNAKE_CASE	): v for k, v in idalabel.items()}
		lowercase__				= idalabel
		lowercase__				= {v: k for k, v in idalabel.items()}

		# size of the architecture
		if swiftformer_name == "swiftformer_xs":
				lowercase__				= [3, 3, 6, 4]
				lowercase__				= [48, 56, 1_12, 2_20]

		elif swiftformer_name == "swiftformer_s":
				lowercase__				= [3, 3, 9, 6]
				lowercase__				= [48, 64, 1_68, 2_24]

		elif swiftformer_name == "swiftformer_l1":
				lowercase__				= [4, 3, 10, 5]
				lowercase__				= [48, 96, 1_92, 3_84]

		elif swiftformer_name == "swiftformer_l3":
				lowercase__				= [4, 4, 12, 6]
				lowercase__				= [64, 1_28, 3_20, 5_12]

		# load state_dict of original model, remove and rename some keys
		if original_ckpt:
				if original_ckpt.startswith('''https'''	):
						lowercase__				= torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE		,							map_location='''cpu'''		,							check_hash=SCREAMING_SNAKE_CASE	)
				else:
						lowercase__				= torch.load(SCREAMING_SNAKE_CASE		,							map_location='''cpu'''	)
		lowercase__				= checkpoint

		lowercase__				= create_rename_keys(SCREAMING_SNAKE_CASE	)
		for rename_key_src, rename_key_dest in rename_keys:
				rename_key(SCREAMING_SNAKE_CASE		,							SCREAMING_SNAKE_CASE		,							SCREAMING_SNAKE_CASE	)

		# load HuggingFace model
		lowercase__				= SwiftFormerForImageClassification(SCREAMING_SNAKE_CASE	).eval()
		hf_model.load_state_dict(SCREAMING_SNAKE_CASE	)

		# prepare test inputs
		lowercase__				= prepare_img()
		lowercase__				= ViTImageProcessor.from_pretrained('''preprocessor_config'''	)
		lowercase__				= processor(images=SCREAMING_SNAKE_CASE		,							return_tensors='''pt'''	)

		# compare outputs from both models
		lowercase__				= get_expected_output(SCREAMING_SNAKE_CASE	)
		lowercase__				= hf_model(inputs['''pixel_values''']	).logits

		assert hf_logits.shape == torch.Size([1, 10_00]	)
		assert torch.allclose(hf_logits[0, 0:5]		,							SCREAMING_SNAKE_CASE		,							atol=1E-3	)

		Path(SCREAMING_SNAKE_CASE	).mkdir(exist_ok=SCREAMING_SNAKE_CASE	)
		print(f'Saving model {swiftformer_name} to {pytorch_dump_folder_path}'	)
		hf_model.save_pretrained(SCREAMING_SNAKE_CASE	)


if __name__ == "__main__":
			lowerCAmelCase			=				argparse.ArgumentParser()
			# Required parameters
			parser.add_argument(
			    '--swiftformer_name',
			    default='swiftformer_xs',
			    choices=['swiftformer_xs', 'swiftformer_s', 'swiftformer_l1', 'swiftformer_l3'],
			    type=str,
			    help='Name of the SwiftFormer model you\'d like to convert.',
			)
			parser.add_argument(
			    '--pytorch_dump_folder_path',
			    default='./converted_outputs/',
			    type=str,
			    help='Path to the output PyTorch model directory.',
			)
			parser.add_argument('--original_ckpt', default=None, type=str, help='Path to the original model checkpoint.')

			lowerCAmelCase			=				parser.parse_args()
			convert_swiftformer_checkpoint(args.swiftformer_name, args.pytorch_dump_folder_path, args.original_ckpt)




"""simple docstring"""

from ...configuration_utils import PretrainedConfig
from ...utils import logging


SCREAMING_SNAKE_CASE_            =    logging.get_logger(__name__)

SCREAMING_SNAKE_CASE_            =    {
    '''tiiuae/falcon-40b''': '''https://huggingface.co/tiiuae/falcon-40b/resolve/main/config.json''',
    '''tiiuae/falcon-7b''': '''https://huggingface.co/tiiuae/falcon-7b/resolve/main/config.json''',
}
class   _UpperCAmelCase				(		SCREAMING_SNAKE_CASE_				):
      __SCREAMING_SNAKE_CASE :				Optional[Any]				   =						"falcon"
      __SCREAMING_SNAKE_CASE :				List[Any]				   =						["past_key_values"]




      def __init__(				self ,							lowercase_=6_5_0_2_4 ,							lowercase_=4_5_4_4 ,							lowercase_=3_2 ,							lowercase_=7_1 ,							lowercase_=1E-5 ,							lowercase_=0.0_2 ,							lowercase_=True ,							lowercase_=0.0 ,							lowercase_=0.0 ,							lowercase_=None ,							lowercase_=False ,							lowercase_=False ,							lowercase_=True ,							lowercase_=True ,							lowercase_=False ,							lowercase_=1_1 ,							lowercase_=1_1 ,							**lowercase_ ,							)    ->       Any:
          UpperCAmelCase          =       vocab_size
          # Backward compatibility with n_embed kwarg
          UpperCAmelCase          =       kwargs.pop('n_embed' ,							lowercase_				)
          UpperCAmelCase          =       hidden_size if n_embed is None else n_embed
          UpperCAmelCase          =       num_hidden_layers
          UpperCAmelCase          =       num_attention_heads
          UpperCAmelCase          =       layer_norm_epsilon
          UpperCAmelCase          =       initializer_range
          UpperCAmelCase          =       use_cache
          UpperCAmelCase          =       hidden_dropout
          UpperCAmelCase          =       attention_dropout

          UpperCAmelCase          =       bos_token_id
          UpperCAmelCase          =       eos_token_id
          UpperCAmelCase          =       num_attention_heads if num_kv_heads is None else num_kv_heads
          UpperCAmelCase          =       alibi
          UpperCAmelCase          =       new_decoder_architecture
          UpperCAmelCase          =       multi_query  # Ignored when new_decoder_architecture is True
          UpperCAmelCase          =       parallel_attn
          UpperCAmelCase          =       bias

          super().__init__(bos_token_id=lowercase_ ,							eos_token_id=lowercase_ ,							**lowercase_				)




      @property
      def      a_   (				self				)    ->       List[Any]:
          return self.hidden_size // self.num_attention_heads




      @property
      def      a_   (				self				)    ->       List[str]:
          return not self.alibi



