





from __future__ import annotations

import unittest

from transformers import RoFormerConfig, 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 import (
           TFRoFormerForCausalLM,
           TFRoFormerForMaskedLM,
           TFRoFormerForMultipleChoice,
           TFRoFormerForQuestionAnswering,
           TFRoFormerForSequenceClassification,
           TFRoFormerForTokenClassification,
           TFRoFormerModel,
       )
       from transformers.models.roformer.modeling_tf_roformer import (
           TFRoFormerSelfAttention,
           TFRoFormerSinusoidalPositionalEmbedding,
       )




class     SCREAMING_SNAKE_CASE__ :



 def __init__(							self,__lowerCamelCase,__lowerCamelCase=13,__lowerCamelCase=7,__lowerCamelCase=True,__lowerCamelCase=True,__lowerCamelCase=True,__lowerCamelCase=True,__lowerCamelCase=99,__lowerCamelCase=32,__lowerCamelCase=2,__lowerCamelCase=4,__lowerCamelCase=37,__lowerCamelCase="gelu",__lowerCamelCase=0.1,__lowerCamelCase=0.1,__lowerCamelCase=512,__lowerCamelCase=16,__lowerCamelCase=2,__lowerCamelCase=0.02,__lowerCamelCase=3,__lowerCamelCase=4,__lowerCamelCase=None,):
        A__											=						parent
        A__											=						13
        A__											=						7
        A__											=						True
        A__											=						True
        A__											=						True
        A__											=						True
        A__											=						99
        A__											=						32
        A__											=						2
        A__											=						4
        A__											=						37
        A__											=						'''gelu'''
        A__											=						0.1
        A__											=						0.1
        A__											=						512
        A__											=						16
        A__											=						2
        A__											=						0.02
        A__											=						3
        A__											=						4
        A__											=						None



 def 				UpperCamelCase			(							self ):
        A__											=						ids_tensor([self.batch_size, self.seq_length],self.vocab_size )

        A__											=						None
        if self.use_input_mask:
               A__											=						random_attention_mask([self.batch_size, self.seq_length] )

        A__											=						None
        if self.use_token_type_ids:
               A__											=						ids_tensor([self.batch_size, self.seq_length],self.type_vocab_size )

        A__											=						None
        A__											=						None
        A__											=						None
        if self.use_labels:
               A__											=						ids_tensor([self.batch_size],self.type_sequence_label_size )
               A__											=						ids_tensor([self.batch_size, self.seq_length],self.num_labels )
               A__											=						ids_tensor([self.batch_size],self.num_choices )

        A__											=						RoFormerConfig(
            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,type_vocab_size=self.type_vocab_size,initializer_range=self.initializer_range,return_dict=__lowerCamelCase,)

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



 def 				UpperCamelCase			(							self,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase ):
        A__											=						TFRoFormerModel(config=__lowerCamelCase )
        A__											=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}

        A__											=						[input_ids, input_mask]
        A__											=						model(__lowerCamelCase )

        A__											=						model(__lowerCamelCase )

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



 def 				UpperCamelCase			(							self,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase ):
        A__											=						True
        A__											=						TFRoFormerForCausalLM(config=__lowerCamelCase )
        A__											=						{
            '''input_ids''': input_ids,
            '''attention_mask''': input_mask,
            '''token_type_ids''': token_type_ids,
        }
        A__											=						model(__lowerCamelCase )['''logits''']
        self.parent.assertListEqual(
            list(prediction_scores.numpy().shape ),[self.batch_size, self.seq_length, self.vocab_size] )



 def 				UpperCamelCase			(							self,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase ):
        A__											=						TFRoFormerForMaskedLM(config=__lowerCamelCase )
        A__											=						{
            '''input_ids''': input_ids,
            '''attention_mask''': input_mask,
            '''token_type_ids''': token_type_ids,
        }
        A__											=						model(__lowerCamelCase )
        self.parent.assertEqual(result.logits.shape,(self.batch_size, self.seq_length, self.vocab_size) )



 def 				UpperCamelCase			(							self,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase ):
        A__											=						self.num_labels
        A__											=						TFRoFormerForSequenceClassification(config=__lowerCamelCase )
        A__											=						{
            '''input_ids''': input_ids,
            '''attention_mask''': input_mask,
            '''token_type_ids''': token_type_ids,
        }

        A__											=						model(__lowerCamelCase )
        self.parent.assertEqual(result.logits.shape,(self.batch_size, self.num_labels) )



 def 				UpperCamelCase			(							self,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase ):
        A__											=						self.num_choices
        A__											=						TFRoFormerForMultipleChoice(config=__lowerCamelCase )
        A__											=						tf.tile(tf.expand_dims(__lowerCamelCase,1 ),(1, self.num_choices, 1) )
        A__											=						tf.tile(tf.expand_dims(__lowerCamelCase,1 ),(1, self.num_choices, 1) )
        A__											=						tf.tile(tf.expand_dims(__lowerCamelCase,1 ),(1, self.num_choices, 1) )
        A__											=						{
            '''input_ids''': multiple_choice_inputs_ids,
            '''attention_mask''': multiple_choice_input_mask,
            '''token_type_ids''': multiple_choice_token_type_ids,
        }
        A__											=						model(__lowerCamelCase )
        self.parent.assertEqual(result.logits.shape,(self.batch_size, self.num_choices) )



 def 				UpperCamelCase			(							self,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase ):
        A__											=						self.num_labels
        A__											=						TFRoFormerForTokenClassification(config=__lowerCamelCase )
        A__											=						{
            '''input_ids''': input_ids,
            '''attention_mask''': input_mask,
            '''token_type_ids''': token_type_ids,
        }
        A__											=						model(__lowerCamelCase )
        self.parent.assertEqual(result.logits.shape,(self.batch_size, self.seq_length, self.num_labels) )



 def 				UpperCamelCase			(							self,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase ):
        A__											=						TFRoFormerForQuestionAnswering(config=__lowerCamelCase )
        A__											=						{
            '''input_ids''': input_ids,
            '''attention_mask''': input_mask,
            '''token_type_ids''': token_type_ids,
        }

        A__											=						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 				UpperCamelCase			(							self ):
        A__											=						self.prepare_config_and_inputs()
        (
            (
            A__
        )					,	(
            A__
        )					,	(
            A__
        )					,	(
            A__
        )					,	(
            A__
        )					,	(
            A__
        )					,	(
            A__
        )					,	
        )											=						config_and_inputs
        A__											=						{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
        return config, inputs_dict




@require_tf
class     SCREAMING_SNAKE_CASE__ (						UpperCamelCase__				,					UpperCamelCase__				,					unittest.TestCase      ):
 __SCREAMING_SNAKE_CASE							       =	(
     (
         TFRoFormerModel,
         TFRoFormerForCausalLM,
         TFRoFormerForMaskedLM,
         TFRoFormerForQuestionAnswering,
         TFRoFormerForSequenceClassification,
         TFRoFormerForTokenClassification,
         TFRoFormerForMultipleChoice,
     )
     if is_tf_available()
     else ()
 )
 __SCREAMING_SNAKE_CASE							       =	(
     {
         '''feature-extraction''': TFRoFormerModel,
         '''fill-mask''': TFRoFormerForMaskedLM,
         '''question-answering''': TFRoFormerForQuestionAnswering,
         '''text-classification''': TFRoFormerForSequenceClassification,
         '''text-generation''': TFRoFormerForCausalLM,
         '''token-classification''': TFRoFormerForTokenClassification,
         '''zero-shot''': TFRoFormerForSequenceClassification,
     }
     if is_tf_available()
     else {}
 )

 __SCREAMING_SNAKE_CASE							       =	False
 __SCREAMING_SNAKE_CASE							       =	False



 def 				UpperCamelCase			(							self,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase ):
        if pipeline_test_casse_name == "TextGenerationPipelineTests":
               return True

        return False



 def 				UpperCamelCase			(							self ):
        A__											=						TFRoFormerModelTester(self )
        A__											=						ConfigTester(self,config_class=__lowerCamelCase,hidden_size=37 )



 def 				UpperCamelCase			(							self ):
        self.config_tester.run_common_tests()



 def 				UpperCamelCase			(							self ):
        A__											=						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__lowerCamelCase )



 def 				UpperCamelCase			(							self ):
        A__											=						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*__lowerCamelCase )



 def 				UpperCamelCase			(							self ):
        A__											=						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_lm_head(*__lowerCamelCase )



 def 				UpperCamelCase			(							self ):
        A__											=						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_multiple_choice(*__lowerCamelCase )



 def 				UpperCamelCase			(							self ):
        A__											=						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_question_answering(*__lowerCamelCase )



 def 				UpperCamelCase			(							self ):
        A__											=						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_sequence_classification(*__lowerCamelCase )



 def 				UpperCamelCase			(							self ):
        A__											=						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*__lowerCamelCase )




 @slow
 def 				UpperCamelCase			(							self ):
        A__											=						TFRoFormerModel.from_pretrained('''junnyu/roformer_chinese_base''' )
        self.assertIsNotNone(__lowerCamelCase )




