






"""simple docstring"""
from __future__ import annotations

import unittest

from transformers import DistilBertConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
			import tensorflow as tf

			from transformers.models.distilbert.modeling_tf_distilbert import (
			    TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
			    TFDistilBertForMaskedLM,
			    TFDistilBertForMultipleChoice,
			    TFDistilBertForQuestionAnswering,
			    TFDistilBertForSequenceClassification,
			    TFDistilBertForTokenClassification,
			    TFDistilBertModel,
			)







class      UpperCamelCase_			:





	def __init__( self   :					Optional[Any]      ,      lowerCAmelCase_   :					int      ,      )		->     Tuple:
				UpperCAmelCase_	:					Optional[int]     	=   parent
				UpperCAmelCase_	:					Tuple     	=   13
				UpperCAmelCase_	:					int     	=   7
				UpperCAmelCase_	:					Optional[Any]     	=   True
				UpperCAmelCase_	:					int     	=   True
				UpperCAmelCase_	:					Tuple     	=   False
				UpperCAmelCase_	:					Any     	=   True
				UpperCAmelCase_	:					List[Any]     	=   99
				UpperCAmelCase_	:					Any     	=   32
				UpperCAmelCase_	:					str     	=   2
				UpperCAmelCase_	:					Tuple     	=   4
				UpperCAmelCase_	:					Dict     	=   37
				UpperCAmelCase_	:					Optional[int]     	=   "gelu"
				UpperCAmelCase_	:					List[str]     	=   0.1
				UpperCAmelCase_	:					Tuple     	=   0.1
				UpperCAmelCase_	:					Any     	=   512
				UpperCAmelCase_	:					List[str]     	=   16
				UpperCAmelCase_	:					Optional[Any]     	=   2
				UpperCAmelCase_	:					Union[str, Any]     	=   0.0_2
				UpperCAmelCase_	:					str     	=   3
				UpperCAmelCase_	:					Optional[int]     	=   4
				UpperCAmelCase_	:					Dict     	=   None





	def 	_SCREAMING_SNAKE_CASE ( self   :					int		)		->     List[str]:
				UpperCAmelCase_	:					Tuple     	=   ids_tensor([self.batch_size, self.seq_length]      ,      self.vocab_size		)

				UpperCAmelCase_	:					Optional[Any]     	=   None
				if self.use_input_mask:
							UpperCAmelCase_	:					Optional[int]     	=   random_attention_mask([self.batch_size, self.seq_length]		)

				UpperCAmelCase_	:					Tuple     	=   None
				UpperCAmelCase_	:					List[str]     	=   None
				UpperCAmelCase_	:					Dict     	=   None
				if self.use_labels:
							UpperCAmelCase_	:					Optional[Any]     	=   ids_tensor([self.batch_size]      ,      self.type_sequence_label_size		)
							UpperCAmelCase_	:					Any     	=   ids_tensor([self.batch_size, self.seq_length]      ,      self.num_labels		)
							UpperCAmelCase_	:					Tuple     	=   ids_tensor([self.batch_size]      ,      self.num_choices		)

				UpperCAmelCase_	:					Dict     	=   DistilBertConfig(
				    vocab_size=self.vocab_size      ,      dim=self.hidden_size      ,      n_layers=self.num_hidden_layers      ,      n_heads=self.num_attention_heads      ,      hidden_dim=self.intermediate_size      ,      hidden_act=self.hidden_act      ,      dropout=self.hidden_dropout_prob      ,      attention_dropout=self.attention_probs_dropout_prob      ,      max_position_embeddings=self.max_position_embeddings      ,      initializer_range=self.initializer_range      ,      )

				return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels





	def 	_SCREAMING_SNAKE_CASE ( self   :					Optional[Any]      ,      lowerCAmelCase_   :					Tuple      ,      lowerCAmelCase_   :					Tuple      ,      lowerCAmelCase_   :					List[Any]      ,      lowerCAmelCase_   :					List[Any]      ,      lowerCAmelCase_   :					Tuple      ,      lowerCAmelCase_   :					Any		)		->     str:
				UpperCAmelCase_	:					Optional[Any]     	=   TFDistilBertModel(config=lowerCAmelCase_		)
				UpperCAmelCase_	:					Any     	=   {"input_ids": input_ids, "attention_mask": input_mask}

				UpperCAmelCase_	:					str     	=   model(lowerCAmelCase_		)

				UpperCAmelCase_	:					List[Any]     	=   [input_ids, input_mask]

				UpperCAmelCase_	:					Optional[int]     	=   model(lowerCAmelCase_		)

				self.parent.assertEqual(result.last_hidden_state.shape      ,      (self.batch_size, self.seq_length, self.hidden_size)		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					Tuple      ,      lowerCAmelCase_   :					Optional[Any]      ,      lowerCAmelCase_   :					Optional[int]      ,      lowerCAmelCase_   :					Dict      ,      lowerCAmelCase_   :					Optional[Any]      ,      lowerCAmelCase_   :					str      ,      lowerCAmelCase_   :					Union[str, Any]		)		->     List[str]:
				UpperCAmelCase_	:					Dict     	=   TFDistilBertForMaskedLM(config=lowerCAmelCase_		)
				UpperCAmelCase_	:					int     	=   {"input_ids": input_ids, "attention_mask": input_mask}
				UpperCAmelCase_	:					List[str]     	=   model(lowerCAmelCase_		)
				self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.seq_length, self.vocab_size)		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					Union[str, Any]      ,      lowerCAmelCase_   :					int      ,      lowerCAmelCase_   :					List[str]      ,      lowerCAmelCase_   :					Dict      ,      lowerCAmelCase_   :					int      ,      lowerCAmelCase_   :					List[Any]      ,      lowerCAmelCase_   :					Dict		)		->     Dict:
				UpperCAmelCase_	:					Optional[Any]     	=   TFDistilBertForQuestionAnswering(config=lowerCAmelCase_		)
				UpperCAmelCase_	:					List[Any]     	=   {
				    "input_ids": input_ids,
				    "attention_mask": input_mask,
				}
				UpperCAmelCase_	:					Any     	=   model(lowerCAmelCase_		)
				self.parent.assertEqual(result.start_logits.shape      ,      (self.batch_size, self.seq_length)		)
				self.parent.assertEqual(result.end_logits.shape      ,      (self.batch_size, self.seq_length)		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					int      ,      lowerCAmelCase_   :					Optional[int]      ,      lowerCAmelCase_   :					Optional[int]      ,      lowerCAmelCase_   :					str      ,      lowerCAmelCase_   :					List[Any]      ,      lowerCAmelCase_   :					Dict      ,      lowerCAmelCase_   :					Optional[Any]		)		->     Optional[int]:
				UpperCAmelCase_	:					List[Any]     	=   self.num_labels
				UpperCAmelCase_	:					Tuple     	=   TFDistilBertForSequenceClassification(lowerCAmelCase_		)
				UpperCAmelCase_	:					int     	=   {"input_ids": input_ids, "attention_mask": input_mask}
				UpperCAmelCase_	:					Optional[Any]     	=   model(lowerCAmelCase_		)
				self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.num_labels)		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					Optional[int]      ,      lowerCAmelCase_   :					str      ,      lowerCAmelCase_   :					Optional[int]      ,      lowerCAmelCase_   :					Optional[Any]      ,      lowerCAmelCase_   :					Tuple      ,      lowerCAmelCase_   :					Tuple      ,      lowerCAmelCase_   :					Union[str, Any]		)		->     Optional[Any]:
				UpperCAmelCase_	:					Union[str, Any]     	=   self.num_choices
				UpperCAmelCase_	:					List[Any]     	=   TFDistilBertForMultipleChoice(lowerCAmelCase_		)
				UpperCAmelCase_	:					List[Any]     	=   tf.tile(tf.expand_dims(lowerCAmelCase_      ,      1		)      ,      (1, self.num_choices, 1)		)
				UpperCAmelCase_	:					List[Any]     	=   tf.tile(tf.expand_dims(lowerCAmelCase_      ,      1		)      ,      (1, self.num_choices, 1)		)
				UpperCAmelCase_	:					Any     	=   {
				    "input_ids": multiple_choice_inputs_ids,
				    "attention_mask": multiple_choice_input_mask,
				}
				UpperCAmelCase_	:					List[Any]     	=   model(lowerCAmelCase_		)
				self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.num_choices)		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					Tuple      ,      lowerCAmelCase_   :					Union[str, Any]      ,      lowerCAmelCase_   :					Tuple      ,      lowerCAmelCase_   :					List[Any]      ,      lowerCAmelCase_   :					Any      ,      lowerCAmelCase_   :					Union[str, Any]      ,      lowerCAmelCase_   :					Union[str, Any]		)		->     Tuple:
				UpperCAmelCase_	:					Union[str, Any]     	=   self.num_labels
				UpperCAmelCase_	:					List[str]     	=   TFDistilBertForTokenClassification(lowerCAmelCase_		)
				UpperCAmelCase_	:					int     	=   {"input_ids": input_ids, "attention_mask": input_mask}
				UpperCAmelCase_	:					Optional[Any]     	=   model(lowerCAmelCase_		)
				self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.seq_length, self.num_labels)		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					Any		)		->     Dict:
				UpperCAmelCase_	:					Optional[int]     	=   self.prepare_config_and_inputs()
				((UpperCAmelCase_)	,      (UpperCAmelCase_)	,      (UpperCAmelCase_)	,      (UpperCAmelCase_)	,      (UpperCAmelCase_)	,      (UpperCAmelCase_))	:					str     	=   config_and_inputs
				UpperCAmelCase_	:					List[str]     	=   {"input_ids": input_ids, "attention_mask": input_mask}
				return config, inputs_dict







