




import argparse
import json
from collections import OrderedDict
from pathlib import Path

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

from transformers import PoolFormerConfig, PoolFormerForImageClassification, PoolFormerImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
lowerCamelCase       :	List[str]									=     logging.get_logger(__name__)

def     _SCREAMING_SNAKE_CASE    (							lowercase			: Optional[int]			,							lowercase			: int			,							lowercase			: Dict			,							lowercase			: Tuple  ):


      '''simple docstring'''



      lowerCamelCase_	  =					original_name.split('.'  )[0]
      lowerCamelCase_	  =					key.split('.'  )
      lowerCamelCase_	  =					int(key_list[key_list.index(lowercase  ) - 2]  )
      lowerCamelCase_	  =					int(key_list[key_list.index(lowercase  ) - 1]  )
      lowerCamelCase_	  =					orig_block_num - offset

      lowerCamelCase_	  =					key.replace(f"""{orig_block_num}.{layer_num}.{original_name}"""			,							f"""block.{new_block_num}.{layer_num}.{new_name}"""  )
      return key

def     _SCREAMING_SNAKE_CASE    (							lowercase			: Tuple  ):


      '''simple docstring'''



      lowerCamelCase_	  =					OrderedDict()
      lowerCamelCase_  ,       lowerCamelCase_	  =					0, 0
      for key, value in state_dict.items():
            if key.startswith('network'  ):
                  lowerCamelCase_	  =					key.replace('network'			,							'poolformer.encoder'  )
            if "proj" in key:
                  # Works for the first embedding as well as the internal embedding layers
                  if key.endswith('bias'  ) and "patch_embed" not in key:
                        patch_emb_offset += 1
                  lowerCamelCase_	  =					key[: key.find('proj'  )]
                  lowerCamelCase_	  =					key.replace(lowercase			,							f"""patch_embeddings.{total_embed_found}."""  )
                  lowerCamelCase_	  =					key.replace('proj'			,							'projection'  )
                  if key.endswith('bias'  ):
                        total_embed_found += 1
            if "patch_embeddings" in key:
                  lowerCamelCase_	  =					'poolformer.encoder.' + key
            if "mlp.fc1" in key:
                  lowerCamelCase_	  =					replace_key_with_offset(lowercase			,							lowercase			,							'mlp.fc1'			,							'output.conv1'  )
            if "mlp.fc2" in key:
                  lowerCamelCase_	  =					replace_key_with_offset(lowercase			,							lowercase			,							'mlp.fc2'			,							'output.conv2'  )
            if "norm1" in key:
                  lowerCamelCase_	  =					replace_key_with_offset(lowercase			,							lowercase			,							'norm1'			,							'before_norm'  )
            if "norm2" in key:
                  lowerCamelCase_	  =					replace_key_with_offset(lowercase			,							lowercase			,							'norm2'			,							'after_norm'  )
            if "layer_scale_1" in key:
                  lowerCamelCase_	  =					replace_key_with_offset(lowercase			,							lowercase			,							'layer_scale_1'			,							'layer_scale_1'  )
            if "layer_scale_2" in key:
                  lowerCamelCase_	  =					replace_key_with_offset(lowercase			,							lowercase			,							'layer_scale_2'			,							'layer_scale_2'  )
            if "head" in key:
                  lowerCamelCase_	  =					key.replace('head'			,							'classifier'  )
            lowerCamelCase_	  =					value
      return new_state_dict

def     _SCREAMING_SNAKE_CASE    (							):


      '''simple docstring'''



      lowerCamelCase_	  =					'http://images.cocodataset.org/val2017/000000039769.jpg'
      lowerCamelCase_	  =					Image.open(requests.get(lowercase			,							stream=lowercase  ).raw  )

      return image

