






'''simple docstring'''

from __future__ import annotations

import unittest

from transformers import 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 numpy
   import tensorflow as tf

   from transformers import (
       TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
       TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
       TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
       BertConfig,
       DPRConfig,
       TFDPRContextEncoder,
       TFDPRQuestionEncoder,
       TFDPRReader,
   )







class 	SCREAMING_SNAKE_CASE	:




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


            """simple docstring"""







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




       def 					a_					(					self      : List[str]					):


            """simple docstring"""







            __lowerCamelCase						:       Tuple				   =				ids_tensor([self.batch_size, self.seq_length]	,    self.vocab_size					)

            __lowerCamelCase						:       Dict				   =				None
            if self.use_input_mask:
                 # follow test_modeling_tf_ctrl.py
                 __lowerCamelCase						:       List[Any]				   =				random_attention_mask([self.batch_size, self.seq_length]					)

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

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

            __lowerCamelCase						:       Optional[Any]				   =				BertConfig(
                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	,    is_decoder=A__	,    initializer_range=self.initializer_range	,    )
            __lowerCamelCase						:       Optional[Any]				   =				DPRConfig(projection_dim=self.projection_dim	,    **config.to_dict()					)

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




       def 					a_					(					self      : Tuple	,    A__      : Optional[Any]	,    A__      : str	,    A__      : Any	,    A__      : List[str]	,    A__      : str	,    A__      : str	,    A__      : Tuple					):


            """simple docstring"""







            __lowerCamelCase						:       int				   =				TFDPRContextEncoder(config=A__					)
            __lowerCamelCase						:       str				   =				model(A__	,    attention_mask=A__	,    token_type_ids=A__					)
            __lowerCamelCase						:       str				   =				model(A__	,    token_type_ids=A__					)
            __lowerCamelCase						:       Dict				   =				model(A__					)
            self.parent.assertEqual(result.pooler_output.shape	,    (self.batch_size, self.projection_dim or self.hidden_size)					)




       def 					a_					(					self      : str	,    A__      : Union[str, Any]	,    A__      : Tuple	,    A__      : Optional[int]	,    A__      : Optional[int]	,    A__      : Union[str, Any]	,    A__      : List[str]	,    A__      : Optional[int]					):


            """simple docstring"""







            __lowerCamelCase						:       Optional[Any]				   =				TFDPRQuestionEncoder(config=A__					)
            __lowerCamelCase						:       str				   =				model(A__	,    attention_mask=A__	,    token_type_ids=A__					)
            __lowerCamelCase						:       Dict				   =				model(A__	,    token_type_ids=A__					)
            __lowerCamelCase						:       str				   =				model(A__					)
            self.parent.assertEqual(result.pooler_output.shape	,    (self.batch_size, self.projection_dim or self.hidden_size)					)




       def 					a_					(					self      : int	,    A__      : Union[str, Any]	,    A__      : Optional[int]	,    A__      : Optional[Any]	,    A__      : List[Any]	,    A__      : Tuple	,    A__      : Optional[Any]	,    A__      : Tuple					):


            """simple docstring"""







            __lowerCamelCase						:       List[str]				   =				TFDPRReader(config=A__					)
            __lowerCamelCase						:       Optional[Any]				   =				model(A__	,    attention_mask=A__					)

            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)					)
            self.parent.assertEqual(result.relevance_logits.shape	,    (self.batch_size,)					)




       def 					a_					(					self      : str					):


            """simple docstring"""







            __lowerCamelCase						:       Optional[int]				   =				self.prepare_config_and_inputs()
            (
                (
                __lowerCamelCase
            )   ,      (
                __lowerCamelCase
            )   ,      (
                __lowerCamelCase
            )   ,      (
                __lowerCamelCase
            )   ,      (
                __lowerCamelCase
            )   ,      (
                __lowerCamelCase
            )   ,      (
                __lowerCamelCase
            )   ,      
            )						:       Union[str, Any]				   =				config_and_inputs
            __lowerCamelCase						:       str				   =				{"""input_ids""": input_ids}
            return config, inputs_dict