@require_tf
class      UpperCamelCase_			(__A							,		__A							,		unittest.TestCase	):
	__magic_name__       =							(
	    (
	        TFDistilBertModel,
	        TFDistilBertForMaskedLM,
	        TFDistilBertForQuestionAnswering,
	        TFDistilBertForSequenceClassification,
	        TFDistilBertForTokenClassification,
	        TFDistilBertForMultipleChoice,
	    )
	    if is_tf_available()
	    else None
	)
	__magic_name__       =							(
	    {
	        '''feature-extraction''': TFDistilBertModel,
	        '''fill-mask''': TFDistilBertForMaskedLM,
	        '''question-answering''': TFDistilBertForQuestionAnswering,
	        '''text-classification''': TFDistilBertForSequenceClassification,
	        '''token-classification''': TFDistilBertForTokenClassification,
	        '''zero-shot''': TFDistilBertForSequenceClassification,
	    }
	    if is_tf_available()
	    else {}
	)
	__magic_name__       =							False
	__magic_name__       =							False





	def 	_SCREAMING_SNAKE_CASE ( self   :					List[str]		)		->     Tuple:
				UpperCAmelCase_	:					Any     	=   TFDistilBertModelTester(self		)
				UpperCAmelCase_	:					List[str]     	=   ConfigTester(self      ,      config_class=lowerCAmelCase_      ,      dim=37		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					Optional[Any]		)		->     Optional[int]:
				self.config_tester.run_common_tests()





	def 	_SCREAMING_SNAKE_CASE ( self   :					List[str]		)		->     int:
				UpperCAmelCase_	:					str     	=   self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_distilbert_model(*lowerCAmelCase_		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					int		)		->     Union[str, Any]:
				UpperCAmelCase_	:					List[Any]     	=   self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_distilbert_for_masked_lm(*lowerCAmelCase_		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					int		)		->     Optional[Any]:
				UpperCAmelCase_	:					str     	=   self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_distilbert_for_question_answering(*lowerCAmelCase_		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					Any		)		->     Union[str, Any]:
				UpperCAmelCase_	:					Optional[Any]     	=   self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_distilbert_for_sequence_classification(*lowerCAmelCase_		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					Union[str, Any]		)		->     int:
				UpperCAmelCase_	:					Optional[Any]     	=   self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_distilbert_for_multiple_choice(*lowerCAmelCase_		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					List[str]		)		->     Tuple:
				UpperCAmelCase_	:					str     	=   self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_distilbert_for_token_classification(*lowerCAmelCase_		)





	@slow
	def 	_SCREAMING_SNAKE_CASE ( self   :					List[Any]		)		->     Dict:
				for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]		):
							UpperCAmelCase_	:					Tuple     	=   TFDistilBertModel.from_pretrained(lowerCAmelCase_		)
							self.assertIsNotNone(lowerCAmelCase_		)