@torch.no_grad()
def     _SCREAMING_SNAKE_CASE    (							lowercase			: List[str]			,							lowercase			: str			,							lowercase			: List[str]  ):


      '''simple docstring'''



      lowerCamelCase_	  =					PoolFormerConfig()

      # set attributes based on model_name
      lowerCamelCase_	  =					'huggingface/label-files'
      lowerCamelCase_	  =					model_name[-3:]
      lowerCamelCase_	  =					10_00
      lowerCamelCase_	  =					'imagenet-1k-id2label.json'
      lowerCamelCase_	  =					(1, 10_00)

      # set config attributes
      lowerCamelCase_	  =					json.load(open(hf_hub_download(lowercase			,							lowercase			,							repo_type='dataset'  )			,							'r'  )  )
      lowerCamelCase_	  =					{int(lowercase  ): v for k, v in idalabel.items()}
      lowerCamelCase_	  =					idalabel
      lowerCamelCase_	  =					{v: k for k, v in idalabel.items()}
      if size == "s12":
            lowerCamelCase_	  =					[2, 2, 6, 2]
            lowerCamelCase_	  =					[64, 1_28, 3_20, 5_12]
            lowerCamelCase_	  =					4.0
            lowerCamelCase_	  =					0.9
      elif size == "s24":
            lowerCamelCase_	  =					[4, 4, 12, 4]
            lowerCamelCase_	  =					[64, 1_28, 3_20, 5_12]
            lowerCamelCase_	  =					4.0
            lowerCamelCase_	  =					0.9
      elif size == "s36":
            lowerCamelCase_	  =					[6, 6, 18, 6]
            lowerCamelCase_	  =					[64, 1_28, 3_20, 5_12]
            lowerCamelCase_	  =					4.0
            lowerCamelCase_	  =					1e-6
            lowerCamelCase_	  =					0.9
      elif size == "m36":
            lowerCamelCase_	  =					[6, 6, 18, 6]
            lowerCamelCase_	  =					[96, 1_92, 3_84, 7_68]
            lowerCamelCase_	  =					4.0
            lowerCamelCase_	  =					1e-6
            lowerCamelCase_	  =					0.95
      elif size == "m48":
            lowerCamelCase_	  =					[8, 8, 24, 8]
            lowerCamelCase_	  =					[96, 1_92, 3_84, 7_68]
            lowerCamelCase_	  =					4.0
            lowerCamelCase_	  =					1e-6
            lowerCamelCase_	  =					0.95
      else:
            raise ValueError(f"""Size {size} not supported"""  )

      # load image processor
      lowerCamelCase_	  =					PoolFormerImageProcessor(crop_pct=lowercase  )

      # Prepare image
      lowerCamelCase_	  =					prepare_img()
      lowerCamelCase_	  =					image_processor(images=lowercase			,							return_tensors='pt'  ).pixel_values

      logger.info(f"""Converting model {model_name}..."""  )

      # load original state dict
      lowerCamelCase_	  =					torch.load(lowercase			,							map_location=torch.device('cpu'  )  )

      # rename keys
      lowerCamelCase_	  =					rename_keys(lowercase  )

      # create HuggingFace model and load state dict
      lowerCamelCase_	  =					PoolFormerForImageClassification(lowercase  )
      model.load_state_dict(lowercase  )
      model.eval()

      # Define image processor
      lowerCamelCase_	  =					PoolFormerImageProcessor(crop_pct=lowercase  )
      lowerCamelCase_	  =					image_processor(images=prepare_img()			,							return_tensors='pt'  ).pixel_values

      # forward pass
      lowerCamelCase_	  =					model(lowercase  )
      lowerCamelCase_	  =					outputs.logits

      # define expected logit slices for different models
      if size == "s12":
            lowerCamelCase_	  =					torch.tensor([-0.3045, -0.6758, -0.4869]  )
      elif size == "s24":
            lowerCamelCase_	  =					torch.tensor([0.4402, -0.1374, -0.8045]  )
      elif size == "s36":
            lowerCamelCase_	  =					torch.tensor([-0.6080, -0.5133, -0.5898]  )
      elif size == "m36":
            lowerCamelCase_	  =					torch.tensor([0.3952, 0.2263, -1.2668]  )
      elif size == "m48":
            lowerCamelCase_	  =					torch.tensor([0.1167, -0.0656, -0.3423]  )
      else:
            raise ValueError(f"""Size {size} not supported"""  )

      # verify logits
      assert logits.shape == expected_shape
      assert torch.allclose(logits[0, :3]			,							lowercase			,							atol=1e-2  )

      # finally, save model and image processor
      logger.info(f"""Saving PyTorch model and image processor to {pytorch_dump_folder_path}..."""  )
      Path(lowercase  ).mkdir(exist_ok=lowercase  )
      model.save_pretrained(lowercase  )
      print(f"""Saving image processor to {pytorch_dump_folder_path}"""  )
      image_processor.save_pretrained(lowercase  )


if __name__ == "__main__":
 lowerCamelCase       :	List[Any]									=     argparse.ArgumentParser()

 parser.add_argument(
     "--model_name",
     default="poolformer_s12",
     type=str,
     help="Name of the model you'd like to convert.",
 )
 parser.add_argument(
     "--checkpoint_path", default=None, type=str, help="Path to the original PyTorch checkpoint (.pth file)."
 )
 parser.add_argument(
     "--pytorch_dump_folder_path", default=None, type=str, help="Path to the folder to output PyTorch model."
 )
 lowerCamelCase       :	Union[str, Any]									=     parser.parse_args()
 convert_poolformer_checkpoint(args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path)









