





'''simple docstring'''




import unittest

from transformers import SqueezeBertConfig, is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import (
      SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
      SqueezeBertForMaskedLM,
      SqueezeBertForMultipleChoice,
      SqueezeBertForQuestionAnswering,
      SqueezeBertForSequenceClassification,
      SqueezeBertForTokenClassification,
      SqueezeBertModel,
  )


class 					SCREAMING_SNAKE_CASE       (_lowerCAmelCase				):


     def __init__(		self		,      _UpperCAmelCase		,      _UpperCAmelCase=13		,      _UpperCAmelCase=7		,      _UpperCAmelCase=True		,      _UpperCAmelCase=True		,      _UpperCAmelCase=False		,      _UpperCAmelCase=True		,      _UpperCAmelCase=99		,      _UpperCAmelCase=32		,      _UpperCAmelCase=5		,      _UpperCAmelCase=4		,      _UpperCAmelCase=64		,      _UpperCAmelCase="gelu"		,      _UpperCAmelCase=0.1		,      _UpperCAmelCase=0.1		,      _UpperCAmelCase=512		,      _UpperCAmelCase=16		,      _UpperCAmelCase=2		,      _UpperCAmelCase=0.02		,      _UpperCAmelCase=3		,      _UpperCAmelCase=4		,      _UpperCAmelCase=None		,      _UpperCAmelCase=2		,      _UpperCAmelCase=2		,      _UpperCAmelCase=2		,      _UpperCAmelCase=2		,      _UpperCAmelCase=4		,      _UpperCAmelCase=1		,      ):

          '''simple docstring'''






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


     def 				SCREAMING_SNAKE_CASE    (		self):

          '''simple docstring'''






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

          __A		:					Optional[Any]						=			None
          if self.use_input_mask:
               __A		:					Tuple						=			random_attention_mask([self.batch_size, self.seq_length])

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

          __A		:					Optional[int]						=			self.get_config()

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


     def 				SCREAMING_SNAKE_CASE    (		self):

          '''simple docstring'''






          return SqueezeBertConfig(
              embedding_size=self.hidden_size		,      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		,      attention_probs_dropout_prob=self.hidden_dropout_prob		,      attention_dropout=self.attention_probs_dropout_prob		,      max_position_embeddings=self.max_position_embeddings		,      initializer_range=self.initializer_range		,      q_groups=self.q_groups		,      k_groups=self.k_groups		,      v_groups=self.v_groups		,      post_attention_groups=self.post_attention_groups		,      intermediate_groups=self.intermediate_groups		,      output_groups=self.output_groups		,      )


     def 				SCREAMING_SNAKE_CASE    (		self		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase):

          '''simple docstring'''






          __A		:					Tuple						=			SqueezeBertModel(config=_UpperCAmelCase)
          model.to(_UpperCAmelCase)
          model.eval()
          __A		:					List[Any]						=			model(_UpperCAmelCase		,      _UpperCAmelCase)
          __A		:					str						=			model(_UpperCAmelCase)
          self.parent.assertEqual(result.last_hidden_state.shape		,      (self.batch_size, self.seq_length, self.hidden_size))


     def 				SCREAMING_SNAKE_CASE    (		self		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase):

          '''simple docstring'''






          __A		:					Optional[int]						=			SqueezeBertForMaskedLM(config=_UpperCAmelCase)
          model.to(_UpperCAmelCase)
          model.eval()
          __A		:					Union[str, Any]						=			model(_UpperCAmelCase		,      attention_mask=_UpperCAmelCase		,      labels=_UpperCAmelCase)
          self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.seq_length, self.vocab_size))


     def 				SCREAMING_SNAKE_CASE    (		self		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase):

          '''simple docstring'''






          __A		:					str						=			SqueezeBertForQuestionAnswering(config=_UpperCAmelCase)
          model.to(_UpperCAmelCase)
          model.eval()
          __A		:					Any						=			model(
              _UpperCAmelCase		,      attention_mask=_UpperCAmelCase		,      start_positions=_UpperCAmelCase		,      end_positions=_UpperCAmelCase)
          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		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase):

          '''simple docstring'''






          __A		:					int						=			self.num_labels
          __A		:					Optional[int]						=			SqueezeBertForSequenceClassification(_UpperCAmelCase)
          model.to(_UpperCAmelCase)
          model.eval()
          __A		:					int						=			model(_UpperCAmelCase		,      attention_mask=_UpperCAmelCase		,      labels=_UpperCAmelCase)
          self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.num_labels))


     def 				SCREAMING_SNAKE_CASE    (		self		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase):

          '''simple docstring'''






          __A		:					Dict						=			self.num_labels
          __A		:					str						=			SqueezeBertForTokenClassification(config=_UpperCAmelCase)
          model.to(_UpperCAmelCase)
          model.eval()

          __A		:					List[Any]						=			model(_UpperCAmelCase		,      attention_mask=_UpperCAmelCase		,      labels=_UpperCAmelCase)
          self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.seq_length, self.num_labels))


     def 				SCREAMING_SNAKE_CASE    (		self		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase):

          '''simple docstring'''






          __A		:					Dict						=			self.num_choices
          __A		:					str						=			SqueezeBertForMultipleChoice(config=_UpperCAmelCase)
          model.to(_UpperCAmelCase)
          model.eval()
          __A		:					List[Any]						=			input_ids.unsqueeze(1).expand(-1		,      self.num_choices		,      -1).contiguous()
          __A		:					Any						=			input_mask.unsqueeze(1).expand(-1		,      self.num_choices		,      -1).contiguous()
          __A		:					Optional[Any]						=			model(
              _UpperCAmelCase		,      attention_mask=_UpperCAmelCase		,      labels=_UpperCAmelCase		,      )
          self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.num_choices))


     def 				SCREAMING_SNAKE_CASE    (		self):

          '''simple docstring'''






          __A		:					Tuple						=			self.prepare_config_and_inputs()
          ((__A)					,(__A)					,(__A)					,(__A)					,(__A)					,(__A))		:					Optional[Any]						=			config_and_inputs
          __A		:					Optional[Any]						=			{'input_ids': input_ids, 'attention_mask': input_mask}
          return config, inputs_dict




