



import unittest

from transformers import MraConfig, 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


if is_torch_available():
 import torch

 from transformers import (
     MraForMaskedLM,
     MraForMultipleChoice,
     MraForQuestionAnswering,
     MraForSequenceClassification,
     MraForTokenClassification,
     MraModel,
 )
 from transformers.models.mra.modeling_mra import MRA_PRETRAINED_MODEL_ARCHIVE_LIST

class    _A	:



       def __init__(		self							:					Any					,		__SCREAMING_SNAKE_CASE							:					Dict					,		__SCREAMING_SNAKE_CASE							:					Union[str, Any]=2					,		__SCREAMING_SNAKE_CASE							:					List[str]=8					,		__SCREAMING_SNAKE_CASE							:					Union[str, Any]=True					,		__SCREAMING_SNAKE_CASE							:					Optional[int]=True					,		__SCREAMING_SNAKE_CASE							:					Optional[int]=True					,		__SCREAMING_SNAKE_CASE							:					int=True					,		__SCREAMING_SNAKE_CASE							:					str=99					,		__SCREAMING_SNAKE_CASE							:					str=16					,		__SCREAMING_SNAKE_CASE							:					Optional[Any]=5					,		__SCREAMING_SNAKE_CASE							:					Optional[Any]=2					,		__SCREAMING_SNAKE_CASE							:					Optional[int]=36					,		__SCREAMING_SNAKE_CASE							:					Tuple="gelu"					,		__SCREAMING_SNAKE_CASE							:					Optional[Any]=0.0					,		__SCREAMING_SNAKE_CASE							:					List[Any]=0.0					,		__SCREAMING_SNAKE_CASE							:					Union[str, Any]=512					,		__SCREAMING_SNAKE_CASE							:					Tuple=16					,		__SCREAMING_SNAKE_CASE							:					List[Any]=2					,		__SCREAMING_SNAKE_CASE							:					str=0.02					,		__SCREAMING_SNAKE_CASE							:					List[str]=3					,		__SCREAMING_SNAKE_CASE							:					Union[str, Any]=4					,		__SCREAMING_SNAKE_CASE							:					str=None					,		):


           '''simple docstring'''

           __a							=    parent
           __a							=    batch_size
           __a							=    seq_length
           __a							=    is_training
           __a							=    use_input_mask
           __a							=    use_token_type_ids
           __a							=    use_labels
           __a							=    vocab_size
           __a							=    hidden_size
           __a							=    num_hidden_layers
           __a							=    num_attention_heads
           __a							=    intermediate_size
           __a							=    hidden_act
           __a							=    hidden_dropout_prob
           __a							=    attention_probs_dropout_prob
           __a							=    max_position_embeddings
           __a							=    type_vocab_size
           __a							=    type_sequence_label_size
           __a							=    initializer_range
           __a							=    num_labels
           __a							=    num_choices
           __a							=    scope



       def 		_lowerCamelCase      (		self							:					List[str]):


           '''simple docstring'''

           __a							=    ids_tensor([self.batch_size, self.seq_length]					,		self.vocab_size)

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

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

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

           __a							=    self.get_config()

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



       def 		_lowerCamelCase      (		self							:					Dict):


           '''simple docstring'''

           return MraConfig(
               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=__SCREAMING_SNAKE_CASE					,		initializer_range=self.initializer_range					,		)



       def 		_lowerCamelCase      (		self							:					Optional[Any]):


           '''simple docstring'''

           __a							=    self.get_config()
           __a							=    300
           return config



       def 		_lowerCamelCase      (		self							:					Union[str, Any]):


           '''simple docstring'''

           (
               (
               __a       
           )	,       (
               __a       
           )	,       (
               __a       
           )	,       (
               __a       
           )	,       (
               __a       
           )	,       (
               __a       
           )	,       (
               __a       
           )	,              
           )							=    self.prepare_config_and_inputs()

           __a							=    True
           __a							=    floats_tensor([self.batch_size, self.seq_length, self.hidden_size])
           __a							=    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							:					int					,		__SCREAMING_SNAKE_CASE							:					Optional[Any]					,		__SCREAMING_SNAKE_CASE							:					Dict					,		__SCREAMING_SNAKE_CASE							:					Optional[int]					,		__SCREAMING_SNAKE_CASE							:					Optional[int]					,		__SCREAMING_SNAKE_CASE							:					Optional[Any]					,		__SCREAMING_SNAKE_CASE							:					List[str]					,		__SCREAMING_SNAKE_CASE							:					List[str]):


           '''simple docstring'''

           __a							=    MraModel(config=__SCREAMING_SNAKE_CASE)
           model.to(__SCREAMING_SNAKE_CASE)
           model.eval()
           __a							=    model(__SCREAMING_SNAKE_CASE					,		attention_mask=__SCREAMING_SNAKE_CASE					,		token_type_ids=__SCREAMING_SNAKE_CASE)
           __a							=    model(__SCREAMING_SNAKE_CASE					,		token_type_ids=__SCREAMING_SNAKE_CASE)
           __a							=    model(__SCREAMING_SNAKE_CASE)
           self.parent.assertEqual(result.last_hidden_state.shape					,		(self.batch_size, self.seq_length, self.hidden_size))



       def 		_lowerCamelCase      (		self							:					Optional[int]					,		__SCREAMING_SNAKE_CASE							:					List[str]					,		__SCREAMING_SNAKE_CASE							:					int					,		__SCREAMING_SNAKE_CASE							:					Tuple					,		__SCREAMING_SNAKE_CASE							:					Optional[int]					,		__SCREAMING_SNAKE_CASE							:					Union[str, Any]					,		__SCREAMING_SNAKE_CASE							:					List[str]					,		__SCREAMING_SNAKE_CASE							:					Optional[int]					,		__SCREAMING_SNAKE_CASE							:					Any					,		__SCREAMING_SNAKE_CASE							:					Any					,		):


           '''simple docstring'''

           __a							=    True
           __a							=    MraModel(__SCREAMING_SNAKE_CASE)
           model.to(__SCREAMING_SNAKE_CASE)
           model.eval()
           __a							=    model(
               __SCREAMING_SNAKE_CASE					,		attention_mask=__SCREAMING_SNAKE_CASE					,		token_type_ids=__SCREAMING_SNAKE_CASE					,		encoder_hidden_states=__SCREAMING_SNAKE_CASE					,		encoder_attention_mask=__SCREAMING_SNAKE_CASE					,		)
           __a							=    model(
               __SCREAMING_SNAKE_CASE					,		attention_mask=__SCREAMING_SNAKE_CASE					,		token_type_ids=__SCREAMING_SNAKE_CASE					,		encoder_hidden_states=__SCREAMING_SNAKE_CASE					,		)
           __a							=    model(__SCREAMING_SNAKE_CASE					,		attention_mask=__SCREAMING_SNAKE_CASE					,		token_type_ids=__SCREAMING_SNAKE_CASE)
           self.parent.assertEqual(result.last_hidden_state.shape					,		(self.batch_size, self.seq_length, self.hidden_size))



       def 		_lowerCamelCase      (		self							:					int					,		__SCREAMING_SNAKE_CASE							:					Dict					,		__SCREAMING_SNAKE_CASE							:					Any					,		__SCREAMING_SNAKE_CASE							:					Tuple					,		__SCREAMING_SNAKE_CASE							:					int					,		__SCREAMING_SNAKE_CASE							:					List[Any]					,		__SCREAMING_SNAKE_CASE							:					List[str]					,		__SCREAMING_SNAKE_CASE							:					Optional[int]):


           '''simple docstring'''

           __a							=    MraForMaskedLM(config=__SCREAMING_SNAKE_CASE)
           model.to(__SCREAMING_SNAKE_CASE)
           model.eval()
           __a							=    model(__SCREAMING_SNAKE_CASE					,		attention_mask=__SCREAMING_SNAKE_CASE					,		token_type_ids=__SCREAMING_SNAKE_CASE					,		labels=__SCREAMING_SNAKE_CASE)
           self.parent.assertEqual(result.logits.shape					,		(self.batch_size, self.seq_length, self.vocab_size))



       def 		_lowerCamelCase      (		self							:					Any					,		__SCREAMING_SNAKE_CASE							:					Tuple					,		__SCREAMING_SNAKE_CASE							:					Optional[Any]					,		__SCREAMING_SNAKE_CASE							:					List[Any]					,		__SCREAMING_SNAKE_CASE							:					Dict					,		__SCREAMING_SNAKE_CASE							:					Optional[int]					,		__SCREAMING_SNAKE_CASE							:					Union[str, Any]					,		__SCREAMING_SNAKE_CASE							:					Union[str, Any]):


           '''simple docstring'''

           __a							=    MraForQuestionAnswering(config=__SCREAMING_SNAKE_CASE)
           model.to(__SCREAMING_SNAKE_CASE)
           model.eval()
           __a							=    model(
               __SCREAMING_SNAKE_CASE					,		attention_mask=__SCREAMING_SNAKE_CASE					,		token_type_ids=__SCREAMING_SNAKE_CASE					,		start_positions=__SCREAMING_SNAKE_CASE					,		end_positions=__SCREAMING_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 		_lowerCamelCase      (		self							:					str					,		__SCREAMING_SNAKE_CASE							:					List[str]					,		__SCREAMING_SNAKE_CASE							:					List[str]					,		__SCREAMING_SNAKE_CASE							:					str					,		__SCREAMING_SNAKE_CASE							:					Any					,		__SCREAMING_SNAKE_CASE							:					Optional[int]					,		__SCREAMING_SNAKE_CASE							:					Any					,		__SCREAMING_SNAKE_CASE							:					Dict):


           '''simple docstring'''

           __a							=    self.num_labels
           __a							=    MraForSequenceClassification(__SCREAMING_SNAKE_CASE)
           model.to(__SCREAMING_SNAKE_CASE)
           model.eval()
           __a							=    model(__SCREAMING_SNAKE_CASE					,		attention_mask=__SCREAMING_SNAKE_CASE					,		token_type_ids=__SCREAMING_SNAKE_CASE					,		labels=__SCREAMING_SNAKE_CASE)
           self.parent.assertEqual(result.logits.shape					,		(self.batch_size, self.num_labels))



       def 		_lowerCamelCase      (		self							:					str					,		__SCREAMING_SNAKE_CASE							:					List[str]					,		__SCREAMING_SNAKE_CASE							:					int					,		__SCREAMING_SNAKE_CASE							:					int					,		__SCREAMING_SNAKE_CASE							:					Any					,		__SCREAMING_SNAKE_CASE							:					List[Any]					,		__SCREAMING_SNAKE_CASE							:					List[str]					,		__SCREAMING_SNAKE_CASE							:					int):


           '''simple docstring'''

           __a							=    self.num_labels
           __a							=    MraForTokenClassification(config=__SCREAMING_SNAKE_CASE)
           model.to(__SCREAMING_SNAKE_CASE)
           model.eval()
           __a							=    model(__SCREAMING_SNAKE_CASE					,		attention_mask=__SCREAMING_SNAKE_CASE					,		token_type_ids=__SCREAMING_SNAKE_CASE					,		labels=__SCREAMING_SNAKE_CASE)
           self.parent.assertEqual(result.logits.shape					,		(self.batch_size, self.seq_length, self.num_labels))



       def 		_lowerCamelCase      (		self							:					Optional[int]					,		__SCREAMING_SNAKE_CASE							:					int					,		__SCREAMING_SNAKE_CASE							:					str					,		__SCREAMING_SNAKE_CASE							:					Optional[int]					,		__SCREAMING_SNAKE_CASE							:					Dict					,		__SCREAMING_SNAKE_CASE							:					Optional[int]					,		__SCREAMING_SNAKE_CASE							:					Tuple					,		__SCREAMING_SNAKE_CASE							:					Optional[Any]):


           '''simple docstring'''

           __a							=    self.num_choices
           __a							=    MraForMultipleChoice(config=__SCREAMING_SNAKE_CASE)
           model.to(__SCREAMING_SNAKE_CASE)
           model.eval()
           __a							=    input_ids.unsqueeze(1).expand(-1					,		self.num_choices					,		-1).contiguous()
           __a							=    token_type_ids.unsqueeze(1).expand(-1					,		self.num_choices					,		-1).contiguous()
           __a							=    input_mask.unsqueeze(1).expand(-1					,		self.num_choices					,		-1).contiguous()
           __a							=    model(
               __SCREAMING_SNAKE_CASE					,		attention_mask=__SCREAMING_SNAKE_CASE					,		token_type_ids=__SCREAMING_SNAKE_CASE					,		labels=__SCREAMING_SNAKE_CASE					,		)
           self.parent.assertEqual(result.logits.shape					,		(self.batch_size, self.num_choices))





       def 		_lowerCamelCase      (		self							:					Any):


           '''simple docstring'''

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


@require_torch
class    _A	(     __UpperCAmelCase				,unittest.TestCase			):
       UpperCamelCase__  :   Optional[int]						 =						(
           (
               MraModel,
               MraForMaskedLM,
               MraForMultipleChoice,
               MraForQuestionAnswering,
               MraForSequenceClassification,
               MraForTokenClassification,
           )
           if is_torch_available()
           else ()
       )
       UpperCamelCase__  :   Union[str, Any]						 =						False
       UpperCamelCase__  :   str						 =						False
       UpperCamelCase__  :   Union[str, Any]						 =						False
       UpperCamelCase__  :   Union[str, Any]						 =						False

       UpperCamelCase__  :   Any						 =						()



       def 		_lowerCamelCase      (		self							:					Dict):


           '''simple docstring'''

           __a							=    MraModelTester(self)
           __a							=    ConfigTester(self					,		config_class=__SCREAMING_SNAKE_CASE					,		hidden_size=37)



       def 		_lowerCamelCase      (		self							:					List[Any]):


           '''simple docstring'''

           self.config_tester.run_common_tests()



       def 		_lowerCamelCase      (		self							:					List[Any]):


           '''simple docstring'''

           __a							=    self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*__SCREAMING_SNAKE_CASE)



       def 		_lowerCamelCase      (		self							:					Optional[int]):


           '''simple docstring'''

           __a							=    self.model_tester.prepare_config_and_inputs()
           for type in ["absolute", "relative_key", "relative_key_query"]:
               __a							=    type
               self.model_tester.create_and_check_model(*__SCREAMING_SNAKE_CASE)



       def 		_lowerCamelCase      (		self							:					List[Any]):


           '''simple docstring'''

           __a							=    self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_lm(*__SCREAMING_SNAKE_CASE)



       def 		_lowerCamelCase      (		self							:					Tuple):


           '''simple docstring'''

           __a							=    self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_multiple_choice(*__SCREAMING_SNAKE_CASE)



       def 		_lowerCamelCase      (		self							:					int):


           '''simple docstring'''

           __a							=    self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_question_answering(*__SCREAMING_SNAKE_CASE)



       def 		_lowerCamelCase      (		self							:					List[str]):


           '''simple docstring'''

           __a							=    self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_sequence_classification(*__SCREAMING_SNAKE_CASE)



       def 		_lowerCamelCase      (		self							:					List[str]):


           '''simple docstring'''

           __a							=    self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_token_classification(*__SCREAMING_SNAKE_CASE)



       @slow
       def 		_lowerCamelCase      (		self							:					Optional[int]):


           '''simple docstring'''

           for model_name in MRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               __a							=    MraModel.from_pretrained(__SCREAMING_SNAKE_CASE)
               self.assertIsNotNone(__SCREAMING_SNAKE_CASE)





       @unittest.skip(reason='''MRA does not output attentions''')
       def 		_lowerCamelCase      (		self							:					List[str]):


           '''simple docstring'''

           return