from __future__ import annotations

import copy
import tempfile
import unittest

from transformers import CONFIG_MAPPING, AutoConfig, BertConfig, GPTaConfig, TaConfig, TapasConfig, is_tf_available
from transformers.testing_utils import (
    DUMMY_UNKNOWN_IDENTIFIER,
    SMALL_MODEL_IDENTIFIER,
    RequestCounter,
    require_tensorflow_probability,
    require_tf,
    slow,
)

from ..bert.test_modeling_bert import BertModelTester


if is_tf_available():
 from transformers import (
     TFAutoModel,
     TFAutoModelForCausalLM,
     TFAutoModelForMaskedLM,
     TFAutoModelForPreTraining,
     TFAutoModelForQuestionAnswering,
     TFAutoModelForSeqaSeqLM,
     TFAutoModelForSequenceClassification,
     TFAutoModelForTableQuestionAnswering,
     TFAutoModelForTokenClassification,
     TFAutoModelWithLMHead,
     TFBertForMaskedLM,
     TFBertForPreTraining,
     TFBertForQuestionAnswering,
     TFBertForSequenceClassification,
     TFBertModel,
     TFFunnelBaseModel,
     TFFunnelModel,
     TFGPTaLMHeadModel,
     TFRobertaForMaskedLM,
     TFTaForConditionalGeneration,
     TFTapasForQuestionAnswering,
 )
 from transformers.models.auto.modeling_tf_auto import (
     TF_MODEL_FOR_CAUSAL_LM_MAPPING,
     TF_MODEL_FOR_MASKED_LM_MAPPING,
     TF_MODEL_FOR_PRETRAINING_MAPPING,
     TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
     TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
     TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
     TF_MODEL_MAPPING,
 )
 from transformers.models.bert.modeling_tf_bert import TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST
 from transformers.models.gpta.modeling_tf_gpta import TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST
 from transformers.models.ta.modeling_tf_ta import TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST
 from transformers.models.tapas.modeling_tf_tapas import TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST







class 						A(		UpperCamelCase     ):



 '''simple docstring'''


 UpperCamelCase		     =		'''new-model'''


if is_tf_available():







 class 						A(		UpperCamelCase     ):



  '''simple docstring'''


  UpperCamelCase		     =		NewModelConfig







