


import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

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


if is_torch_available():
    import torch

    from transformers import (
        RemBertConfig,
        RemBertForCausalLM,
        RemBertForMaskedLM,
        RemBertForMultipleChoice,
        RemBertForQuestionAnswering,
        RemBertForSequenceClassification,
        RemBertForTokenClassification,
        RemBertModel,
    )
    from transformers.models.rembert.modeling_rembert import REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST







class 					lowerCAmelCase__	:




     '''simple docstring'''


     def __init__(						self			,						__lowerCamelCase			,						__lowerCamelCase=1_3			,						__lowerCamelCase=7			,						__lowerCamelCase=True			,						__lowerCamelCase=True			,						__lowerCamelCase=True			,						__lowerCamelCase=True			,						__lowerCamelCase=9_9			,						__lowerCamelCase=3_2			,						__lowerCamelCase=1_8			,						__lowerCamelCase=4_3			,						__lowerCamelCase=5			,						__lowerCamelCase=4			,						__lowerCamelCase=3_7			,						__lowerCamelCase="gelu"			,						__lowerCamelCase=0.1			,						__lowerCamelCase=0.1			,						__lowerCamelCase=5_1_2			,						__lowerCamelCase=1_6			,						__lowerCamelCase=2			,						__lowerCamelCase=0.0_2			,						__lowerCamelCase=3			,						__lowerCamelCase=4			,						__lowerCamelCase=None			,						)      ->    Union[str, Any]:
            _A				:	int       =			parent
            _A				:	Tuple       =			batch_size
            _A				:	Any       =			seq_length
            _A				:	List[Any]       =			is_training
            _A				:	Dict       =			use_input_mask
            _A				:	List[Any]       =			use_token_type_ids
            _A				:	Union[str, Any]       =			use_labels
            _A				:	Any       =			vocab_size
            _A				:	Optional[int]       =			hidden_size
            _A				:	Optional[Any]       =			input_embedding_size
            _A				:	Dict       =			output_embedding_size
            _A				:	List[Any]       =			num_hidden_layers
            _A				:	List[Any]       =			num_attention_heads
            _A				:	Optional[Any]       =			intermediate_size
            _A				:	int       =			hidden_act
            _A				:	int       =			hidden_dropout_prob
            _A				:	Dict       =			attention_probs_dropout_prob
            _A				:	str       =			max_position_embeddings
            _A				:	List[str]       =			type_vocab_size
            _A				:	List[str]       =			type_sequence_label_size
            _A				:	Any       =			initializer_range
            _A				:	Optional[int]       =			num_labels
            _A				:	Optional[int]       =			num_choices
            _A				:	List[str]       =			scope


     def 			_lowerCamelCase			(						self)      ->    Optional[int]:
            _A				:	int       =			ids_tensor([self.batch_size, self.seq_length]			,						self.vocab_size)

            _A				:	Optional[int]       =			None
            if self.use_input_mask:
                   _A				:	Union[str, Any]       =			random_attention_mask([self.batch_size, self.seq_length])

            _A				:	Optional[int]       =			None
            if self.use_token_type_ids:
                   _A				:	str       =			ids_tensor([self.batch_size, self.seq_length]			,						self.type_vocab_size)

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

            _A				:	Tuple       =			RemBertConfig(
                vocab_size=self.vocab_size			,						hidden_size=self.hidden_size			,						input_embedding_size=self.input_embedding_size			,						output_embedding_size=self.output_embedding_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=__lowerCamelCase			,						initializer_range=self.initializer_range			,						)

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


     def 			_lowerCamelCase			(						self)      ->    Any:
            (
                (
                _A
            )   ,      (
                _A
            )   ,      (
                _A
            )   ,      (
                _A
            )   ,      (
                _A
            )   ,      (
                _A
            )   ,      (
                _A
            )   ,      
            )				:	Any       =			self.prepare_config_and_inputs()

            _A				:	Dict       =			True
            _A				:	int       =			floats_tensor([self.batch_size, self.seq_length, self.hidden_size])
            _A				:	Optional[int]       =			ids_tensor([self.batch_size, self.seq_length]			,						vocab_size=2)

            return (
                config,
                input_ids,
                token_type_ids,
                input_mask,
                sequence_labels,
                token_labels,
                choice_labels,
                encoder_hidden_states,
                encoder_attention_mask,
            )


     def 			_lowerCamelCase			(						self			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase)      ->    List[Any]:
            _A				:	List[str]       =			RemBertModel(config=__lowerCamelCase)
            model.to(__lowerCamelCase)
            model.eval()
            _A				:	List[str]       =			model(__lowerCamelCase			,						attention_mask=__lowerCamelCase			,						token_type_ids=__lowerCamelCase)
            _A				:	List[Any]       =			model(__lowerCamelCase			,						token_type_ids=__lowerCamelCase)
            _A				:	str       =			model(__lowerCamelCase)
            self.parent.assertEqual(result.last_hidden_state.shape			,						(self.batch_size, self.seq_length, self.hidden_size))


     def 			_lowerCamelCase			(						self			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						)      ->    str:
            _A				:	Tuple       =			True
            _A				:	List[Any]       =			RemBertModel(__lowerCamelCase)
            model.to(__lowerCamelCase)
            model.eval()
            _A				:	int       =			model(
                __lowerCamelCase			,						attention_mask=__lowerCamelCase			,						token_type_ids=__lowerCamelCase			,						encoder_hidden_states=__lowerCamelCase			,						encoder_attention_mask=__lowerCamelCase			,						)
            _A				:	Optional[Any]       =			model(
                __lowerCamelCase			,						attention_mask=__lowerCamelCase			,						token_type_ids=__lowerCamelCase			,						encoder_hidden_states=__lowerCamelCase			,						)
            _A				:	Dict       =			model(__lowerCamelCase			,						attention_mask=__lowerCamelCase			,						token_type_ids=__lowerCamelCase)
            self.parent.assertEqual(result.last_hidden_state.shape			,						(self.batch_size, self.seq_length, self.hidden_size))


     def 			_lowerCamelCase			(						self			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						)      ->    int:
            _A				:	Tuple       =			RemBertForCausalLM(config=__lowerCamelCase)
            model.to(__lowerCamelCase)
            model.eval()
            _A				:	str       =			model(__lowerCamelCase			,						attention_mask=__lowerCamelCase			,						token_type_ids=__lowerCamelCase			,						labels=__lowerCamelCase)
            self.parent.assertEqual(result.logits.shape			,						(self.batch_size, self.seq_length, self.vocab_size))


     def 			_lowerCamelCase			(						self			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase)      ->    Dict:
            _A				:	Optional[Any]       =			RemBertForMaskedLM(config=__lowerCamelCase)
            model.to(__lowerCamelCase)
            model.eval()
            _A				:	Dict       =			model(__lowerCamelCase			,						attention_mask=__lowerCamelCase			,						token_type_ids=__lowerCamelCase			,						labels=__lowerCamelCase)
            self.parent.assertEqual(result.logits.shape			,						(self.batch_size, self.seq_length, self.vocab_size))


     def 			_lowerCamelCase			(						self			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						)      ->    Dict:
            _A				:	Optional[int]       =			True
            _A				:	str       =			True
            _A				:	Any       =			RemBertForCausalLM(config=__lowerCamelCase)
            model.to(__lowerCamelCase)
            model.eval()

            # first forward pass
            _A				:	Dict       =			model(
                __lowerCamelCase			,						attention_mask=__lowerCamelCase			,						encoder_hidden_states=__lowerCamelCase			,						encoder_attention_mask=__lowerCamelCase			,						use_cache=__lowerCamelCase			,						)
            _A				:	List[Any]       =			outputs.past_key_values

            # create hypothetical multiple next token and extent to next_input_ids
            _A				:	str       =			ids_tensor((self.batch_size, 3)			,						config.vocab_size)
            _A				:	Union[str, Any]       =			ids_tensor((self.batch_size, 3)			,						vocab_size=2)

            # append to next input_ids and
            _A				:	Union[str, Any]       =			torch.cat([input_ids, next_tokens]			,						dim=-1)
            _A				:	int       =			torch.cat([input_mask, next_mask]			,						dim=-1)

            _A				:	str       =			model(
                __lowerCamelCase			,						attention_mask=__lowerCamelCase			,						encoder_hidden_states=__lowerCamelCase			,						encoder_attention_mask=__lowerCamelCase			,						output_hidden_states=__lowerCamelCase			,						)["hidden_states"][0]
            _A				:	List[Any]       =			model(
                __lowerCamelCase			,						attention_mask=__lowerCamelCase			,						encoder_hidden_states=__lowerCamelCase			,						encoder_attention_mask=__lowerCamelCase			,						past_key_values=__lowerCamelCase			,						output_hidden_states=__lowerCamelCase			,						)["hidden_states"][0]

            # select random slice
            _A				:	Dict       =			ids_tensor((1,)			,						output_from_past.shape[-1]).item()
            _A				:	Dict       =			output_from_no_past[:, -3:, random_slice_idx].detach()
            _A				:	int       =			output_from_past[:, :, random_slice_idx].detach()

            self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1])

            # test that outputs are equal for slice
            self.parent.assertTrue(torch.allclose(__lowerCamelCase			,						__lowerCamelCase			,						atol=1e-3))


     def 			_lowerCamelCase			(						self			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase)      ->    Dict:
            _A				:	int       =			RemBertForQuestionAnswering(config=__lowerCamelCase)
            model.to(__lowerCamelCase)
            model.eval()
            _A				:	Optional[int]       =			model(
                __lowerCamelCase			,						attention_mask=__lowerCamelCase			,						token_type_ids=__lowerCamelCase			,						start_positions=__lowerCamelCase			,						end_positions=__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 			_lowerCamelCase			(						self			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase)      ->    Tuple:
            _A				:	Tuple       =			self.num_labels
            _A				:	str       =			RemBertForSequenceClassification(__lowerCamelCase)
            model.to(__lowerCamelCase)
            model.eval()
            _A				:	List[str]       =			model(__lowerCamelCase			,						attention_mask=__lowerCamelCase			,						token_type_ids=__lowerCamelCase			,						labels=__lowerCamelCase)
            self.parent.assertEqual(result.logits.shape			,						(self.batch_size, self.num_labels))


     def 			_lowerCamelCase			(						self			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase)      ->    Optional[Any]:
            _A				:	List[Any]       =			self.num_labels
            _A				:	List[str]       =			RemBertForTokenClassification(config=__lowerCamelCase)
            model.to(__lowerCamelCase)
            model.eval()
            _A				:	Dict       =			model(__lowerCamelCase			,						attention_mask=__lowerCamelCase			,						token_type_ids=__lowerCamelCase			,						labels=__lowerCamelCase)
            self.parent.assertEqual(result.logits.shape			,						(self.batch_size, self.seq_length, self.num_labels))


     def 			_lowerCamelCase			(						self			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase)      ->    int:
            _A				:	List[str]       =			self.num_choices
            _A				:	Optional[Any]       =			RemBertForMultipleChoice(config=__lowerCamelCase)
            model.to(__lowerCamelCase)
            model.eval()
            _A				:	int       =			input_ids.unsqueeze(1).expand(-1			,						self.num_choices			,						-1).contiguous()
            _A				:	Tuple       =			token_type_ids.unsqueeze(1).expand(-1			,						self.num_choices			,						-1).contiguous()
            _A				:	Union[str, Any]       =			input_mask.unsqueeze(1).expand(-1			,						self.num_choices			,						-1).contiguous()
            _A				:	Dict       =			model(
                __lowerCamelCase			,						attention_mask=__lowerCamelCase			,						token_type_ids=__lowerCamelCase			,						labels=__lowerCamelCase			,						)
            self.parent.assertEqual(result.logits.shape			,						(self.batch_size, self.num_choices))


     def 			_lowerCamelCase			(						self)      ->    Any:
            _A				:	Optional[Any]       =			self.prepare_config_and_inputs()
            (
                (
                _A
            )   ,      (
                _A
            )   ,      (
                _A
            )   ,      (
                _A
            )   ,      (
                _A
            )   ,      (
                _A
            )   ,      (
                _A
            )   ,      
            )				:	List[str]       =			config_and_inputs
            _A				:	Any       =			{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
            return config, inputs_dict







@require_torch
class 					lowerCAmelCase__	(    a					,   a					,   unittest.TestCase):




     '''simple docstring'''
     __SCREAMING_SNAKE_CASE        =			(
         (
             RemBertModel,
             RemBertForMaskedLM,
             RemBertForCausalLM,
             RemBertForMultipleChoice,
             RemBertForQuestionAnswering,
             RemBertForSequenceClassification,
             RemBertForTokenClassification,
         )
         if is_torch_available()
         else ()
     )
     __SCREAMING_SNAKE_CASE        =			(RemBertForCausalLM,) if is_torch_available() else ()
     __SCREAMING_SNAKE_CASE        =			(
         {
             "feature-extraction": RemBertModel,
             "fill-mask": RemBertForMaskedLM,
             "question-answering": RemBertForQuestionAnswering,
             "text-classification": RemBertForSequenceClassification,
             "text-generation": RemBertForCausalLM,
             "token-classification": RemBertForTokenClassification,
             "zero-shot": RemBertForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )


     def 			_lowerCamelCase			(						self)      ->    int:
            _A				:	Union[str, Any]       =			RemBertModelTester(self)
            _A				:	List[Any]       =			ConfigTester(self			,						config_class=__lowerCamelCase			,						hidden_size=3_7)


     def 			_lowerCamelCase			(						self)      ->    Union[str, Any]:
            self.config_tester.run_common_tests()


     def 			_lowerCamelCase			(						self)      ->    Dict:
            _A				:	Optional[int]       =			self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_model(*__lowerCamelCase)


     def 			_lowerCamelCase			(						self)      ->    Optional[int]:
            _A				:	List[str]       =			self.model_tester.prepare_config_and_inputs()
            for type in ["absolute", "relative_key", "relative_key_query"]:
                   _A				:	Optional[int]       =			type
                   self.model_tester.create_and_check_model(*__lowerCamelCase)


     def 			_lowerCamelCase			(						self)      ->    str:
            _A				:	Any       =			self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_masked_lm(*__lowerCamelCase)


     def 			_lowerCamelCase			(						self)      ->    Optional[Any]:
            _A				:	Any       =			self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_multiple_choice(*__lowerCamelCase)


     def 			_lowerCamelCase			(						self)      ->    List[str]:
            _A				:	List[Any]       =			self.model_tester.prepare_config_and_inputs_for_decoder()
            self.model_tester.create_and_check_decoder_model_past_large_inputs(*__lowerCamelCase)


     def 			_lowerCamelCase			(						self)      ->    Optional[Any]:
            _A				:	List[str]       =			self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_question_answering(*__lowerCamelCase)


     def 			_lowerCamelCase			(						self)      ->    List[Any]:
            _A				:	List[Any]       =			self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_sequence_classification(*__lowerCamelCase)


     def 			_lowerCamelCase			(						self)      ->    str:
            _A				:	List[Any]       =			self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_token_classification(*__lowerCamelCase)


     def 			_lowerCamelCase			(						self)      ->    Dict:
            _A				:	Union[str, Any]       =			self.model_tester.prepare_config_and_inputs_for_decoder()
            self.model_tester.create_and_check_model_as_decoder(*__lowerCamelCase)


     def 			_lowerCamelCase			(						self)      ->    int:
            # This regression test was failing with PyTorch < 1.3
            (
                (
                _A
            )   ,      (
                _A
            )   ,      (
                _A
            )   ,      (
                _A
            )   ,      (
                _A
            )   ,      (
                _A
            )   ,      (
                _A
            )   ,      (
                _A
            )   ,      (
                _A
            )   ,      
            )				:	Optional[Any]       =			self.model_tester.prepare_config_and_inputs_for_decoder()

            _A				:	Union[str, Any]       =			None

            self.model_tester.create_and_check_model_as_decoder(
                __lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						__lowerCamelCase			,						)


     @slow
     def 			_lowerCamelCase			(						self)      ->    str:
            for model_name in REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                   _A				:	str       =			RemBertModel.from_pretrained(__lowerCamelCase)
                   self.assertIsNotNone(__lowerCamelCase)







@require_torch
class 					lowerCAmelCase__	(    unittest.TestCase):




     '''simple docstring'''


     @slow
     def 			_lowerCamelCase			(						self)      ->    Optional[Any]:
            # Test exact values at the last hidden layer
            _A				:	Union[str, Any]       =			RemBertModel.from_pretrained("google/rembert")
            _A				:	int       =			torch.tensor([[3_1_2, 5_6_4_9_8, 3_1_3, 2_1_2_5, 3_1_3]])
            _A				:	int       =			torch.tensor([[0, 0, 0, 1, 1]])
            with torch.no_grad():
                   _A				:	List[str]       =			model(__lowerCamelCase			,						token_type_ids=__lowerCamelCase			,						output_hidden_states=__lowerCamelCase)

            _A				:	str       =			1_1_5_2

            _A				:	Tuple       =			torch.Size((1, 5, hidden_size))
            self.assertEqual(output["last_hidden_state"].shape			,						__lowerCamelCase)

            _A				:	Dict       =			torch.tensor(
                [
                    [
                        [0.0_7_5_4, -0.2_0_2_2, 0.1_9_0_4],
                        [-0.3_3_5_4, -0.3_6_9_2, -0.4_7_9_1],
                        [-0.2_3_1_4, -0.6_7_2_9, -0.0_7_4_9],
                        [-0.0_3_9_6, -0.3_1_0_5, -0.4_2_3_4],
                        [-0.1_5_7_1, -0.0_5_2_5, 0.5_3_5_3],
                    ]
                ])

            # Running on the original tf implementation gives slightly different results here.
            # Not clear why this variations is present
            # TODO: Find reason for discrepancy
            # expected_original_implementation = [[
            #     [0.07630594074726105, -0.20146065950393677, 0.19107051193714142],
            #     [-0.3405614495277405, -0.36971670389175415, -0.4808273911476135],
            #     [-0.22587086260318756, -0.6656315922737122, -0.07844287157058716],
            #     [-0.04145475849509239, -0.3077218234539032, -0.42316967248916626],
            #     [-0.15887849032878876, -0.054529931396245956, 0.5356100797653198]
            # ]]

            self.assertTrue(torch.allclose(output["last_hidden_state"][:, :, :3]			,						__lowerCamelCase			,						atol=1e-4))




import warnings

from .generation import FlaxGenerationMixin







class 					lowerCAmelCase__	(    a):




     '''simple docstring'''
     # warning at import time
     warnings.warn(
         "Importing `FlaxGenerationMixin` from `src/transformers/generation_flax_utils.py` is deprecated and will "
         "be removed in Transformers v5. Import as `from transformers import FlaxGenerationMixin` instead."					,   a					,   )