@require_torch
class    _A	(     unittest.TestCase			):



       @slow
       def 		_lowerCamelCase      (		self							:					Optional[int]):


           '''simple docstring'''

           __a							=    MraModel.from_pretrained('''uw-madison/mra-base-512-4''')
           __a							=    torch.arange(256).unsqueeze(0)

           with torch.no_grad():
               __a							=    model(__SCREAMING_SNAKE_CASE)[0]

           __a							=    torch.Size((1, 256, 768))
           self.assertEqual(output.shape					,		__SCREAMING_SNAKE_CASE)

           __a							=    torch.tensor(
               [[[-0.01_40, 0.08_30, -0.03_81], [0.15_46, 0.14_02, 0.02_20], [0.11_62, 0.08_51, 0.01_65]]])

           self.assertTrue(torch.allclose(output[:, :3, :3]					,		__SCREAMING_SNAKE_CASE					,		atol=1E-4))



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


           '''simple docstring'''

           __a							=    MraForMaskedLM.from_pretrained('''uw-madison/mra-base-512-4''')
           __a							=    torch.arange(256).unsqueeze(0)

           with torch.no_grad():
               __a							=    model(__SCREAMING_SNAKE_CASE)[0]

           __a							=    50_265

           __a							=    torch.Size((1, 256, vocab_size))
           self.assertEqual(output.shape					,		__SCREAMING_SNAKE_CASE)

           __a							=    torch.tensor(
               [[[9.25_95, -3.60_38, 11.88_19], [9.38_69, -3.26_93, 11.09_56], [11.85_24, -3.49_38, 13.12_10]]])

           self.assertTrue(torch.allclose(output[:, :3, :3]					,		__SCREAMING_SNAKE_CASE					,		atol=1E-4))





       @slow
       def 		_lowerCamelCase      (		self							:					str):


           '''simple docstring'''

           __a							=    MraForMaskedLM.from_pretrained('''uw-madison/mra-base-4096-8-d3''')
           __a							=    torch.arange(4_096).unsqueeze(0)

           with torch.no_grad():
               __a							=    model(__SCREAMING_SNAKE_CASE)[0]

           __a							=    50_265

           __a							=    torch.Size((1, 4_096, vocab_size))
           self.assertEqual(output.shape					,		__SCREAMING_SNAKE_CASE)

           __a							=    torch.tensor(
               [[[5.47_89, -2.35_64, 7.50_64], [7.90_67, -1.33_69, 9.96_68], [9.07_12, -1.81_06, 7.03_80]]])

           self.assertTrue(torch.allclose(output[:, :3, :3]					,		__SCREAMING_SNAKE_CASE					,		atol=1E-4))










