





from __future__ import annotations

import unittest

from transformers import MPNetConfig, 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.mpnet.modeling_tf_mpnet import (
     TFMPNetForMaskedLM,
     TFMPNetForMultipleChoice,
     TFMPNetForQuestionAnswering,
     TFMPNetForSequenceClassification,
     TFMPNetForTokenClassification,
     TFMPNetModel,
 )


class 					__magic_name__  :





     """simple docstring"""



     def __init__(							self       :Union[str, Any]			,      snake_case       :Dict			,      snake_case       :Union[str, Any]=13			,      snake_case       :Tuple=7			,      snake_case       :Union[str, Any]=True			,      snake_case       :Any=True			,      snake_case       :Union[str, Any]=False			,      snake_case       :Tuple=True			,      snake_case       :Any=99			,      snake_case       :Any=64			,      snake_case       :int=2			,      snake_case       :str=4			,      snake_case       :List[str]=64			,      snake_case       :Optional[Any]="gelu"			,      snake_case       :int=0.1			,      snake_case       :int=0.1			,      snake_case       :Optional[int]=512			,      snake_case       :Optional[Any]=16			,      snake_case       :Any=2			,      snake_case       :List[str]=0.02			,      snake_case       :Any=3			,      snake_case       :Optional[int]=4			,      snake_case       :Union[str, Any]=None			,      ):





        '''simple docstring'''




        A_					:				int          =			parent
        A_					:				Optional[Any]          =			batch_size
        A_					:				List[str]          =			seq_length
        A_					:				Dict          =			is_training
        A_					:				List[Any]          =			use_input_mask
        A_					:				str          =			use_token_type_ids
        A_					:				Dict          =			use_labels
        A_					:				List[str]          =			vocab_size
        A_					:				List[Any]          =			hidden_size
        A_					:				Dict          =			num_hidden_layers
        A_					:				str          =			num_attention_heads
        A_					:				int          =			intermediate_size
        A_					:				Tuple          =			hidden_act
        A_					:				str          =			hidden_dropout_prob
        A_					:				List[str]          =			attention_probs_dropout_prob
        A_					:				Optional[int]          =			max_position_embeddings
        A_					:				Optional[int]          =			type_vocab_size
        A_					:				Tuple          =			type_sequence_label_size
        A_					:				int          =			initializer_range
        A_					:				Union[str, Any]          =			num_labels
        A_					:				List[Any]          =			num_choices
        A_					:				Optional[int]          =			scope



     def 						SCREAMING_SNAKE_CASE		(							self       :str  ):





        '''simple docstring'''




        A_					:				List[str]          =			ids_tensor([self.batch_size, self.seq_length]			,      self.vocab_size  )

        A_					:				List[str]          =			None
        if self.use_input_mask:
           A_					:				Dict          =			random_attention_mask([self.batch_size, self.seq_length]  )

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

        A_					:				Union[str, Any]          =			MPNetConfig(
            vocab_size=self.vocab_size			,      hidden_size=self.hidden_size			,      num_hidden_layers=self.num_hidden_layers			,      num_attention_heads=self.num_attention_heads			,      intermediate_size=self.intermediate_size			,      hidden_act=self.hidden_act			,      hidden_dropout_prob=self.hidden_dropout_prob			,      attention_probs_dropout_prob=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[int]			,      snake_case       :Any			,      snake_case       :Union[str, Any]			,      snake_case       :List[Any]			,      snake_case       :List[Any]			,      snake_case       :int			,      snake_case       :Any  ):





        '''simple docstring'''




        A_					:				Dict          =			TFMPNetModel(config=snake_case  )
        A_					:				str          =			{"input_ids": input_ids, "attention_mask": input_mask}
        A_					:				str          =			model(snake_case  )
        A_					:				int          =			[input_ids, input_mask]
        A_					:				Optional[int]          =			model(snake_case  )
        self.parent.assertEqual(result.last_hidden_state.shape			,      (self.batch_size, self.seq_length, self.hidden_size)  )



     def 						SCREAMING_SNAKE_CASE		(							self       :str			,      snake_case       :List[str]			,      snake_case       :Dict			,      snake_case       :List[Any]			,      snake_case       :Tuple			,      snake_case       :Union[str, Any]			,      snake_case       :Optional[Any]  ):





        '''simple docstring'''




        A_					:				str          =			TFMPNetForMaskedLM(config=snake_case  )
        A_					:				int          =			{"input_ids": input_ids, "attention_mask": input_mask}
        A_					:				Union[str, Any]          =			model(snake_case  )
        self.parent.assertEqual(result.logits.shape			,      (self.batch_size, self.seq_length, self.vocab_size)  )



     def 						SCREAMING_SNAKE_CASE		(							self       :Tuple			,      snake_case       :Optional[Any]			,      snake_case       :Any			,      snake_case       :Tuple			,      snake_case       :List[str]			,      snake_case       :Tuple			,      snake_case       :Tuple  ):





        '''simple docstring'''




        A_					:				int          =			TFMPNetForQuestionAnswering(config=snake_case  )
        A_					:				List[Any]          =			{
            "input_ids": input_ids,
            "attention_mask": input_mask,
        }
        A_					:				List[Any]          =			model(snake_case  )
        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       :List[str]			,      snake_case       :str			,      snake_case       :Tuple			,      snake_case       :List[str]			,      snake_case       :Dict			,      snake_case       :Any			,      snake_case       :Optional[Any]  ):





        '''simple docstring'''




        A_					:				int          =			self.num_labels
        A_					:				Optional[Any]          =			TFMPNetForSequenceClassification(snake_case  )
        A_					:				str          =			{"input_ids": input_ids, "attention_mask": input_mask}
        A_					:				List[str]          =			model(snake_case  )
        self.parent.assertEqual(result.logits.shape			,      (self.batch_size, self.num_labels)  )



     def 						SCREAMING_SNAKE_CASE		(							self       :Optional[int]			,      snake_case       :Union[str, Any]			,      snake_case       :List[str]			,      snake_case       :str			,      snake_case       :Dict			,      snake_case       :List[str]			,      snake_case       :int  ):





        '''simple docstring'''




        A_					:				Optional[Any]          =			self.num_choices
        A_					:				Optional[Any]          =			TFMPNetForMultipleChoice(snake_case  )
        A_					:				Tuple          =			tf.tile(tf.expand_dims(snake_case			,      1  )			,      (1, self.num_choices, 1)  )
        A_					:				int          =			tf.tile(tf.expand_dims(snake_case			,      1  )			,      (1, self.num_choices, 1)  )
        A_					:				Union[str, Any]          =			{
            "input_ids": multiple_choice_inputs_ids,
            "attention_mask": multiple_choice_input_mask,
        }
        A_					:				Optional[int]          =			model(snake_case  )
        self.parent.assertEqual(result.logits.shape			,      (self.batch_size, self.num_choices)  )



     def 						SCREAMING_SNAKE_CASE		(							self       :List[Any]			,      snake_case       :str			,      snake_case       :List[Any]			,      snake_case       :Union[str, Any]			,      snake_case       :List[str]			,      snake_case       :Any			,      snake_case       :Optional[Any]  ):





        '''simple docstring'''




        A_					:				Optional[Any]          =			self.num_labels
        A_					:				Union[str, Any]          =			TFMPNetForTokenClassification(snake_case  )
        A_					:				Tuple          =			{"input_ids": input_ids, "attention_mask": input_mask}
        A_					:				Tuple          =			model(snake_case  )
        self.parent.assertEqual(result.logits.shape			,      (self.batch_size, self.seq_length, self.num_labels)  )





     def 						SCREAMING_SNAKE_CASE		(							self       :Dict  ):





        '''simple docstring'''




        A_					:				Optional[int]          =			self.prepare_config_and_inputs()
        ((A_)      ,						(A_)      ,						(A_)      ,						(A_)      ,						(A_)      ,						(A_))					:				int          =			config_and_inputs
        A_					:				Dict          =			{"input_ids": input_ids, "attention_mask": input_mask}
        return config, inputs_dict