@require_torch
class 					SCREAMING_SNAKE_CASE       (_lowerCAmelCase   ,						_lowerCAmelCase   ,						unittest.TestCase				):
     lowerCAmelCase						     =			(
         (
             SqueezeBertModel,
             SqueezeBertForMaskedLM,
             SqueezeBertForMultipleChoice,
             SqueezeBertForQuestionAnswering,
             SqueezeBertForSequenceClassification,
             SqueezeBertForTokenClassification,
         )
         if is_torch_available()
         else None
     )
     lowerCAmelCase						     =			(
         {
             '''feature-extraction''': SqueezeBertModel,
             '''fill-mask''': SqueezeBertForMaskedLM,
             '''question-answering''': SqueezeBertForQuestionAnswering,
             '''text-classification''': SqueezeBertForSequenceClassification,
             '''token-classification''': SqueezeBertForTokenClassification,
             '''zero-shot''': SqueezeBertForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     lowerCAmelCase						     =			False
     lowerCAmelCase						     =			True
     lowerCAmelCase						     =			False


     def 				SCREAMING_SNAKE_CASE    (		self):

          '''simple docstring'''






          __A		:					Dict						=			SqueezeBertModelTester(self)
          __A		:					Union[str, Any]						=			ConfigTester(self		,      config_class=_UpperCAmelCase		,      dim=37)


     def 				SCREAMING_SNAKE_CASE    (		self):

          '''simple docstring'''






          self.config_tester.run_common_tests()


     def 				SCREAMING_SNAKE_CASE    (		self):

          '''simple docstring'''






          __A		:					List[Any]						=			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_squeezebert_model(*_UpperCAmelCase)


     def 				SCREAMING_SNAKE_CASE    (		self):

          '''simple docstring'''






          __A		:					List[str]						=			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_squeezebert_for_masked_lm(*_UpperCAmelCase)


     def 				SCREAMING_SNAKE_CASE    (		self):

          '''simple docstring'''






          __A		:					Tuple						=			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_squeezebert_for_question_answering(*_UpperCAmelCase)


     def 				SCREAMING_SNAKE_CASE    (		self):

          '''simple docstring'''






          __A		:					List[str]						=			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_squeezebert_for_sequence_classification(*_UpperCAmelCase)


     def 				SCREAMING_SNAKE_CASE    (		self):

          '''simple docstring'''






          __A		:					Optional[int]						=			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_squeezebert_for_token_classification(*_UpperCAmelCase)


     def 				SCREAMING_SNAKE_CASE    (		self):

          '''simple docstring'''






          __A		:					int						=			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_squeezebert_for_multiple_choice(*_UpperCAmelCase)


     @slow
     def 				SCREAMING_SNAKE_CASE    (		self):

          '''simple docstring'''






          for model_name in SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               __A		:					Dict						=			SqueezeBertModel.from_pretrained(_UpperCAmelCase)
               self.assertIsNotNone(_UpperCAmelCase)




@require_sentencepiece
@require_tokenizers
@require_torch
class 					SCREAMING_SNAKE_CASE       (unittest.TestCase				):


     @slow
     def 				SCREAMING_SNAKE_CASE    (		self):

          '''simple docstring'''






          __A		:					Union[str, Any]						=			SqueezeBertForSequenceClassification.from_pretrained('squeezebert/squeezebert-mnli')

          __A		:					Union[str, Any]						=			torch.tensor([[1, 2_9414, 232, 328, 740, 1140, 1_2695, 69, 13, 1588, 2]])
          __A		:					Any						=			model(_UpperCAmelCase)[0]
          __A		:					Tuple						=			torch.Size((1, 3))
          self.assertEqual(output.shape		,      _UpperCAmelCase)
          __A		:					Optional[Any]						=			torch.tensor([[0.6401, -0.0349, -0.6041]])
          self.assertTrue(torch.allclose(_UpperCAmelCase		,      _UpperCAmelCase		,      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  __UpperCAmelCase						(					SCREAMING_SNAKE_CASE__:						Optional[Any],       SCREAMING_SNAKE_CASE__:						Tuple,       SCREAMING_SNAKE_CASE__:						Dict     )	->      List[Any]:

   """simple docstring"""
   # Initialise PyTorch model
   __a					  =     RemBertConfig.from_json_file(SCREAMING_SNAKE_CASE__     )
   print('Building PyTorch model from configuration: {}'.format(str(SCREAMING_SNAKE_CASE__     )     )     )
   __a					  =     RemBertModel(SCREAMING_SNAKE_CASE__     )

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

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


if __name__ == "__main__":
 __UpperCamelCase   :							List[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(
     """--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[str]         =   parser.parse_args()
 convert_rembert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.rembert_config_file, args.pytorch_dump_path)