@require_tf
class      UpperCamelCase_			(unittest.TestCase	):





	@slow
	def 	_SCREAMING_SNAKE_CASE ( self   :					Dict		)		->     Optional[int]:
				UpperCAmelCase_	:					Optional[int]     	=   TFDistilBertModel.from_pretrained("distilbert-base-uncased"		)
				UpperCAmelCase_	:					Tuple     	=   tf.constant([[0, 1, 2, 3, 4, 5]]		)
				UpperCAmelCase_	:					Dict     	=   model(lowerCAmelCase_		)[0]

				UpperCAmelCase_	:					Union[str, Any]     	=   [1, 6, 768]
				self.assertEqual(output.shape      ,      lowerCAmelCase_		)

				UpperCAmelCase_	:					Dict     	=   tf.constant(
				    [
				        [
				            [0.1_9_2_6_1_8_8_5, -0.1_3_7_3_2_9_5_5, 0.4_1_1_9_7_9_9],
				            [0.2_2_1_5_0_1_5_6, -0.0_7_4_2_2_6_6_1, 0.3_9_0_3_7_2_0_4],
				            [0.2_2_7_5_6_0_1_8, -0.0_8_9_6_4_1_4, 0.3_7_0_1_4_6_7],
				        ]
				    ]		)
				tf.debugging.assert_near(output[:, :3, :3]      ,      lowerCAmelCase_      ,      atol=1e-4		)