@require_tf
class 					__magic_name__  (							lowerCamelCase__    ,				lowerCamelCase__    ,				unittest.TestCase							):





     """simple docstring"""



     __UpperCamelCase					    =     (
         (
             TFMPNetForMaskedLM,
             TFMPNetForMultipleChoice,
             TFMPNetForQuestionAnswering,
             TFMPNetForSequenceClassification,
             TFMPNetForTokenClassification,
             TFMPNetModel,
         )
         if is_tf_available()
         else ()
     )
     __UpperCamelCase					    =     (
         {
             '''feature-extraction''': TFMPNetModel,
             '''fill-mask''': TFMPNetForMaskedLM,
             '''question-answering''': TFMPNetForQuestionAnswering,
             '''text-classification''': TFMPNetForSequenceClassification,
             '''token-classification''': TFMPNetForTokenClassification,
             '''zero-shot''': TFMPNetForSequenceClassification,
         }
         if is_tf_available()
         else {}
     )
     __UpperCamelCase					    =     False
     __UpperCamelCase					    =     False



     def 						SCREAMING_SNAKE_CASE		(							self       :Any  ):





        '''simple docstring'''




        A_					:				Optional[Any]          =			TFMPNetModelTester(self  )
        A_					:				Dict          =			ConfigTester(self			,      config_class=snake_case			,      hidden_size=37  )



     def 						SCREAMING_SNAKE_CASE		(							self       :Any  ):





        '''simple docstring'''




        self.config_tester.run_common_tests()



     def 						SCREAMING_SNAKE_CASE		(							self       :List[Any]  ):





        '''simple docstring'''




        A_					:				Tuple          =			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mpnet_model(*snake_case  )



     def 						SCREAMING_SNAKE_CASE		(							self       :List[Any]  ):





        '''simple docstring'''




        A_					:				Dict          =			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mpnet_for_masked_lm(*snake_case  )



     def 						SCREAMING_SNAKE_CASE		(							self       :str  ):





        '''simple docstring'''




        A_					:				str          =			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mpnet_for_question_answering(*snake_case  )



     def 						SCREAMING_SNAKE_CASE		(							self       :int  ):





        '''simple docstring'''




        A_					:				Tuple          =			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mpnet_for_sequence_classification(*snake_case  )



     def 						SCREAMING_SNAKE_CASE		(							self       :Any  ):





        '''simple docstring'''




        A_					:				Optional[int]          =			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mpnet_for_multiple_choice(*snake_case  )



     def 						SCREAMING_SNAKE_CASE		(							self       :Tuple  ):





        '''simple docstring'''




        A_					:				Any          =			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mpnet_for_token_classification(*snake_case  )





     @slow
     def 						SCREAMING_SNAKE_CASE		(							self       :Optional[Any]  ):





        '''simple docstring'''




        for model_name in ["microsoft/mpnet-base"]:
           A_					:				Optional[Any]          =			TFMPNetModel.from_pretrained(snake_case  )
           self.assertIsNotNone(snake_case  )




