





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







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






     '''simple docstring'''





     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__			       =  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  :Tuple ):






     '''simple docstring'''





     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__			       =  self.get_config()

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







  def   UpperCAmelCase					(      self  :Union[str, 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  :Tuple ):






     '''simple docstring'''





     lowercase__			       =  self.get_config()
     lowercase__			       =  3_00
     return config







  def   UpperCAmelCase					(      self  :Dict ):






     '''simple docstring'''





     (
         (
         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   UpperCAmelCase					(      self  :Union[str, Any]   ,							_lowercase  :Union[str, Any]   ,							_lowercase  :Optional[Any]   ,							_lowercase  :str   ,							_lowercase  :Optional[Any]   ,							_lowercase  :Any   ,							_lowercase  :Union[str, Any]   ,							_lowercase  :str ):






     '''simple docstring'''





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







  def   UpperCAmelCase					(      self  :Optional[Any]   ,							_lowercase  :Optional[Any]   ,							_lowercase  :str   ,							_lowercase  :Optional[Any]   ,							_lowercase  :Any   ,							_lowercase  :str   ,							_lowercase  :List[str]   ,							_lowercase  :List[Any]   ,							_lowercase  :int   ,							_lowercase  :str   ,							):






     '''simple docstring'''





     lowercase__			       =  True
     lowercase__			       =  MraModel(_lowercase )
     model.to(_lowercase )
     model.eval()
     lowercase__			       =  model(
         _lowercase   ,							attention_mask=_lowercase   ,							token_type_ids=_lowercase   ,							encoder_hidden_states=_lowercase   ,							encoder_attention_mask=_lowercase   ,							)
     lowercase__			       =  model(
         _lowercase   ,							attention_mask=_lowercase   ,							token_type_ids=_lowercase   ,							encoder_hidden_states=_lowercase   ,							)
     lowercase__			       =  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  :str   ,							_lowercase  :Any   ,							_lowercase  :Optional[Any]   ,							_lowercase  :List[Any]   ,							_lowercase  :List[Any]   ,							_lowercase  :Dict   ,							_lowercase  :List[Any]   ,							_lowercase  :Dict ):






     '''simple docstring'''





     lowercase__			       =  MraForMaskedLM(config=_lowercase )
     model.to(_lowercase )
     model.eval()
     lowercase__			       =  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  :int   ,							_lowercase  :Union[str, Any]   ,							_lowercase  :Tuple   ,							_lowercase  :List[Any]   ,							_lowercase  :Any   ,							_lowercase  :str   ,							_lowercase  :Dict   ,							_lowercase  :Dict ):






     '''simple docstring'''





     lowercase__			       =  MraForQuestionAnswering(config=_lowercase )
     model.to(_lowercase )
     model.eval()
     lowercase__			       =  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  :Any   ,							_lowercase  :int   ,							_lowercase  :Tuple   ,							_lowercase  :List[str]   ,							_lowercase  :List[Any]   ,							_lowercase  :Optional[int]   ,							_lowercase  :Optional[int]   ,							_lowercase  :Dict ):






     '''simple docstring'''





     lowercase__			       =  self.num_labels
     lowercase__			       =  MraForSequenceClassification(_lowercase )
     model.to(_lowercase )
     model.eval()
     lowercase__			       =  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  :Union[str, Any]   ,							_lowercase  :Optional[Any]   ,							_lowercase  :Union[str, Any]   ,							_lowercase  :int   ,							_lowercase  :Dict   ,							_lowercase  :Union[str, Any]   ,							_lowercase  :str   ,							_lowercase  :Optional[int] ):






     '''simple docstring'''





     lowercase__			       =  self.num_labels
     lowercase__			       =  MraForTokenClassification(config=_lowercase )
     model.to(_lowercase )
     model.eval()
     lowercase__			       =  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  :Dict   ,							_lowercase  :Optional[int]   ,							_lowercase  :Optional[int]   ,							_lowercase  :Dict   ,							_lowercase  :Dict   ,							_lowercase  :int   ,							_lowercase  :Any   ,							_lowercase  :Optional[Any] ):






     '''simple docstring'''





     lowercase__			       =  self.num_choices
     lowercase__			       =  MraForMultipleChoice(config=_lowercase )
     model.to(_lowercase )
     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(
         _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  :Union[str, Any] ):






     '''simple docstring'''





     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  lowerCAmelCase							(       lowercase_							,							unittest.TestCase ):
  __lowerCamelCase		=      (
      (
          MraModel,
          MraForMaskedLM,
          MraForMultipleChoice,
          MraForQuestionAnswering,
          MraForSequenceClassification,
          MraForTokenClassification,
      )
      if is_torch_available()
      else ()
  )
  __lowerCamelCase		=      False
  __lowerCamelCase		=      False
  __lowerCamelCase		=      False
  __lowerCamelCase		=      False

  __lowerCamelCase		=      ()







  def   UpperCAmelCase					(      self  :int ):






     '''simple docstring'''





     lowercase__			       =  MraModelTester(self )
     lowercase__			       =  ConfigTester(self   ,							config_class=_lowercase   ,							hidden_size=37 )







  def   UpperCAmelCase					(      self  :Optional[int] ):






     '''simple docstring'''





     self.config_tester.run_common_tests()







  def   UpperCAmelCase					(      self  :Tuple ):






     '''simple docstring'''





     lowercase__			       =  self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*_lowercase )







  def   UpperCAmelCase					(      self  :Dict ):






     '''simple docstring'''





     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(*_lowercase )







  def   UpperCAmelCase					(      self  :Dict ):






     '''simple docstring'''





     lowercase__			       =  self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_masked_lm(*_lowercase )







  def   UpperCAmelCase					(      self  :int ):






     '''simple docstring'''





     lowercase__			       =  self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_multiple_choice(*_lowercase )







  def   UpperCAmelCase					(      self  :Tuple ):






     '''simple docstring'''





     lowercase__			       =  self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_question_answering(*_lowercase )







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






     '''simple docstring'''





     lowercase__			       =  self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_sequence_classification(*_lowercase )







  def   UpperCAmelCase					(      self  :List[str] ):






     '''simple docstring'''





     lowercase__			       =  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]:
        lowercase__			       =  MraModel.from_pretrained(_lowercase )
        self.assertIsNotNone(_lowercase )







  @unittest.skip(reason="MRA does not output attentions" )
  def   UpperCAmelCase					(      self  :Union[str, Any] ):






     '''simple docstring'''





     return