"""simple docstring"""
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 							snake_case (  ):
			UpperCAmelCase_	:					str     	=   "http://images.cocodataset.org/val2017/000000039769.jpg"
			UpperCAmelCase_	:					str     	=   Image.open(requests.get(A__      ,stream=A__		).raw		)
			return im

def 							snake_case (  A__		):
			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 							snake_case (  A__      ,A__      ,A__		):
			UpperCAmelCase_	:					Tuple     	=   dct.pop(A__		)
			UpperCAmelCase_	:					Optional[Any]     	=   val

def 							snake_case (  A__		):
			UpperCAmelCase_	:					List[str]     	=   []
			for k in state_dict.keys():
						UpperCAmelCase_	:					Union[str, Any]     	=   k
						if ".pwconv" in k:
									UpperCAmelCase_	:					Dict     	=   k_new.replace(".pwconv"      ,".point_wise_conv"		)
						if ".dwconv" in k:
									UpperCAmelCase_	:					Any     	=   k_new.replace(".dwconv"      ,".depth_wise_conv"		)
						if ".Proj." in k:
									UpperCAmelCase_	:					Dict     	=   k_new.replace(".Proj."      ,".proj."		)
						if "patch_embed" in k_new:
									UpperCAmelCase_	:					Tuple     	=   k_new.replace("patch_embed"      ,"swiftformer.patch_embed.patch_embedding"		)
						if "network" in k_new:
									UpperCAmelCase_	:					List[Any]     	=   k_new.split("."		)
									if ls[2].isdigit():
												UpperCAmelCase_	:					Tuple     	=   "swiftformer.encoder.network." + ls[1] + ".blocks." + ls[2] + "." + ".".join(ls[3:]		)
									else:
												UpperCAmelCase_	:					Optional[Any]     	=   k_new.replace("network"      ,"swiftformer.encoder.network"		)
						rename_keys.append((k, k_new)		)
			return rename_keys

@torch.no_grad()
def 							snake_case (  A__      ,A__      ,A__		):
			UpperCAmelCase_	:					Optional[int]     	=   SwiftFormerConfig()

			# dataset (ImageNet-21k only or also fine-tuned on ImageNet 2012), patch_size and image_size
			UpperCAmelCase_	:					Optional[Any]     	=   10_00
			UpperCAmelCase_	:					str     	=   "huggingface/label-files"
			UpperCAmelCase_	:					str     	=   "imagenet-1k-id2label.json"
			UpperCAmelCase_	:					List[str]     	=   json.load(open(hf_hub_download(A__      ,A__      ,repo_type="dataset"		)      ,"r"		)		)
			UpperCAmelCase_	:					Tuple     	=   {int(A__		): v for k, v in idalabel.items()}
			UpperCAmelCase_	:					List[Any]     	=   idalabel
			UpperCAmelCase_	:					Optional[Any]     	=   {v: k for k, v in idalabel.items()}

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

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

			elif swiftformer_name == "swiftformer_l1":
						UpperCAmelCase_	:					int     	=   [4, 3, 10, 5]
						UpperCAmelCase_	:					Union[str, Any]     	=   [48, 96, 1_92, 3_84]

			elif swiftformer_name == "swiftformer_l3":
						UpperCAmelCase_	:					Dict     	=   [4, 4, 12, 6]
						UpperCAmelCase_	:					Optional[int]     	=   [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"		):
									UpperCAmelCase_	:					List[Any]     	=   torch.hub.load_state_dict_from_url(A__      ,map_location="cpu"      ,check_hash=A__		)
						else:
									UpperCAmelCase_	:					Any     	=   torch.load(A__      ,map_location="cpu"		)
			UpperCAmelCase_	:					List[str]     	=   checkpoint

			UpperCAmelCase_	:					Dict     	=   create_rename_keys(A__		)
			for rename_key_src, rename_key_dest in rename_keys:
						rename_key(A__      ,A__      ,A__		)

			# load HuggingFace model
			UpperCAmelCase_	:					Optional[int]     	=   SwiftFormerForImageClassification(A__		).eval()
			hf_model.load_state_dict(A__		)

			# prepare test inputs
			UpperCAmelCase_	:					Tuple     	=   prepare_img()
			UpperCAmelCase_	:					int     	=   ViTImageProcessor.from_pretrained("preprocessor_config"		)
			UpperCAmelCase_	:					int     	=   processor(images=A__      ,return_tensors="pt"		)

			# compare outputs from both models
			UpperCAmelCase_	:					List[Any]     	=   get_expected_output(A__		)
			UpperCAmelCase_	:					int     	=   hf_model(inputs["pixel_values"]		).logits

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

			Path(A__		).mkdir(exist_ok=A__		)
			print(F"""Saving model {swiftformer_name} to {pytorch_dump_folder_path}"""		)
			hf_model.save_pretrained(A__		)


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)