@require_tf
class 					__magic_name__  (							unittest.TestCase							):





     """simple docstring"""



     @slow
     def 						SCREAMING_SNAKE_CASE		(							self       :Dict  ):





        '''simple docstring'''




        A_					:				Tuple          =			TFMPNetModel.from_pretrained("microsoft/mpnet-base"  )
        A_					:				str          =			tf.constant([[0, 1, 2, 3, 4, 5]]  )
        A_					:				Any          =			model(snake_case  )[0]

        A_					:				str          =			[1, 6, 768]
        self.assertEqual(output.shape			,      snake_case  )

        A_					:				int          =			tf.constant(
            [
                [
                    [-0.1067172, 0.08216473, 0.0024543],
                    [-0.03465879, 0.8354118, -0.03252288],
                    [-0.06569476, -0.12424111, -0.0494436],
                ]
            ]  )
        tf.debugging.assert_near(output[:, :3, :3]			,      snake_case			,      atol=1e-4  )







import argparse

import torch

from transformers import ChineseCLIPConfig, ChineseCLIPModel

def      __snake_case						(					_lowerCAmelCase : str				,				_lowerCAmelCase : Optional[int]				,				_lowerCAmelCase : Optional[Any]	) ->						List[str]:
   A_      ,						A_      ,						A_					:				Dict          =			pt_weights[f"{prefix}.in_proj_weight"].chunk(3				,				dim=0	)
   A_      ,						A_      ,						A_					:				Union[str, Any]          =			pt_weights[f"{prefix}.in_proj_bias"].chunk(3				,				dim=0	)

   A_					:				int          =			pt_weights[f"{prefix}.out_proj.weight"]
   A_					:				Dict          =			pt_weights[f"{prefix}.out_proj.bias"]

   A_					:				Optional[int]          =			q_proj
   A_					:				Optional[int]          =			q_proj_bias

   A_					:				Tuple          =			k_proj
   A_					:				str          =			k_proj_bias

   A_					:				Optional[int]          =			v_proj
   A_					:				int          =			v_proj_bias

   A_					:				int          =			out_proj_weights
   A_					:				List[Any]          =			out_proj_bias

def      __snake_case						(					_lowerCAmelCase : int				,				_lowerCAmelCase : Dict				,				_lowerCAmelCase : int	) ->						Dict:
   copy_linear(hf_mlp.fca				,				_lowerCAmelCase				,				f"{prefix}.c_fc"	)
   copy_linear(hf_mlp.fca				,				_lowerCAmelCase				,				f"{prefix}.c_proj"	)

