
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 	A_       :




	'''simple docstring'''


	def __init__(	self		,   snake_case		,   snake_case=13		,   snake_case=7		,   snake_case=True		,   snake_case=True		,   snake_case=True		,   snake_case=True		,   snake_case=99		,   snake_case=32		,   snake_case=18		,   snake_case=43		,   snake_case=5		,   snake_case=4		,   snake_case=37		,   snake_case="gelu"		,   snake_case=0.1		,   snake_case=0.1		,   snake_case=512		,   snake_case=16		,   snake_case=2		,   snake_case=0.02		,   snake_case=3		,   snake_case=4		,   snake_case=None		,   ):
		lowercase					  =   parent
		lowercase					  =   batch_size
		lowercase					  =   seq_length
		lowercase					  =   is_training
		lowercase					  =   use_input_mask
		lowercase					  =   use_token_type_ids
		lowercase					  =   use_labels
		lowercase					  =   vocab_size
		lowercase					  =   hidden_size
		lowercase					  =   input_embedding_size
		lowercase					  =   output_embedding_size
		lowercase					  =   num_hidden_layers
		lowercase					  =   num_attention_heads
		lowercase					  =   intermediate_size
		lowercase					  =   hidden_act
		lowercase					  =   hidden_dropout_prob
		lowercase					  =   attention_probs_dropout_prob
		lowercase					  =   max_position_embeddings
		lowercase					  =   type_vocab_size
		lowercase					  =   type_sequence_label_size
		lowercase					  =   initializer_range
		lowercase					  =   num_labels
		lowercase					  =   num_choices
		lowercase					  =   scope





	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   ids_tensor([self.batch_size, self.seq_length]		,   self.vocab_size	)

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

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

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

		lowercase					  =   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=snake_case__		,   initializer_range=self.initializer_range		,   )

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





	def      SCREAMING_SNAKE_CASE__						(	self	):
		(
		    (
		    lowercase      
		)  , (
		    lowercase      
		)  , (
		    lowercase      
		)  , (
		    lowercase      
		)  , (
		    lowercase      
		)  , (
		    lowercase      
		)  , (
		    lowercase      
		)  ,       
		)					  =   self.prepare_config_and_inputs()

		lowercase					  =   True
		lowercase					  =   floats_tensor([self.batch_size, self.seq_length, self.hidden_size]	)
		lowercase					  =   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      SCREAMING_SNAKE_CASE__						(	self		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case	):
		lowercase					  =   RemBertModel(config=snake_case__	)
		model.to(snake_case__	)
		model.eval()
		lowercase					  =   model(snake_case__		,   attention_mask=snake_case__		,   token_type_ids=snake_case__	)
		lowercase					  =   model(snake_case__		,   token_type_ids=snake_case__	)
		lowercase					  =   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		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   ):
		lowercase					  =   True
		lowercase					  =   RemBertModel(snake_case__	)
		model.to(snake_case__	)
		model.eval()
		lowercase					  =   model(
		    snake_case__		,   attention_mask=snake_case__		,   token_type_ids=snake_case__		,   encoder_hidden_states=snake_case__		,   encoder_attention_mask=snake_case__		,   )
		lowercase					  =   model(
		    snake_case__		,   attention_mask=snake_case__		,   token_type_ids=snake_case__		,   encoder_hidden_states=snake_case__		,   )
		lowercase					  =   model(snake_case__		,   attention_mask=snake_case__		,   token_type_ids=snake_case__	)
		self.parent.assertEqual(result.last_hidden_state.shape		,   (self.batch_size, self.seq_length, self.hidden_size)	)





	def      SCREAMING_SNAKE_CASE__						(	self		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   ):
		lowercase					  =   RemBertForCausalLM(config=snake_case__	)
		model.to(snake_case__	)
		model.eval()
		lowercase					  =   model(snake_case__		,   attention_mask=snake_case__		,   token_type_ids=snake_case__		,   labels=snake_case__	)
		self.parent.assertEqual(result.logits.shape		,   (self.batch_size, self.seq_length, self.vocab_size)	)





	def      SCREAMING_SNAKE_CASE__						(	self		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case	):
		lowercase					  =   RemBertForMaskedLM(config=snake_case__	)
		model.to(snake_case__	)
		model.eval()
		lowercase					  =   model(snake_case__		,   attention_mask=snake_case__		,   token_type_ids=snake_case__		,   labels=snake_case__	)
		self.parent.assertEqual(result.logits.shape		,   (self.batch_size, self.seq_length, self.vocab_size)	)





	def      SCREAMING_SNAKE_CASE__						(	self		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   ):
		lowercase					  =   True
		lowercase					  =   True
		lowercase					  =   RemBertForCausalLM(config=snake_case__	)
		model.to(snake_case__	)
		model.eval()

		# first forward pass
		lowercase					  =   model(
		    snake_case__		,   attention_mask=snake_case__		,   encoder_hidden_states=snake_case__		,   encoder_attention_mask=snake_case__		,   use_cache=snake_case__		,   )
		lowercase					  =   outputs.past_key_values

		# create hypothetical multiple next token and extent to next_input_ids
		lowercase					  =   ids_tensor((self.batch_size, 3)		,   config.vocab_size	)
		lowercase					  =   ids_tensor((self.batch_size, 3)		,   vocab_size=2	)

		# append to next input_ids and
		lowercase					  =   torch.cat([input_ids, next_tokens]		,   dim=-1	)
		lowercase					  =   torch.cat([input_mask, next_mask]		,   dim=-1	)

		lowercase					  =   model(
		    snake_case__		,   attention_mask=snake_case__		,   encoder_hidden_states=snake_case__		,   encoder_attention_mask=snake_case__		,   output_hidden_states=snake_case__		,   )['hidden_states'][0]
		lowercase					  =   model(
		    snake_case__		,   attention_mask=snake_case__		,   encoder_hidden_states=snake_case__		,   encoder_attention_mask=snake_case__		,   past_key_values=snake_case__		,   output_hidden_states=snake_case__		,   )['hidden_states'][0]

		# select random slice
		lowercase					  =   ids_tensor((1,)		,   output_from_past.shape[-1]	).item()
		lowercase					  =   output_from_no_past[:, -3:, random_slice_idx].detach()
		lowercase					  =   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(snake_case__		,   snake_case__		,   atol=1E-3	)	)





	def      SCREAMING_SNAKE_CASE__						(	self		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case	):
		lowercase					  =   RemBertForQuestionAnswering(config=snake_case__	)
		model.to(snake_case__	)
		model.eval()
		lowercase					  =   model(
		    snake_case__		,   attention_mask=snake_case__		,   token_type_ids=snake_case__		,   start_positions=snake_case__		,   end_positions=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		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case	):
		lowercase					  =   self.num_labels
		lowercase					  =   RemBertForSequenceClassification(snake_case__	)
		model.to(snake_case__	)
		model.eval()
		lowercase					  =   model(snake_case__		,   attention_mask=snake_case__		,   token_type_ids=snake_case__		,   labels=snake_case__	)
		self.parent.assertEqual(result.logits.shape		,   (self.batch_size, self.num_labels)	)





	def      SCREAMING_SNAKE_CASE__						(	self		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case	):
		lowercase					  =   self.num_labels
		lowercase					  =   RemBertForTokenClassification(config=snake_case__	)
		model.to(snake_case__	)
		model.eval()
		lowercase					  =   model(snake_case__		,   attention_mask=snake_case__		,   token_type_ids=snake_case__		,   labels=snake_case__	)
		self.parent.assertEqual(result.logits.shape		,   (self.batch_size, self.seq_length, self.num_labels)	)





	def      SCREAMING_SNAKE_CASE__						(	self		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case	):
		lowercase					  =   self.num_choices
		lowercase					  =   RemBertForMultipleChoice(config=snake_case__	)
		model.to(snake_case__	)
		model.eval()
		lowercase					  =   input_ids.unsqueeze(1	).expand(-1		,   self.num_choices		,   -1	).contiguous()
		lowercase					  =   token_type_ids.unsqueeze(1	).expand(-1		,   self.num_choices		,   -1	).contiguous()
		lowercase					  =   input_mask.unsqueeze(1	).expand(-1		,   self.num_choices		,   -1	).contiguous()
		lowercase					  =   model(
		    snake_case__		,   attention_mask=snake_case__		,   token_type_ids=snake_case__		,   labels=snake_case__		,   )
		self.parent.assertEqual(result.logits.shape		,   (self.batch_size, self.num_choices)	)





	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   self.prepare_config_and_inputs()
		(
		    (
		    lowercase      
		)  , (
		    lowercase      
		)  , (
		    lowercase      
		)  , (
		    lowercase      
		)  , (
		    lowercase      
		)  , (
		    lowercase      
		)  , (
		    lowercase      
		)  ,       
		)					  =   config_and_inputs
		lowercase					  =   {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
		return config, inputs_dict







@require_torch
class 	A_       (      a_				,     a_				,     unittest.TestCase       ):




	'''simple docstring'''


	_UpperCamelCase					:					str										=	(
	    (
	        RemBertModel,
	        RemBertForMaskedLM,
	        RemBertForCausalLM,
	        RemBertForMultipleChoice,
	        RemBertForQuestionAnswering,
	        RemBertForSequenceClassification,
	        RemBertForTokenClassification,
	    )
	    if is_torch_available()
	    else ()
	)
	_UpperCamelCase					:					List[str]										=	(RemBertForCausalLM,) if is_torch_available() else ()
	_UpperCamelCase					:					int										=	(
	    {
	        """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      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   RemBertModelTester(self	)
		lowercase					  =   ConfigTester(self		,   config_class=snake_case__		,   hidden_size=37	)





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





	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_model(*snake_case__	)





	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   self.model_tester.prepare_config_and_inputs()
		for type in ["absolute", "relative_key", "relative_key_query"]:
			lowercase					  =   type
			self.model_tester.create_and_check_model(*snake_case__	)





	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_for_masked_lm(*snake_case__	)





	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_for_multiple_choice(*snake_case__	)





	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   self.model_tester.prepare_config_and_inputs_for_decoder()
		self.model_tester.create_and_check_decoder_model_past_large_inputs(*snake_case__	)





	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_for_question_answering(*snake_case__	)





	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_for_sequence_classification(*snake_case__	)





	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_for_token_classification(*snake_case__	)





	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   self.model_tester.prepare_config_and_inputs_for_decoder()
		self.model_tester.create_and_check_model_as_decoder(*snake_case__	)





	def      SCREAMING_SNAKE_CASE__						(	self	):
		(
		    (
		    lowercase      
		)  , (
		    lowercase      
		)  , (
		    lowercase      
		)  , (
		    lowercase      
		)  , (
		    lowercase      
		)  , (
		    lowercase      
		)  , (
		    lowercase      
		)  , (
		    lowercase      
		)  , (
		    lowercase      
		)  ,       
		)					  =   self.model_tester.prepare_config_and_inputs_for_decoder()

		lowercase					  =   None

		self.model_tester.create_and_check_model_as_decoder(
		    snake_case__		,   snake_case__		,   snake_case__		,   snake_case__		,   snake_case__		,   snake_case__		,   snake_case__		,   snake_case__		,   snake_case__		,   )





	@slow
	def      SCREAMING_SNAKE_CASE__						(	self	):


		for model_name in REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
			lowercase					  =   RemBertModel.from_pretrained(snake_case__	)
			self.assertIsNotNone(snake_case__	)







@require_torch
class 	A_       (      unittest.TestCase       ):




	'''simple docstring'''


	@slow
	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   RemBertModel.from_pretrained('google/rembert'	)
		lowercase					  =   torch.tensor([[312, 5_6498, 313, 2125, 313]]	)
		lowercase					  =   torch.tensor([[0, 0, 0, 1, 1]]	)
		with torch.no_grad():
			lowercase					  =   model(snake_case__		,   token_type_ids=snake_case__		,   output_hidden_states=snake_case__	)

		lowercase					  =   1152

		lowercase					  =   torch.Size((1, 5, hidden_size)	)
		self.assertEqual(output['last_hidden_state'].shape		,   snake_case__	)

		lowercase					  =   torch.tensor(
		    [
		        [
		            [0.0_754, -0.2_022, 0.1_904],
		            [-0.3_354, -0.3_692, -0.4_791],
		            [-0.2_314, -0.6_729, -0.0_749],
		            [-0.0_396, -0.3_105, -0.4_234],
		            [-0.1_571, -0.0_525, 0.5_353],
		        ]
		    ]	)

		# 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]		,   snake_case__		,   atol=1E-4	)	)



from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCAmelCase     				=			{
    '''configuration_dpr''': ['''DPR_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''DPRConfig'''],
    '''tokenization_dpr''': [
        '''DPRContextEncoderTokenizer''',
        '''DPRQuestionEncoderTokenizer''',
        '''DPRReaderOutput''',
        '''DPRReaderTokenizer''',
    ],
}


try:
	if not is_tokenizers_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	UpperCAmelCase     				=			[
	    '''DPRContextEncoderTokenizerFast''',
	    '''DPRQuestionEncoderTokenizerFast''',
	    '''DPRReaderTokenizerFast''',
	]

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	UpperCAmelCase     				=			[
	    '''DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST''',
	    '''DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST''',
	    '''DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST''',
	    '''DPRContextEncoder''',
	    '''DPRPretrainedContextEncoder''',
	    '''DPRPreTrainedModel''',
	    '''DPRPretrainedQuestionEncoder''',
	    '''DPRPretrainedReader''',
	    '''DPRQuestionEncoder''',
	    '''DPRReader''',
	]

try:
	if not is_tf_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	UpperCAmelCase     				=			[
	    '''TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST''',
	    '''TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST''',
	    '''TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST''',
	    '''TFDPRContextEncoder''',
	    '''TFDPRPretrainedContextEncoder''',
	    '''TFDPRPretrainedQuestionEncoder''',
	    '''TFDPRPretrainedReader''',
	    '''TFDPRQuestionEncoder''',
	    '''TFDPRReader''',
	]


if TYPE_CHECKING:
	from .configuration_dpr import DPR_PRETRAINED_CONFIG_ARCHIVE_MAP, DPRConfig
	from .tokenization_dpr import (
	    DPRContextEncoderTokenizer,
	    DPRQuestionEncoderTokenizer,
	    DPRReaderOutput,
	    DPRReaderTokenizer,
	)

	try:
		if not is_tokenizers_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .tokenization_dpr_fast import (
		    DPRContextEncoderTokenizerFast,
		    DPRQuestionEncoderTokenizerFast,
		    DPRReaderTokenizerFast,
		)

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_dpr import (
		    DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
		    DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
		    DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
		    DPRContextEncoder,
		    DPRPretrainedContextEncoder,
		    DPRPreTrainedModel,
		    DPRPretrainedQuestionEncoder,
		    DPRPretrainedReader,
		    DPRQuestionEncoder,
		    DPRReader,
		)

	try:
		if not is_tf_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_tf_dpr import (
		    TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFDPRContextEncoder,
		    TFDPRPretrainedContextEncoder,
		    TFDPRPretrainedQuestionEncoder,
		    TFDPRPretrainedReader,
		    TFDPRQuestionEncoder,
		    TFDPRReader,
		)

else:
	import sys

	UpperCAmelCase     				=			_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)