@require_tf
class 	SCREAMING_SNAKE_CASE	(      lowerCAmelCase_    ,   lowerCAmelCase_    ,   unittest.TestCase  ):
       snake_case__    :     Optional[int]   		=			(
           (
               TFDPRContextEncoder,
               TFDPRQuestionEncoder,
               TFDPRReader,
           )
           if is_tf_available()
           else ()
       )
       snake_case__    :     int   		=			{'feature-extraction': TFDPRQuestionEncoder} if is_tf_available() else {}

       snake_case__    :     Optional[Any]   		=			False
       snake_case__    :     Any   		=			False
       snake_case__    :     Any   		=			False
       snake_case__    :     List[Any]   		=			False
       snake_case__    :     List[Any]   		=			False




       def 					a_					(					self      : Dict					):


            """simple docstring"""







            __lowerCamelCase						:       Optional[Any]				   =				TFDPRModelTester(self					)
            __lowerCamelCase						:       Optional[Any]				   =				ConfigTester(self	,    config_class=A__	,    hidden_size=37					)




       def 					a_					(					self      : str					):


            """simple docstring"""







            self.config_tester.run_common_tests()




       def 					a_					(					self      : Any					):


            """simple docstring"""







            __lowerCamelCase						:       Tuple				   =				self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_dpr_context_encoder(*A__					)




       def 					a_					(					self      : Dict					):


            """simple docstring"""







            __lowerCamelCase						:       int				   =				self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_dpr_question_encoder(*A__					)




       def 					a_					(					self      : str					):


            """simple docstring"""







            __lowerCamelCase						:       Optional[Any]				   =				self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_dpr_reader(*A__					)




       @slow
       def 					a_					(					self      : str					):


            """simple docstring"""







            for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 __lowerCamelCase						:       Tuple				   =				TFDPRContextEncoder.from_pretrained(A__					)
                 self.assertIsNotNone(A__					)

            for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 __lowerCamelCase						:       Union[str, Any]				   =				TFDPRContextEncoder.from_pretrained(A__					)
                 self.assertIsNotNone(A__					)

            for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 __lowerCamelCase						:       int				   =				TFDPRQuestionEncoder.from_pretrained(A__					)
                 self.assertIsNotNone(A__					)

            for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 __lowerCamelCase						:       Dict				   =				TFDPRReader.from_pretrained(A__					)
                 self.assertIsNotNone(A__					)







@require_tf
class 	SCREAMING_SNAKE_CASE	(      unittest.TestCase  ):




       @slow
       def 					a_					(					self      : List[str]					):


            """simple docstring"""







            __lowerCamelCase						:       str				   =				TFDPRQuestionEncoder.from_pretrained("""facebook/dpr-question_encoder-single-nq-base"""					)

            __lowerCamelCase						:       Optional[Any]				   =				tf.constant(
                [[101, 7592, 1010, 2003, 2026, 3899, 10140, 1029, 102]]					)  # [CLS] hello, is my dog cute? [SEP]
            __lowerCamelCase						:       Tuple				   =				model(A__					)[0]  # embedding shape = (1, 768)
            # compare the actual values for a slice.
            __lowerCamelCase						:       Optional[int]				   =				tf.constant(
                [
                    [
                        0.0323_6253,
                        0.1275_3335,
                        0.1681_8509,
                        0.0027_9786,
                        0.389_6933,
                        0.2426_4945,
                        0.217_8971,
                        -0.0233_5227,
                        -0.0848_1959,
                        -0.1432_4117,
                    ]
                ]					)
            self.assertTrue(numpy.allclose(output[:, :10].numpy()	,    expected_slice.numpy()	,    atol=1e-4					)					)











'''simple docstring'''

import argparse

import torch

from transformers import RemBertConfig, RemBertModel, load_tf_weights_in_rembert
from transformers.utils import logging


logging.set_verbosity_info()
def     __lowercase  (_lowercase,						_lowercase,						_lowercase     )   ->     Optional[Any]:


     """simple docstring"""
     # Initialise PyTorch model
     __lowerCamelCase						:       str				   =				RemBertConfig.from_json_file(_lowercase     )
     print("""Building PyTorch model from configuration: {}""".format(str(_lowercase     )     )     )
     __lowerCamelCase						:       List[Any]				   =				RemBertModel(_lowercase     )

     # Load weights from tf checkpoint
     load_tf_weights_in_rembert(_lowercase,						_lowercase,						_lowercase     )

     # Save pytorch-model
     print("""Save PyTorch model to {}""".format(_lowercase     )     )
     torch.save(model.state_dict(),						_lowercase     )


if __name__ == "__main__":
   UpperCAmelCase__						:Tuple          =			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(
       """--rembert_config_file""",
       default=None,
       type=str,
       required=True,
       help=(
           """The config json file corresponding to the pre-trained RemBERT 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."""
   )
   UpperCAmelCase__						:List[Any]          =			parser.parse_args()
   convert_rembert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.rembert_config_file, args.pytorch_dump_path)