def      __snake_case						(					_lowerCAmelCase : Tuple				,				_lowerCAmelCase : Optional[Any]				,				_lowerCAmelCase : Optional[int]	) ->						List[str]:
   A_					:				List[str]          =			pt_weights[f"{prefix}.weight"].data
   A_					:				int          =			pt_weights[f"{prefix}.bias"].data

def      __snake_case						(					_lowerCAmelCase : Dict				,				_lowerCAmelCase : str				,				_lowerCAmelCase : Optional[Any]	) ->						str:
   # copy layer norms
   copy_linear(hf_layer.layer_norma				,				_lowerCAmelCase				,				f"{prefix}.ln_1"	)
   copy_linear(hf_layer.layer_norma				,				_lowerCAmelCase				,				f"{prefix}.ln_2"	)

   # copy MLP
   copy_mlp(hf_layer.mlp				,				_lowerCAmelCase				,				f"{prefix}.mlp"	)

   # copy attn
   copy_attn_layer(hf_layer.self_attn				,				_lowerCAmelCase				,				f"{prefix}.attn"	)

def      __snake_case						(					_lowerCAmelCase : Optional[Any]				,				_lowerCAmelCase : int				,				_lowerCAmelCase : Optional[Any]	) ->						List[Any]:
   for layer_id, hf_layer in enumerate(_lowerCAmelCase	):
      copy_layer(_lowerCAmelCase				,				_lowerCAmelCase				,				f"{prefix}.{layer_id}"	)

def      __snake_case						(					_lowerCAmelCase : Union[str, Any]				,				_lowerCAmelCase : Any	) ->						Dict:
   # copy projection
   A_					:				Optional[Any]          =			pt_weights["text_projection"].data.T

   # copy text encoder
   for name, param in hf_model.text_model.named_parameters():
      A_					:				Any          =			pt_weights[f"bert.{name}"].data

def      __snake_case						(					_lowerCAmelCase : int				,				_lowerCAmelCase : Any	) ->						List[str]:
   # copy projection
   A_					:				Tuple          =			pt_weights["visual.proj"].data.T

   # copy layer norms
   copy_linear(hf_model.vision_model.pre_layrnorm				,				_lowerCAmelCase				,				"visual.ln_pre"	)
   copy_linear(hf_model.vision_model.post_layernorm				,				_lowerCAmelCase				,				"visual.ln_post"	)

   # copy embeddings
   A_					:				List[Any]          =			pt_weights["visual.conv1.weight"].data
   A_					:				Union[str, Any]          =			pt_weights["visual.class_embedding"].data
   A_					:				Tuple          =			pt_weights["visual.positional_embedding"].data

   # copy encoder
   copy_layers(hf_model.vision_model.encoder.layers				,				_lowerCAmelCase				,				"visual.transformer.resblocks"	)

@torch.no_grad()
def      __snake_case						(					_lowerCAmelCase : List[Any]				,				_lowerCAmelCase : int				,				_lowerCAmelCase : Optional[int]=None	) ->						str:
   assert config_path is not None, "Please specify the ChineseCLIP model config of the corresponding model size."
   A_					:				Optional[Any]          =			ChineseCLIPConfig.from_pretrained(_lowerCAmelCase	)

   A_					:				Tuple          =			ChineseCLIPModel(_lowerCAmelCase	).eval()

   A_					:				Any          =			torch.load(_lowerCAmelCase				,				map_location="cpu"	)["state_dict"]
   A_					:				Optional[Any]          =			{(name[7:] if name.startswith("module."	) else name): value for name, value in pt_weights.items()}

   copy_text_model_and_projection(_lowerCAmelCase				,				_lowerCAmelCase	)
   copy_vision_model_and_projection(_lowerCAmelCase				,				_lowerCAmelCase	)
   A_					:				List[Any]          =			pt_weights["logit_scale"].data

   hf_model.save_pretrained(_lowerCAmelCase	)


if __name__ == "__main__":
 _lowerCAmelCase							:      Optional[Any]						=	argparse.ArgumentParser()
 parser.add_argument(
     '''--pytorch_dump_folder_path''',
     default=None,
     type=str,
     help='''Path to the output folder storing converted hf PyTorch model.''',
 )
 parser.add_argument(
     '''--checkpoint_path''', default=None, type=str, help='''Path to original github format ChineseCLIP checkpoint.'''
 )
 parser.add_argument(
     '''--config_path''', default=None, required=True, type=str, help='''Path to hf config.json of model to convert.'''
 )
 _lowerCAmelCase							:      Tuple						=	parser.parse_args()

 convert_chinese_clip_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path)
 print('''The conversion is finished!''')