@require_torch
class  lowerCAmelCase							(       unittest.TestCase ):







  @slow
  def   UpperCAmelCase					(      self  :Union[str, Any] ):






     '''simple docstring'''





     lowercase__			       =  MraModel.from_pretrained("uw-madison/mra-base-512-4" )
     lowercase__			       =  torch.arange(2_56 ).unsqueeze(0 )

     with torch.no_grad():
        lowercase__			       =  model(_lowercase )[0]

     lowercase__			       =  torch.Size((1, 2_56, 7_68) )
     self.assertEqual(output.shape   ,							_lowercase )

     lowercase__			       =  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  :Tuple ):






     '''simple docstring'''





     lowercase__			       =  MraForMaskedLM.from_pretrained("uw-madison/mra-base-512-4" )
     lowercase__			       =  torch.arange(2_56 ).unsqueeze(0 )

     with torch.no_grad():
        lowercase__			       =  model(_lowercase )[0]

     lowercase__			       =  5_02_65

     lowercase__			       =  torch.Size((1, 2_56, vocab_size) )
     self.assertEqual(output.shape   ,							_lowercase )

     lowercase__			       =  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  :str ):






     '''simple docstring'''





     lowercase__			       =  MraForMaskedLM.from_pretrained("uw-madison/mra-base-4096-8-d3" )
     lowercase__			       =  torch.arange(40_96 ).unsqueeze(0 )

     with torch.no_grad():
        lowercase__			       =  model(_lowercase )[0]

     lowercase__			       =  5_02_65

     lowercase__			       =  torch.Size((1, 40_96, vocab_size) )
     self.assertEqual(output.shape   ,							_lowercase )

     lowercase__			       =  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 ) )













from itertools import count


def _A       (   __magic_name__ = 50	):
   lowercase__			       =  [1] * min_block_length

   for n in count(__magic_name__	):
      fill_count_functions.append(1	)

      for block_length in range(__magic_name__  , n + 1	):
         for block_start in range(n - block_length	):
            fill_count_functions[n] += fill_count_functions[
                n - block_start - block_length - 1
            ]

         fill_count_functions[n] += 1

      if fill_count_functions[n] > 100_0000:
         break

   return n


if __name__ == "__main__":
  print(F"""{solution() = }""")






