

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

   '''simple docstring'''




   def __init__(				self,    lowercase_,    lowercase_=2,    lowercase_=8,    lowercase_=True,    lowercase_=True,    lowercase_=True,    lowercase_=True,    lowercase_=99,    lowercase_=16,    lowercase_=5,    lowercase_=2,    lowercase_=36,    lowercase_="gelu",    lowercase_=0.0,    lowercase_=0.0,    lowercase_=512,    lowercase_=16,    lowercase_=2,    lowercase_=0.02,    lowercase_=3,    lowercase_=4,    lowercase_=None,    )							->    Union[str, Any]:






     """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  _UpperCAmelCase (				self       )							->    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  _UpperCAmelCase (				self       )							->    Any:






     """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=lowercase_,    initializer_range=self.initializer_range,    )



   def  _UpperCAmelCase (				self       )							->    List[str]:






     """simple docstring"""






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



   def  _UpperCAmelCase (				self       )							->    Optional[int]:






     """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  _UpperCAmelCase (				self,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_       )							->    str:






     """simple docstring"""






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



   def  _UpperCAmelCase (				self,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    )							->    str:






     """simple docstring"""






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



   def  _UpperCAmelCase (				self,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_       )							->    Any:






     """simple docstring"""






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



   def  _UpperCAmelCase (				self,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_       )							->    Optional[int]:






     """simple docstring"""






     a__  =MraForQuestionAnswering(config=lowercase_       )
     model.to(lowercase_       )
     model.eval()
     a__  =model(
         lowercase_,    attention_mask=lowercase_,    token_type_ids=lowercase_,    start_positions=lowercase_,    end_positions=lowercase_,    )
     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,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_       )							->    Any:






     """simple docstring"""






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



   def  _UpperCAmelCase (				self,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_       )							->    Any:






     """simple docstring"""






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



   def  _UpperCAmelCase (				self,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_,    lowercase_       )							->    Union[str, Any]:






     """simple docstring"""






     a__  =self.num_choices
     a__  =MraForMultipleChoice(config=lowercase_       )
     model.to(lowercase_       )
     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(
         lowercase_,    attention_mask=lowercase_,    token_type_ids=lowercase_,    labels=lowercase_,    )
     self.parent.assertEqual(result.logits.shape,    (self.batch_size, self.num_choices)       )



   def  _UpperCAmelCase (				self       )							->    str:






     """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 						__magic_name__			(							lowerCamelCase__		,       unittest.TestCase							):

   '''simple docstring'''




   lowerCamelCase__ :	Tuple    =   (
       (
           MraModel,
           MraForMaskedLM,
           MraForMultipleChoice,
           MraForQuestionAnswering,
           MraForSequenceClassification,
           MraForTokenClassification,
       )
       if is_torch_available()
       else ()
   )
   lowerCamelCase__ :	int    =   False
   lowerCamelCase__ :	Optional[int]    =   False
   lowerCamelCase__ :	Optional[Any]    =   False
   lowerCamelCase__ :	Optional[int]    =   False

   lowerCamelCase__ :	int    =   ()



   def  _UpperCAmelCase (				self       )							->    Any:






     """simple docstring"""






     a__  =MraModelTester(self       )
     a__  =ConfigTester(self,    config_class=lowercase_,    hidden_size=37       )



   def  _UpperCAmelCase (				self       )							->    str:






     """simple docstring"""






     self.config_tester.run_common_tests()



   def  _UpperCAmelCase (				self       )							->    Dict:






     """simple docstring"""






     a__  =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*lowercase_       )



   def  _UpperCAmelCase (				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(*lowercase_       )



   def  _UpperCAmelCase (				self       )							->    Union[str, Any]:






     """simple docstring"""






     a__  =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_masked_lm(*lowercase_       )



   def  _UpperCAmelCase (				self       )							->    Optional[int]:






     """simple docstring"""






     a__  =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_multiple_choice(*lowercase_       )



   def  _UpperCAmelCase (				self       )							->    Tuple:






     """simple docstring"""






     a__  =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_question_answering(*lowercase_       )



   def  _UpperCAmelCase (				self       )							->    List[str]:






     """simple docstring"""






     a__  =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_sequence_classification(*lowercase_       )



   def  _UpperCAmelCase (				self       )							->    List[Any]:






     """simple docstring"""






     a__  =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_token_classification(*lowercase_       )



   @slow
   def  _UpperCAmelCase (				self       )							->    int:






     """simple docstring"""






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



   @unittest.skip(reason='''MRA does not output attentions'''       )
   def  _UpperCAmelCase (				self       )							->    int:






     """simple docstring"""






     return







@require_torch
class 						__magic_name__			(							unittest.TestCase							):

   '''simple docstring'''




   @slow
   def  _UpperCAmelCase (				self       )							->    Any:






     """simple docstring"""






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

     with torch.no_grad():
       a__  =model(lowercase_       )[0]

     a__  =torch.Size((1, 256, 768)       )
     self.assertEqual(output.shape,    lowercase_       )

     a__  =torch.tensor(
         [[[-0.0140, 0.0830, -0.0381], [0.1546, 0.1402, 0.0220], [0.1162, 0.0851, 0.0165]]]       )

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



   @slow
   def  _UpperCAmelCase (				self       )							->    int:






     """simple docstring"""






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

     with torch.no_grad():
       a__  =model(lowercase_       )[0]

     a__  =50265

     a__  =torch.Size((1, 256, vocab_size)       )
     self.assertEqual(output.shape,    lowercase_       )

     a__  =torch.tensor(
         [[[9.2595, -3.6038, 11.8819], [9.3869, -3.2693, 11.0956], [11.8524, -3.4938, 13.1210]]]       )

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



   @slow
   def  _UpperCAmelCase (				self       )							->    Any:






     """simple docstring"""






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

     with torch.no_grad():
       a__  =model(lowercase_       )[0]

     a__  =50265

     a__  =torch.Size((1, 4096, vocab_size)       )
     self.assertEqual(output.shape,    lowercase_       )

     a__  =torch.tensor(
         [[[5.4789, -2.3564, 7.5064], [7.9067, -1.3369, 9.9668], [9.0712, -1.8106, 7.0380]]]       )

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



def 			UpperCAmelCase__				(			_A      : dict							):





  '''simple docstring'''







  a__  =set()
  # To detect a back edge, keep track of vertices currently in the recursion stack
  a__  =set()
  return any(
      node not in visited and depth_first_search(_A	,      _A	,      _A	,      _A							)
      for node in graph							)




def 			UpperCAmelCase__				(			_A      : dict	,      _A      : int	,      _A      : set	,      _A      : set							):





  '''simple docstring'''







  visited.add(_A							)
  rec_stk.add(_A							)

  for node in graph[vertex]:
    if node not in visited:
      if depth_first_search(_A	,      _A	,      _A	,      _A							):
        return True
    elif node in rec_stk:
      return True

    # The node needs to be removed from recursion stack before function ends
  rec_stk.remove(_A							)
  return False


if __name__ == "__main__":
       from doctest import testmod

       testmod()
