





import inspect
import unittest

from transformers import ConvNextConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import ConvNextBackbone, ConvNextForImageClassification, ConvNextModel
  from transformers.models.convnext.modeling_convnext import CONVNEXT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
  from PIL import Image

  from transformers import AutoImageProcessor

class 		__a				:



       def __init__(  self						: Tuple       ,      UpperCAmelCase						: Optional[int]       ,      UpperCAmelCase						: Tuple=13       ,      UpperCAmelCase						: int=32       ,      UpperCAmelCase						: List[str]=3       ,      UpperCAmelCase						: int=4       ,      UpperCAmelCase						: Any=[10, 20, 30, 40]       ,      UpperCAmelCase						: Tuple=[2, 2, 3, 2]       ,      UpperCAmelCase						: List[str]=True       ,      UpperCAmelCase						: Dict=True       ,      UpperCAmelCase						: Any=37       ,      UpperCAmelCase						: Optional[Any]="gelu"       ,      UpperCAmelCase						: Optional[int]=10       ,      UpperCAmelCase						: Any=0.02       ,      UpperCAmelCase						: Union[str, Any]=["stage2", "stage3", "stage4"]       ,      UpperCAmelCase						: Optional[int]=[2, 3, 4]       ,      UpperCAmelCase						: int=None       ,      ):
           lowerCAmelCase_       :						List[Any]							= parent
           lowerCAmelCase_       :						Tuple							= batch_size
           lowerCAmelCase_       :						Tuple							= image_size
           lowerCAmelCase_       :						Any							= num_channels
           lowerCAmelCase_       :						List[Any]							= num_stages
           lowerCAmelCase_       :						Optional[int]							= hidden_sizes
           lowerCAmelCase_       :						Optional[int]							= depths
           lowerCAmelCase_       :						Union[str, Any]							= is_training
           lowerCAmelCase_       :						Optional[int]							= use_labels
           lowerCAmelCase_       :						Dict							= intermediate_size
           lowerCAmelCase_       :						List[Any]							= hidden_act
           lowerCAmelCase_       :						Optional[int]							= num_labels
           lowerCAmelCase_       :						Any							= initializer_range
           lowerCAmelCase_       :						str							= out_features
           lowerCAmelCase_       :						Dict							= out_indices
           lowerCAmelCase_       :						Dict							= scope



       def 			A				(  self						: Optional[int]      ):
           lowerCAmelCase_       :						Optional[Any]							= floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

           lowerCAmelCase_       :						Tuple							= None
           if self.use_labels:
               lowerCAmelCase_       :						int							= ids_tensor([self.batch_size]       ,      self.num_labels      )

           lowerCAmelCase_       :						Dict							= self.get_config()
           return config, pixel_values, labels



       def 			A				(  self						: List[str]      ):
           return ConvNextConfig(
               num_channels=self.num_channels       ,      hidden_sizes=self.hidden_sizes       ,      depths=self.depths       ,      num_stages=self.num_stages       ,      hidden_act=self.hidden_act       ,      is_decoder=lowerCamelCase__       ,      initializer_range=self.initializer_range       ,      out_features=self.out_features       ,      out_indices=self.out_indices       ,      num_labels=self.num_labels       ,      )



       def 			A				(  self						: Dict       ,      UpperCAmelCase						: List[str]       ,      UpperCAmelCase						: Tuple       ,      UpperCAmelCase						: Tuple      ):
           lowerCAmelCase_       :						Tuple							= ConvNextModel(config=lowerCamelCase__      )
           model.to(lowerCamelCase__      )
           model.eval()
           lowerCAmelCase_       :						List[Any]							= model(lowerCamelCase__      )
           # expected last hidden states: B, C, H // 32, W // 32
           self.parent.assertEqual(
               result.last_hidden_state.shape       ,      (self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)       ,      )



       def 			A				(  self						: Tuple       ,      UpperCAmelCase						: Union[str, Any]       ,      UpperCAmelCase						: Optional[Any]       ,      UpperCAmelCase						: List[str]      ):
           lowerCAmelCase_       :						List[str]							= ConvNextForImageClassification(lowerCamelCase__      )
           model.to(lowerCamelCase__      )
           model.eval()
           lowerCAmelCase_       :						int							= model(lowerCamelCase__       ,      labels=lowerCamelCase__      )
           self.parent.assertEqual(result.logits.shape       ,      (self.batch_size, self.num_labels)      )



       def 			A				(  self						: int       ,      UpperCAmelCase						: Tuple       ,      UpperCAmelCase						: List[Any]       ,      UpperCAmelCase						: Optional[int]      ):
           lowerCAmelCase_       :						Optional[int]							= ConvNextBackbone(config=lowerCamelCase__      )
           model.to(lowerCamelCase__      )
           model.eval()
           lowerCAmelCase_       :						Union[str, Any]							= model(lowerCamelCase__      )

           # verify hidden states
           self.parent.assertEqual(len(result.feature_maps      )       ,      len(config.out_features      )      )
           self.parent.assertListEqual(list(result.feature_maps[0].shape      )       ,      [self.batch_size, self.hidden_sizes[1], 4, 4]      )

           # verify channels
           self.parent.assertEqual(len(model.channels      )       ,      len(config.out_features      )      )
           self.parent.assertListEqual(model.channels       ,      config.hidden_sizes[1:]      )

           # verify backbone works with out_features=None
           lowerCAmelCase_       :						Optional[int]							= None
           lowerCAmelCase_       :						List[str]							= ConvNextBackbone(config=lowerCamelCase__      )
           model.to(lowerCamelCase__      )
           model.eval()
           lowerCAmelCase_       :						int							= model(lowerCamelCase__      )

           # verify feature maps
           self.parent.assertEqual(len(result.feature_maps      )       ,      1      )
           self.parent.assertListEqual(list(result.feature_maps[0].shape      )       ,      [self.batch_size, self.hidden_sizes[-1], 1, 1]      )

           # verify channels
           self.parent.assertEqual(len(model.channels      )       ,      1      )
           self.parent.assertListEqual(model.channels       ,      [config.hidden_sizes[-1]]      )



       def 			A				(  self						: List[Any]      ):
           lowerCAmelCase_       :						Tuple							= self.prepare_config_and_inputs()
           lowerCAmelCase_       :						int							= config_and_inputs
           lowerCAmelCase_       :						List[str]							= {'pixel_values': pixel_values}
           return config, inputs_dict