@require_tf
class     SCREAMING_SNAKE_CASE__ (						unittest.TestCase      ):



 @slow
 def 				UpperCamelCase			(							self ):
        A__											=						TFRoFormerForMaskedLM.from_pretrained('''junnyu/roformer_chinese_base''' )
        A__											=						tf.constant([[0, 1, 2, 3, 4, 5]] )
        A__											=						model(__lowerCamelCase )[0]

        # TODO Replace vocab size
        A__											=						5_0000

        A__											=						[1, 6, vocab_size]
        self.assertEqual(output.shape,__lowerCamelCase )

        print(output[:, :3, :3] )

        # TODO Replace values below with what was printed above.
        A__											=						tf.constant(
            [
                [
                    [-0.12053341, -1.0264901, 0.29221946],
                    [-1.5133783, 0.197433, 0.15190607],
                    [-5.0135403, -3.900256, -0.84038764],
                ]
            ] )
        tf.debugging.assert_near(output[:, :3, :3],__lowerCamelCase,atol=1E-4 )




@require_tf
class     SCREAMING_SNAKE_CASE__ (						unittest.TestCase      ):
 __SCREAMING_SNAKE_CASE							       =	1E-4



 def 				UpperCamelCase			(							self ):
        A__											=						tf.constant([[4, 10]] )
        A__											=						TFRoFormerSinusoidalPositionalEmbedding(num_positions=6,embedding_dim=6 )

        A__											=						emba(input_ids.shape )
        A__											=						tf.constant(
            [[0.0000, 0.0000, 0.0000, 1.0000, 1.0000, 1.0000], [0.8415, 0.0464, 0.0022, 0.5403, 0.9989, 1.0000]] )

        tf.debugging.assert_near(__lowerCamelCase,__lowerCamelCase,atol=self.tolerance )




 def 				UpperCamelCase			(							self ):
        A__											=						tf.constant(
            [
                [0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
                [0.8415, 0.8219, 0.8020, 0.7819, 0.7617],
                [0.9093, 0.9364, 0.9581, 0.9749, 0.9870],
            ] )
        A__											=						TFRoFormerSinusoidalPositionalEmbedding(num_positions=512,embedding_dim=512 )
        emba([2, 16, 512] )
        A__											=						emba.weight[:3, :5]

        tf.debugging.assert_near(__lowerCamelCase,__lowerCamelCase,atol=self.tolerance )




@require_tf
class     SCREAMING_SNAKE_CASE__ (						unittest.TestCase      ):
 __SCREAMING_SNAKE_CASE							       =	1E-4



 def 				UpperCamelCase			(							self ):
        # 2,12,16,64
        A__											=						tf.reshape(tf.range(2 * 12 * 16 * 64,dtype=tf.floataa ),shape=(2, 12, 16, 64) ) / 100

        A__											=						-tf.reshape(tf.range(2 * 12 * 16 * 64,dtype=tf.floataa ),shape=(2, 12, 16, 64) ) / 100

        A__											=						TFRoFormerSinusoidalPositionalEmbedding(num_positions=32,embedding_dim=64 )
        A__											=						embed_positions([2, 16, 768] )[None, None, :, :]

        A__					,	A__											=						TFRoFormerSelfAttention.apply_rotary_position_embeddings(
            __lowerCamelCase,__lowerCamelCase,__lowerCamelCase )

        A__											=						tf.constant(
            [
                [0.0000, 0.0100, 0.0200, 0.0300, 0.0400, 0.0500, 0.0600, 0.0700],
                [-0.2012, 0.8897, 0.0263, 0.9401, 0.2074, 0.9463, 0.3481, 0.9343],
                [-1.7057, 0.6271, -1.2145, 1.3897, -0.6303, 1.7647, -0.1173, 1.8985],
                [-2.1731, -1.6397, -2.7358, 0.2854, -2.1840, 1.7183, -1.3018, 2.4871],
                [0.2717, -3.6173, -2.9206, -2.1988, -3.6638, 0.3858, -2.9155, 2.2980],
                [3.9859, -2.1580, -0.7984, -4.4904, -4.1181, -2.0252, -4.4782, 1.1253],
            ] )
        A__											=						tf.constant(
            [
                [0.0000, -0.0100, -0.0200, -0.0300, -0.0400, -0.0500, -0.0600, -0.0700],
                [0.2012, -0.8897, -0.0263, -0.9401, -0.2074, -0.9463, -0.3481, -0.9343],
                [1.7057, -0.6271, 1.2145, -1.3897, 0.6303, -1.7647, 0.1173, -1.8985],
                [2.1731, 1.6397, 2.7358, -0.2854, 2.1840, -1.7183, 1.3018, -2.4871],
                [-0.2717, 3.6173, 2.9206, 2.1988, 3.6638, -0.3858, 2.9155, -2.2980],
                [-3.9859, 2.1580, 0.7984, 4.4904, 4.1181, 2.0252, 4.4782, -1.1253],
            ] )

        tf.debugging.assert_near(query_layer[0, 0, :6, :8],__lowerCamelCase,atol=self.tolerance )
        tf.debugging.assert_near(key_layer[0, 0, :6, :8],__lowerCamelCase,atol=self.tolerance )












import argparse

import torch

from transformers import BertConfig, BertForPreTraining, load_tf_weights_in_bert
from transformers.utils import logging


logging.set_verbosity_info()


def 	UpperCamelCase__(       UpperCamelCase__			:       Any ,     UpperCamelCase__			:       int ,     UpperCamelCase__			:       int  )->int:
       # Initialise PyTorch model
       A__											=						BertConfig.from_json_file(UpperCamelCase__  )
       print(f"Building PyTorch model from configuration: {config}"  )
       A__											=						BertForPreTraining(UpperCamelCase__  )

       # Load weights from tf checkpoint
       load_tf_weights_in_bert(UpperCamelCase__ ,     UpperCamelCase__ ,     UpperCamelCase__  )

       # Save pytorch-model
       print(f"Save PyTorch model to {pytorch_dump_path}"  )
       torch.save(model.state_dict() ,     UpperCamelCase__  )


if __name__ == "__main__":
       a__:      str									=		argparse.ArgumentParser()
       # Required parameters
       parser.add_argument(
           '--tf_checkpoint_path', default=None, type=str, required=True, help='Path to the TensorFlow checkpoint path.'
       )
       parser.add_argument(
           '--bert_config_file',
           default=None,
           type=str,
           required=True,
           help=(
               'The config json file corresponding to the pre-trained BERT model. \n'
               'This specifies the model architecture.'
           ),
       )
       parser.add_argument(
           '--pytorch_dump_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
       )
       a__:      Any									=		parser.parse_args()
       convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)