@require_tf
class 						A(		unittest.TestCase     ):



 '''simple docstring'''


 @slow
 def        a__				(       self      :						str						)     ->			Union[str, Any]:

       """simple docstring"""






       lowerCamelCase_	  =					'bert-base-cased'
       lowerCamelCase_	  =					AutoConfig.from_pretrained(A_						)
       self.assertIsNotNone(A_						)
       self.assertIsInstance(A_		,							A_						)

       lowerCamelCase_	  =					TFAutoModel.from_pretrained(A_						)
       self.assertIsNotNone(A_						)
       self.assertIsInstance(A_		,							A_						)

 @slow
 def        a__				(       self      :						List[Any]						)     ->			Dict:

       """simple docstring"""






       lowerCamelCase_	  =					'bert-base-cased'
       lowerCamelCase_	  =					AutoConfig.from_pretrained(A_						)
       self.assertIsNotNone(A_						)
       self.assertIsInstance(A_		,							A_						)

       lowerCamelCase_	  =					TFAutoModelForPreTraining.from_pretrained(A_						)
       self.assertIsNotNone(A_						)
       self.assertIsInstance(A_		,							A_						)

 @slow
 def        a__				(       self      :						Union[str, Any]						)     ->			str:

       """simple docstring"""






       for model_name in TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             lowerCamelCase_	  =					AutoConfig.from_pretrained(A_						)
             self.assertIsNotNone(A_						)
             self.assertIsInstance(A_		,							A_						)

             lowerCamelCase_	  =					TFAutoModelForCausalLM.from_pretrained(A_						)
             lowerCamelCase_  ,       lowerCamelCase_	  =					TFAutoModelForCausalLM.from_pretrained(A_		,							output_loading_info=A_						)
             self.assertIsNotNone(A_						)
             self.assertIsInstance(A_		,							A_						)

 @slow
 def        a__				(       self      :						List[Any]						)     ->			Dict:

       """simple docstring"""






       for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             lowerCamelCase_	  =					AutoConfig.from_pretrained(A_						)
             self.assertIsNotNone(A_						)
             self.assertIsInstance(A_		,							A_						)

             lowerCamelCase_	  =					TFAutoModelWithLMHead.from_pretrained(A_						)
             self.assertIsNotNone(A_						)
             self.assertIsInstance(A_		,							A_						)

 @slow
 def        a__				(       self      :						int						)     ->			str:

       """simple docstring"""






       for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             lowerCamelCase_	  =					AutoConfig.from_pretrained(A_						)
             self.assertIsNotNone(A_						)
             self.assertIsInstance(A_		,							A_						)

             lowerCamelCase_	  =					TFAutoModelForMaskedLM.from_pretrained(A_						)
             lowerCamelCase_  ,       lowerCamelCase_	  =					TFAutoModelForMaskedLM.from_pretrained(A_		,							output_loading_info=A_						)
             self.assertIsNotNone(A_						)
             self.assertIsInstance(A_		,							A_						)

 @slow
 def        a__				(       self      :						Any						)     ->			List[Any]:

       """simple docstring"""






       for model_name in TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             lowerCamelCase_	  =					AutoConfig.from_pretrained(A_						)
             self.assertIsNotNone(A_						)
             self.assertIsInstance(A_		,							A_						)

             lowerCamelCase_	  =					TFAutoModelForSeqaSeqLM.from_pretrained(A_						)
             lowerCamelCase_  ,       lowerCamelCase_	  =					TFAutoModelForSeqaSeqLM.from_pretrained(A_		,							output_loading_info=A_						)
             self.assertIsNotNone(A_						)
             self.assertIsInstance(A_		,							A_						)

 @slow
 def        a__				(       self      :						Tuple						)     ->			str:

       """simple docstring"""






       for model_name in ["bert-base-uncased"]:
             lowerCamelCase_	  =					AutoConfig.from_pretrained(A_						)
             self.assertIsNotNone(A_						)
             self.assertIsInstance(A_		,							A_						)

             lowerCamelCase_	  =					TFAutoModelForSequenceClassification.from_pretrained(A_						)
             self.assertIsNotNone(A_						)
             self.assertIsInstance(A_		,							A_						)

 @slow
 def        a__				(       self      :						List[Any]						)     ->			Any:

       """simple docstring"""






       for model_name in ["bert-base-uncased"]:
             lowerCamelCase_	  =					AutoConfig.from_pretrained(A_						)
             self.assertIsNotNone(A_						)
             self.assertIsInstance(A_		,							A_						)

             lowerCamelCase_	  =					TFAutoModelForQuestionAnswering.from_pretrained(A_						)
             self.assertIsNotNone(A_						)
             self.assertIsInstance(A_		,							A_						)

 @slow
 @require_tensorflow_probability
 def        a__				(       self      :						int						)     ->			Union[str, Any]:

       """simple docstring"""






       for model_name in TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST[5:6]:
             lowerCamelCase_	  =					AutoConfig.from_pretrained(A_						)
             self.assertIsNotNone(A_						)
             self.assertIsInstance(A_		,							A_						)

             lowerCamelCase_	  =					TFAutoModelForTableQuestionAnswering.from_pretrained(A_						)
             lowerCamelCase_  ,       lowerCamelCase_	  =					TFAutoModelForTableQuestionAnswering.from_pretrained(
                 A_		,							output_loading_info=A_						)
             self.assertIsNotNone(A_						)
             self.assertIsInstance(A_		,							A_						)

 def        a__				(       self      :						int						)     ->			Union[str, Any]:

       """simple docstring"""






       lowerCamelCase_	  =					TFAutoModelWithLMHead.from_pretrained(A_						)
       self.assertIsInstance(A_		,							A_						)
       self.assertEqual(model.num_parameters()		,							14410						)
       self.assertEqual(model.num_parameters(only_trainable=A_						)		,							14410						)

 def        a__				(       self      :						Optional[int]						)     ->			List[Any]:

       """simple docstring"""






       lowerCamelCase_	  =					TFAutoModelWithLMHead.from_pretrained(A_						)
       self.assertIsInstance(A_		,							A_						)
       self.assertEqual(model.num_parameters()		,							14410						)
       self.assertEqual(model.num_parameters(only_trainable=A_						)		,							14410						)

 def        a__				(       self      :						Tuple						)     ->			Optional[Any]:

       """simple docstring"""






       lowerCamelCase_	  =					TFAutoModel.from_pretrained('sgugger/funnel-random-tiny'						)
       self.assertIsInstance(A_		,							A_						)

       lowerCamelCase_	  =					copy.deepcopy(model.config						)
       lowerCamelCase_	  =					['FunnelBaseModel']
       lowerCamelCase_	  =					TFAutoModel.from_config(A_						)
       self.assertIsInstance(A_		,							A_						)

       with tempfile.TemporaryDirectory() as tmp_dir:
             model.save_pretrained(A_						)
             lowerCamelCase_	  =					TFAutoModel.from_pretrained(A_						)
             self.assertIsInstance(A_		,							A_						)

 def        a__				(       self      :						Any						)     ->			Tuple:

       """simple docstring"""






       try:
             AutoConfig.register('new-model'		,							A_						)

             lowerCamelCase_	  =					[
                 TFAutoModel,
                 TFAutoModelForCausalLM,
                 TFAutoModelForMaskedLM,
                 TFAutoModelForPreTraining,
                 TFAutoModelForQuestionAnswering,
                 TFAutoModelForSequenceClassification,
                 TFAutoModelForTokenClassification,
             ]

             for auto_class in auto_classes:
                   with self.subTest(auto_class.__name__						):
                         # Wrong config class will raise an error
                         with self.assertRaises(A_						):
                               auto_class.register(A_		,							A_						)
                         auto_class.register(A_		,							A_						)
                         # Trying to register something existing in the Transformers library will raise an error
                         with self.assertRaises(A_						):
                               auto_class.register(A_		,							A_						)

                         # Now that the config is registered, it can be used as any other config with the auto-API
                         lowerCamelCase_	  =					BertModelTester(self						).get_config()
                         lowerCamelCase_	  =					NewModelConfig(**tiny_config.to_dict()						)
                         lowerCamelCase_	  =					auto_class.from_config(A_						)
                         self.assertIsInstance(A_		,							A_						)

                         with tempfile.TemporaryDirectory() as tmp_dir:
                               model.save_pretrained(A_						)
                               lowerCamelCase_	  =					auto_class.from_pretrained(A_						)
                               self.assertIsInstance(A_		,							A_						)

       finally:
             if "new-model" in CONFIG_MAPPING._extra_content:
                   del CONFIG_MAPPING._extra_content["new-model"]
             for mapping in (
                 TF_MODEL_MAPPING,
                 TF_MODEL_FOR_PRETRAINING_MAPPING,
                 TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
                 TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
                 TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
                 TF_MODEL_FOR_CAUSAL_LM_MAPPING,
                 TF_MODEL_FOR_MASKED_LM_MAPPING,
             ):
                   if NewModelConfig in mapping._extra_content:
                         del mapping._extra_content[NewModelConfig]

 def        a__				(       self      :						int						)     ->			int:

       """simple docstring"""






       with self.assertRaisesRegex(
           A_		,							'bert-base is not a local folder and is not a valid model identifier'						):
             lowerCamelCase_	  =					TFAutoModel.from_pretrained('bert-base'						)

 def        a__				(       self      :						Any						)     ->			Dict:

       """simple docstring"""






       with self.assertRaisesRegex(
           A_		,							r'aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'						):
             lowerCamelCase_	  =					TFAutoModel.from_pretrained(A_		,							revision='aaaaaa'						)

 def        a__				(       self      :						str						)     ->			int:

       """simple docstring"""






       with self.assertRaisesRegex(
           A_		,							'hf-internal-testing/config-no-model does not appear to have a file named pytorch_model.bin'		,							):
             lowerCamelCase_	  =					TFAutoModel.from_pretrained('hf-internal-testing/config-no-model'						)

 def        a__				(       self      :						Any						)     ->			List[Any]:

       """simple docstring"""






       with self.assertRaisesRegex(A_		,							'Use `from_pt=True` to load this model'						):
             lowerCamelCase_	  =					TFAutoModel.from_pretrained('hf-internal-testing/tiny-bert-pt-only'						)





 def        a__				(       self      :						Tuple						)     ->			Union[str, Any]:

       """simple docstring"""






       lowerCamelCase_	  =					TFAutoModel.from_pretrained('hf-internal-testing/tiny-random-bert'						)
       with RequestCounter() as counter:
             lowerCamelCase_	  =					TFAutoModel.from_pretrained('hf-internal-testing/tiny-random-bert'						)
             self.assertEqual(counter.get_request_count		,							0						)
             self.assertEqual(counter.head_request_count		,							1						)
             self.assertEqual(counter.other_request_count		,							0						)

       # With a sharded checkpoint
       lowerCamelCase_	  =					TFAutoModel.from_pretrained('ArthurZ/tiny-random-bert-sharded'						)
       with RequestCounter() as counter:
             lowerCamelCase_	  =					TFAutoModel.from_pretrained('ArthurZ/tiny-random-bert-sharded'						)
             self.assertEqual(counter.get_request_count		,							0						)
             self.assertEqual(counter.head_request_count		,							1						)
             self.assertEqual(counter.other_request_count		,							0						)



