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






			def __init__(			self							, __UpperCamelCase							, __UpperCamelCase=13							, __UpperCamelCase=7							, __UpperCamelCase=True							, __UpperCamelCase=True							, __UpperCamelCase=True							, __UpperCamelCase=True							, __UpperCamelCase=99							, __UpperCamelCase=32							, __UpperCamelCase=18							, __UpperCamelCase=43							, __UpperCamelCase=5							, __UpperCamelCase=4							, __UpperCamelCase=37							, __UpperCamelCase="gelu"							, __UpperCamelCase=0.1							, __UpperCamelCase=0.1							, __UpperCamelCase=512							, __UpperCamelCase=16							, __UpperCamelCase=2							, __UpperCamelCase=0.02							, __UpperCamelCase=3							, __UpperCamelCase=4							, __UpperCamelCase=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       UpperCamelCase_							(			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=__UpperCamelCase							, initializer_range=self.initializer_range							, )

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






			def       UpperCamelCase_							(			self     ):
									(
									    _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       UpperCamelCase_							(			self							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase     ):
									_lowercase									=				RemBertModel(config=__UpperCamelCase     )
									model.to(__UpperCamelCase     )
									model.eval()
									_lowercase									=				model(__UpperCamelCase							, attention_mask=__UpperCamelCase							, token_type_ids=__UpperCamelCase     )
									_lowercase									=				model(__UpperCamelCase							, token_type_ids=__UpperCamelCase     )
									_lowercase									=				model(__UpperCamelCase     )
									self.parent.assertEqual(result.last_hidden_state.shape							, (self.batch_size, self.seq_length, self.hidden_size)     )






			def       UpperCamelCase_							(			self							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, ):
									_lowercase									=				True
									_lowercase									=				RemBertModel(__UpperCamelCase     )
									model.to(__UpperCamelCase     )
									model.eval()
									_lowercase									=				model(
									    __UpperCamelCase							, attention_mask=__UpperCamelCase							, token_type_ids=__UpperCamelCase							, encoder_hidden_states=__UpperCamelCase							, encoder_attention_mask=__UpperCamelCase							, )
									_lowercase									=				model(
									    __UpperCamelCase							, attention_mask=__UpperCamelCase							, token_type_ids=__UpperCamelCase							, encoder_hidden_states=__UpperCamelCase							, )
									_lowercase									=				model(__UpperCamelCase							, attention_mask=__UpperCamelCase							, token_type_ids=__UpperCamelCase     )
									self.parent.assertEqual(result.last_hidden_state.shape							, (self.batch_size, self.seq_length, self.hidden_size)     )






			def       UpperCamelCase_							(			self							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, ):
									_lowercase									=				RemBertForCausalLM(config=__UpperCamelCase     )
									model.to(__UpperCamelCase     )
									model.eval()
									_lowercase									=				model(__UpperCamelCase							, attention_mask=__UpperCamelCase							, token_type_ids=__UpperCamelCase							, labels=__UpperCamelCase     )
									self.parent.assertEqual(result.logits.shape							, (self.batch_size, self.seq_length, self.vocab_size)     )






			def       UpperCamelCase_							(			self							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase     ):
									_lowercase									=				RemBertForMaskedLM(config=__UpperCamelCase     )
									model.to(__UpperCamelCase     )
									model.eval()
									_lowercase									=				model(__UpperCamelCase							, attention_mask=__UpperCamelCase							, token_type_ids=__UpperCamelCase							, labels=__UpperCamelCase     )
									self.parent.assertEqual(result.logits.shape							, (self.batch_size, self.seq_length, self.vocab_size)     )






			def       UpperCamelCase_							(			self							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, ):
									_lowercase									=				True
									_lowercase									=				True
									_lowercase									=				RemBertForCausalLM(config=__UpperCamelCase     )
									model.to(__UpperCamelCase     )
									model.eval()

									# first forward pass
									_lowercase									=				model(
									    __UpperCamelCase							, attention_mask=__UpperCamelCase							, encoder_hidden_states=__UpperCamelCase							, encoder_attention_mask=__UpperCamelCase							, use_cache=__UpperCamelCase							, )
									_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(
									    __UpperCamelCase							, attention_mask=__UpperCamelCase							, encoder_hidden_states=__UpperCamelCase							, encoder_attention_mask=__UpperCamelCase							, output_hidden_states=__UpperCamelCase							, )['hidden_states'][0]
									_lowercase									=				model(
									    __UpperCamelCase							, attention_mask=__UpperCamelCase							, encoder_hidden_states=__UpperCamelCase							, encoder_attention_mask=__UpperCamelCase							, past_key_values=__UpperCamelCase							, output_hidden_states=__UpperCamelCase							, )['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(__UpperCamelCase							, __UpperCamelCase							, atol=1E-3     )     )






			def       UpperCamelCase_							(			self							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase     ):
									_lowercase									=				RemBertForQuestionAnswering(config=__UpperCamelCase     )
									model.to(__UpperCamelCase     )
									model.eval()
									_lowercase									=				model(
									    __UpperCamelCase							, attention_mask=__UpperCamelCase							, token_type_ids=__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       UpperCamelCase_							(			self							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase     ):
									_lowercase									=				self.num_labels
									_lowercase									=				RemBertForSequenceClassification(__UpperCamelCase     )
									model.to(__UpperCamelCase     )
									model.eval()
									_lowercase									=				model(__UpperCamelCase							, attention_mask=__UpperCamelCase							, token_type_ids=__UpperCamelCase							, labels=__UpperCamelCase     )
									self.parent.assertEqual(result.logits.shape							, (self.batch_size, self.num_labels)     )






			def       UpperCamelCase_							(			self							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase     ):
									_lowercase									=				self.num_labels
									_lowercase									=				RemBertForTokenClassification(config=__UpperCamelCase     )
									model.to(__UpperCamelCase     )
									model.eval()
									_lowercase									=				model(__UpperCamelCase							, attention_mask=__UpperCamelCase							, token_type_ids=__UpperCamelCase							, labels=__UpperCamelCase     )
									self.parent.assertEqual(result.logits.shape							, (self.batch_size, self.seq_length, self.num_labels)     )






			def       UpperCamelCase_							(			self							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase     ):
									_lowercase									=				self.num_choices
									_lowercase									=				RemBertForMultipleChoice(config=__UpperCamelCase     )
									model.to(__UpperCamelCase     )
									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(
									    __UpperCamelCase							, attention_mask=__UpperCamelCase							, token_type_ids=__UpperCamelCase							, labels=__UpperCamelCase							, )
									self.parent.assertEqual(result.logits.shape							, (self.batch_size, self.num_choices)     )






			def       UpperCamelCase_							(			self     ):
									_lowercase									=				self.prepare_config_and_inputs()
									(
									    _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_   ( snake_case__				, snake_case__				, unittest.TestCase						):
			a  : int				      =       (
			    (
			        RemBertModel,
			        RemBertForMaskedLM,
			        RemBertForCausalLM,
			        RemBertForMultipleChoice,
			        RemBertForQuestionAnswering,
			        RemBertForSequenceClassification,
			        RemBertForTokenClassification,
			    )
			    if is_torch_available()
			    else ()
			)
			a  : Tuple				      =       (RemBertForCausalLM,) if is_torch_available() else ()
			a  : List[Any]				      =       (
			    {
			        '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       UpperCamelCase_							(			self     ):
									_lowercase									=				RemBertModelTester(self     )
									_lowercase									=				ConfigTester(self							, config_class=__UpperCamelCase							, hidden_size=37     )






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






			def       UpperCamelCase_							(			self     ):
									_lowercase									=				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*__UpperCamelCase     )






			def       UpperCamelCase_							(			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(*__UpperCamelCase     )






			def       UpperCamelCase_							(			self     ):
									_lowercase									=				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_masked_lm(*__UpperCamelCase     )






			def       UpperCamelCase_							(			self     ):
									_lowercase									=				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_multiple_choice(*__UpperCamelCase     )






			def       UpperCamelCase_							(			self     ):
									_lowercase									=				self.model_tester.prepare_config_and_inputs_for_decoder()
									self.model_tester.create_and_check_decoder_model_past_large_inputs(*__UpperCamelCase     )






			def       UpperCamelCase_							(			self     ):
									_lowercase									=				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_question_answering(*__UpperCamelCase     )






			def       UpperCamelCase_							(			self     ):
									_lowercase									=				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_sequence_classification(*__UpperCamelCase     )






			def       UpperCamelCase_							(			self     ):
									_lowercase									=				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_token_classification(*__UpperCamelCase     )






			def       UpperCamelCase_							(			self     ):
									_lowercase									=				self.model_tester.prepare_config_and_inputs_for_decoder()
									self.model_tester.create_and_check_model_as_decoder(*__UpperCamelCase     )






			def       UpperCamelCase_							(			self     ):
									# This regression test was failing with PyTorch < 1.3
									(
									    _lowercase			
									)									=				self.model_tester.prepare_config_and_inputs_for_decoder()

									_lowercase									=				None

									self.model_tester.create_and_check_model_as_decoder(
									    __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, __UpperCamelCase							, )






			@slow
			def       UpperCamelCase_							(			self     ):
									for model_name in REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															_lowercase									=				RemBertModel.from_pretrained(__UpperCamelCase     )
															self.assertIsNotNone(__UpperCamelCase     )







@require_torch
class       a_   ( unittest.TestCase						):






			@slow
			def       UpperCamelCase_							(			self     ):
									# Test exact values at the last hidden layer
									_lowercase									=				RemBertModel.from_pretrained("""google/rembert"""     )
									_lowercase									=				torch.tensor([[312, 56_498, 313, 2_125, 313]]     )
									_lowercase									=				torch.tensor([[0, 0, 0, 1, 1]]     )
									with torch.no_grad():
															_lowercase									=				model(__UpperCamelCase							, token_type_ids=__UpperCamelCase							, output_hidden_states=__UpperCamelCase     )

									_lowercase									=				1_152

									_lowercase									=				torch.Size((1, 5, hidden_size)     )
									self.assertEqual(output["""last_hidden_state"""].shape							, __UpperCamelCase     )

									_lowercase									=				torch.tensor(
									    [
									        [
									            [0.07_54, -0.20_22, 0.19_04],
									            [-0.33_54, -0.36_92, -0.47_91],
									            [-0.23_14, -0.67_29, -0.07_49],
									            [-0.03_96, -0.31_05, -0.42_34],
									            [-0.15_71, -0.05_25, 0.53_53],
									        ]
									    ]     )

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

import abc
import copy
import dataclasses
from dataclasses import dataclass
from typing import ClassVar, Dict, Type, TypeVar

from ..features import Features


A		:   Optional[int]         =    TypeVar('''T''', bound='''TaskTemplate''')



@dataclass(frozen=_a						)
class       a_   ( abc.ABC						):
			# `task` is not a ClassVar since we want it to be part of the `asdict` output for JSON serialization
			a  : str
			a  : ClassVar[Features]
			a  : ClassVar[Features]






			def       UpperCamelCase_							(			self							, __UpperCamelCase     ):
									return copy.deepcopy(self     )






			@property
			def       UpperCamelCase_							(			self     ):
									return Features(**self.input_schema							, **self.label_schema     )






			@property
			@abc.abstractmethod
			def       UpperCamelCase_							(			self     ):
									raise NotImplementedError






			@classmethod
			def       UpperCamelCase_							(			cls							, __UpperCamelCase     ):
									_lowercase									=				{f.name for f in dataclasses.fields(cls     )}
									return cls(**{k: v for k, v in template_dict.items() if k in field_names}     )