from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


__snake_case						:Tuple						=					{
    '''configuration_luke''': ['''LUKE_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''LukeConfig'''],
    '''tokenization_luke''': ['''LukeTokenizer'''],
}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __snake_case						:Optional[int]						=					[
     '''LUKE_PRETRAINED_MODEL_ARCHIVE_LIST''',
     '''LukeForEntityClassification''',
     '''LukeForEntityPairClassification''',
     '''LukeForEntitySpanClassification''',
     '''LukeForMultipleChoice''',
     '''LukeForQuestionAnswering''',
     '''LukeForSequenceClassification''',
     '''LukeForTokenClassification''',
     '''LukeForMaskedLM''',
     '''LukeModel''',
     '''LukePreTrainedModel''',
 ]


if TYPE_CHECKING:
 from .configuration_luke import LUKE_PRETRAINED_CONFIG_ARCHIVE_MAP, LukeConfig
 from .tokenization_luke import LukeTokenizer

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_luke import (
      LUKE_PRETRAINED_MODEL_ARCHIVE_LIST,
      LukeForEntityClassification,
      LukeForEntityPairClassification,
      LukeForEntitySpanClassification,
      LukeForMaskedLM,
      LukeForMultipleChoice,
      LukeForQuestionAnswering,
      LukeForSequenceClassification,
      LukeForTokenClassification,
      LukeModel,
      LukePreTrainedModel,
  )

else:
 import sys

 __snake_case						:Dict						=					_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)