@require_torch
class 		__a				(						__UpperCamelCase      ,__UpperCamelCase      ,unittest.TestCase   ):
       __snake_case    : List[str]   =		(
           (
               ConvNextModel,
               ConvNextForImageClassification,
               ConvNextBackbone,
           )
           if is_torch_available()
           else ()
       )
       __snake_case    : Any   =		(
           {"""feature-extraction""": ConvNextModel, """image-classification""": ConvNextForImageClassification}
           if is_torch_available()
           else {}
       )

       __snake_case    : Optional[Any]   =		True
       __snake_case    : Tuple   =		False
       __snake_case    : Optional[Any]   =		False
       __snake_case    : Optional[Any]   =		False
       __snake_case    : Any   =		False



       def 			A				(  self						: Any      ):
           lowerCAmelCase_       :						Optional[int]							= ConvNextModelTester(self      )
           lowerCAmelCase_       :						Union[str, Any]							= ConfigTester(self       ,      config_class=lowerCamelCase__       ,      has_text_modality=lowerCamelCase__       ,      hidden_size=37      )



       def 			A				(  self						: int      ):
           self.create_and_test_config_common_properties()
           self.config_tester.create_and_test_config_to_json_string()
           self.config_tester.create_and_test_config_to_json_file()
           self.config_tester.create_and_test_config_from_and_save_pretrained()
           self.config_tester.create_and_test_config_with_num_labels()
           self.config_tester.check_config_can_be_init_without_params()
           self.config_tester.check_config_arguments_init()



       def 			A				(  self						: int      ):
           return



       @unittest.skip(reason="""ConvNext does not use inputs_embeds"""      )
       def 			A				(  self						: Tuple      ):
           pass



       @unittest.skip(reason="""ConvNext does not support input and output embeddings"""      )
       def 			A				(  self						: Optional[int]      ):
           pass



       @unittest.skip(reason="""ConvNext does not use feedforward chunking"""      )
       def 			A				(  self						: Optional[Any]      ):
           pass



       def 			A				(  self						: int      ):
           lowerCAmelCase_       :						int							= self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
               lowerCAmelCase_       :						Dict							= model_class(lowerCamelCase__      )
               lowerCAmelCase_       :						int							= inspect.signature(model.forward      )
               # signature.parameters is an OrderedDict => so arg_names order is deterministic
               lowerCAmelCase_       :						Any							= [*signature.parameters.keys()]

               lowerCAmelCase_       :						Union[str, Any]							= ['pixel_values']
               self.assertListEqual(arg_names[:1]       ,      lowerCamelCase__      )



       def 			A				(  self						: str      ):
           lowerCAmelCase_       :						Optional[Any]							= self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*lowerCamelCase__      )



       def 			A				(  self						: Optional[int]      ):
           lowerCAmelCase_       :						int							= self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_backbone(*lowerCamelCase__      )



       def 			A				(  self						: Any      ):
           def check_hidden_states_output(UpperCAmelCase						: Optional[int]       ,      UpperCAmelCase						: Any       ,      UpperCAmelCase						: int      ):
               lowerCAmelCase_       :						Tuple							= model_class(lowerCamelCase__      )
               model.to(lowerCamelCase__      )
               model.eval()

               with torch.no_grad():
                   lowerCAmelCase_       :						Tuple							= model(**self._prepare_for_class(lowerCamelCase__       ,      lowerCamelCase__      )      )

               lowerCAmelCase_       :						List[str]							= outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

               lowerCAmelCase_       :						str							= self.model_tester.num_stages
               self.assertEqual(len(lowerCamelCase__      )       ,      expected_num_stages + 1      )

               # ConvNext's feature maps are of shape (batch_size, num_channels, height, width)
               self.assertListEqual(
                   list(hidden_states[0].shape[-2:]      )       ,      [self.model_tester.image_size // 4, self.model_tester.image_size // 4]       ,      )

           lowerCAmelCase_       :						int							= self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
               lowerCAmelCase_       :						List[str]							= True
               check_hidden_states_output(lowerCamelCase__       ,      lowerCamelCase__       ,      lowerCamelCase__      )

               # check that output_hidden_states also work using config
               del inputs_dict["output_hidden_states"]
               lowerCAmelCase_       :						Optional[Any]							= True

               check_hidden_states_output(lowerCamelCase__       ,      lowerCamelCase__       ,      lowerCamelCase__      )



       def 			A				(  self						: str      ):
           lowerCAmelCase_       :						str							= self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_image_classification(*lowerCamelCase__      )



       @slow
       def 			A				(  self						: Optional[int]      ):
           for model_name in CONVNEXT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               lowerCAmelCase_       :						str							= ConvNextModel.from_pretrained(lowerCamelCase__      )
               self.assertIsNotNone(lowerCamelCase__      )






def       __UpperCamelCase					(      )						->   Union[str, Any]:
    '''simple docstring'''





    lowerCAmelCase_       :						Any							= Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""		)
    return image

@require_torch
@require_vision
class 		__a				(						unittest.TestCase   ):



       @cached_property
       def 			A				(  self						: str      ):
           return AutoImageProcessor.from_pretrained("""facebook/convnext-tiny-224"""      ) if is_vision_available() else None



       @slow
       def 			A				(  self						: Tuple      ):
           lowerCAmelCase_       :						List[Any]							= ConvNextForImageClassification.from_pretrained("""facebook/convnext-tiny-224"""      ).to(lowerCamelCase__      )

           lowerCAmelCase_       :						str							= self.default_image_processor
           lowerCAmelCase_       :						int							= prepare_img()
           lowerCAmelCase_       :						int							= image_processor(images=lowerCamelCase__       ,      return_tensors="""pt"""      ).to(lowerCamelCase__      )

           # forward pass
           with torch.no_grad():
               lowerCAmelCase_       :						Tuple							= model(**lowerCamelCase__      )

           # verify the logits
           lowerCAmelCase_       :						Optional[Any]							= torch.Size((1, 10_00)      )
           self.assertEqual(outputs.logits.shape       ,      lowerCamelCase__      )

           lowerCAmelCase_       :						Tuple							= torch.tensor([-0.0260, -0.4739, 0.1911]      ).to(lowerCamelCase__      )

           self.assertTrue(torch.allclose(outputs.logits[0, :3]       ,      lowerCamelCase__       ,      atol=1e-4      )      )






@require_torch
class 		__a				(						unittest.TestCase      ,__UpperCamelCase   ):
       __snake_case    : Optional[Any]   =		(ConvNextBackbone,) if is_torch_available() else ()
       __snake_case    : Optional[Any]   =		ConvNextConfig

       __snake_case    : Optional[int]   =		False



       def 			A				(  self						: Optional[Any]      ):
           lowerCAmelCase_       :						Tuple							= ConvNextModelTester(self      )







'''simple docstring'''
import argparse

import torch

from transformers import YosoConfig, YosoForMaskedLM






def A__       (	UpperCAmelCase_				):
						if "model" in orig_key:
												_UpperCamelCase   :						List[Any]												=  orig_key.replace('model.'    ,    ''				)
						if "norm1" in orig_key:
												_UpperCamelCase   :						List[Any]												=  orig_key.replace('norm1'    ,    'attention.output.LayerNorm'				)
						if "norm2" in orig_key:
												_UpperCamelCase   :						int												=  orig_key.replace('norm2'    ,    'output.LayerNorm'				)
						if "norm" in orig_key:
												_UpperCamelCase   :						str												=  orig_key.replace('norm'    ,    'LayerNorm'				)
						if "transformer" in orig_key:
												_UpperCamelCase   :						int												=  orig_key.split('.'				)[0].split('_'				)[-1]
												_UpperCamelCase   :						Tuple												=  orig_key.replace(f'transformer_{layer_num}'    ,    f'encoder.layer.{layer_num}'				)
						if "mha.attn" in orig_key:
												_UpperCamelCase   :						List[Any]												=  orig_key.replace('mha.attn'    ,    'attention.self'				)
						if "mha" in orig_key:
												_UpperCamelCase   :						Union[str, Any]												=  orig_key.replace('mha'    ,    'attention'				)
						if "W_q" in orig_key:
												_UpperCamelCase   :						Union[str, Any]												=  orig_key.replace('W_q'    ,    'self.query'				)
						if "W_k" in orig_key:
												_UpperCamelCase   :						List[Any]												=  orig_key.replace('W_k'    ,    'self.key'				)
						if "W_v" in orig_key:
												_UpperCamelCase   :						List[Any]												=  orig_key.replace('W_v'    ,    'self.value'				)
						if "ff1" in orig_key:
												_UpperCamelCase   :						Optional[Any]												=  orig_key.replace('ff1'    ,    'intermediate.dense'				)
						if "ff2" in orig_key:
												_UpperCamelCase   :						List[str]												=  orig_key.replace('ff2'    ,    'output.dense'				)
						if "ff" in orig_key:
												_UpperCamelCase   :						Union[str, Any]												=  orig_key.replace('ff'    ,    'output.dense'				)
						if "mlm_class" in orig_key:
												_UpperCamelCase   :						Tuple												=  orig_key.replace('mlm.mlm_class'    ,    'cls.predictions.decoder'				)
						if "mlm" in orig_key:
												_UpperCamelCase   :						int												=  orig_key.replace('mlm'    ,    'cls.predictions.transform'				)
						if "cls" not in orig_key:
												_UpperCamelCase   :						Optional[int]												=  'yoso.' + orig_key

						return orig_key






def A__       (	UpperCAmelCase_    ,    UpperCAmelCase_				):
						for key in orig_state_dict.copy().keys():
												_UpperCamelCase   :						Any												=  orig_state_dict.pop(UpperCAmelCase_				)

												if ("pooler" in key) or ("sen_class" in key):
																		continue
												else:
																		_UpperCamelCase   :						Union[str, Any]												=  val

						_UpperCamelCase   :						Any												=  orig_state_dict['cls.predictions.decoder.bias']
						_UpperCamelCase   :						Union[str, Any]												=  torch.arange(UpperCAmelCase_				).expand((1, -1)				) + 2

						return orig_state_dict






def A__       (	UpperCAmelCase_    ,    UpperCAmelCase_    ,    UpperCAmelCase_				):
						_UpperCamelCase   :						List[Any]												=  torch.load(UpperCAmelCase_    ,    map_location='cpu'				)['model_state_dict']
						_UpperCamelCase   :						List[str]												=  YosoConfig.from_json_file(UpperCAmelCase_				)
						_UpperCamelCase   :						Union[str, Any]												=  YosoForMaskedLM(UpperCAmelCase_				)

						_UpperCamelCase   :						Tuple												=  convert_checkpoint_helper(config.max_position_embeddings    ,    UpperCAmelCase_				)

						print(model.load_state_dict(UpperCAmelCase_				)				)
						model.eval()
						model.save_pretrained(UpperCAmelCase_				)

						print(f'Checkpoint successfuly converted. Model saved at {pytorch_dump_path}'				)


if __name__ == "__main__":
					snake_case_					:      int    			=    argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    '--pytorch_model_path', default=None, type=str, required=True, help='Path to YOSO pytorch checkpoint.'
					)
					parser.add_argument(
					    '--config_file',
					    default=None,
					    type=str,
					    required=True,
					    help='The json file for YOSO model config.',
					)
					parser.add_argument(
					    '--pytorch_dump_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
					)
					snake_case_					:      List[Any]    			=    parser.parse_args()
					convert_yoso_checkpoint(args.pytorch_model_path, args.config_file, args.pytorch_dump_path)


