


'''simple docstring'''

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





  """simple docstring"""




  def __init__(    self			,  UpperCamelCase			,  UpperCamelCase=2			,  UpperCamelCase=8			,  UpperCamelCase=True			,  UpperCamelCase=True			,  UpperCamelCase=True			,  UpperCamelCase=True			,  UpperCamelCase=99			,  UpperCamelCase=16			,  UpperCamelCase=5			,  UpperCamelCase=2			,  UpperCamelCase=36			,  UpperCamelCase="gelu"			,  UpperCamelCase=0.0			,  UpperCamelCase=0.0			,  UpperCamelCase=512			,  UpperCamelCase=16			,  UpperCamelCase=2			,  UpperCamelCase=0.02			,  UpperCamelCase=3			,  UpperCamelCase=4			,  UpperCamelCase=None			,  ):







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






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  ids_tensor([self.batch_size, self.seq_length]			,  self.vocab_size     )

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

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

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

   lowerCamelCase_      =  self.get_config()

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






  def   snake_case		(    self     ):







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






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  self.get_config()
   lowerCamelCase_      =  300
   return config






  def   snake_case		(    self     ):







   """simple docstring"""
   (
       (
       lowerCamelCase_
   ) ,(
       lowerCamelCase_
   ) ,(
       lowerCamelCase_
   ) ,(
       lowerCamelCase_
   ) ,(
       lowerCamelCase_
   ) ,(
       lowerCamelCase_
   ) ,(
       lowerCamelCase_
   ) ,
   )      =  self.prepare_config_and_inputs()

   lowerCamelCase_      =  True
   lowerCamelCase_      =  floats_tensor([self.batch_size, self.seq_length, self.hidden_size]     )
   lowerCamelCase_      =  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   snake_case		(    self			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase     ):







   """simple docstring"""
   lowerCamelCase_      =  MraModel(config=UpperCamelCase     )
   model.to(UpperCamelCase     )
   model.eval()
   lowerCamelCase_      =  model(UpperCamelCase			,  attention_mask=UpperCamelCase			,  token_type_ids=UpperCamelCase     )
   lowerCamelCase_      =  model(UpperCamelCase			,  token_type_ids=UpperCamelCase     )
   lowerCamelCase_      =  model(UpperCamelCase     )
   self.parent.assertEqual(result.last_hidden_state.shape			,  (self.batch_size, self.seq_length, self.hidden_size)     )






  def   snake_case		(    self			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  ):







   """simple docstring"""
   lowerCamelCase_      =  True
   lowerCamelCase_      =  MraModel(UpperCamelCase     )
   model.to(UpperCamelCase     )
   model.eval()
   lowerCamelCase_      =  model(
       UpperCamelCase			,  attention_mask=UpperCamelCase			,  token_type_ids=UpperCamelCase			,  encoder_hidden_states=UpperCamelCase			,  encoder_attention_mask=UpperCamelCase			,  )
   lowerCamelCase_      =  model(
       UpperCamelCase			,  attention_mask=UpperCamelCase			,  token_type_ids=UpperCamelCase			,  encoder_hidden_states=UpperCamelCase			,  )
   lowerCamelCase_      =  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   snake_case		(    self			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase     ):







   """simple docstring"""
   lowerCamelCase_      =  MraForMaskedLM(config=UpperCamelCase     )
   model.to(UpperCamelCase     )
   model.eval()
   lowerCamelCase_      =  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   snake_case		(    self			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase     ):







   """simple docstring"""
   lowerCamelCase_      =  MraForQuestionAnswering(config=UpperCamelCase     )
   model.to(UpperCamelCase     )
   model.eval()
   lowerCamelCase_      =  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   snake_case		(    self			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase     ):







   """simple docstring"""
   lowerCamelCase_      =  self.num_labels
   lowerCamelCase_      =  MraForSequenceClassification(UpperCamelCase     )
   model.to(UpperCamelCase     )
   model.eval()
   lowerCamelCase_      =  model(UpperCamelCase			,  attention_mask=UpperCamelCase			,  token_type_ids=UpperCamelCase			,  labels=UpperCamelCase     )
   self.parent.assertEqual(result.logits.shape			,  (self.batch_size, self.num_labels)     )






  def   snake_case		(    self			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase     ):







   """simple docstring"""
   lowerCamelCase_      =  self.num_labels
   lowerCamelCase_      =  MraForTokenClassification(config=UpperCamelCase     )
   model.to(UpperCamelCase     )
   model.eval()
   lowerCamelCase_      =  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   snake_case		(    self			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase     ):







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







  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  self.prepare_config_and_inputs()
   (
       (
       lowerCamelCase_
   ) ,(
       lowerCamelCase_
   ) ,(
       lowerCamelCase_
   ) ,(
       lowerCamelCase_
   ) ,(
       lowerCamelCase_
   ) ,(
       lowerCamelCase_
   ) ,(
       lowerCamelCase_
   ) ,
   )      =  config_and_inputs
   lowerCamelCase_      =  {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
   return config, inputs_dict





@require_torch
class      snake_case    (    lowercase			,						unittest.TestCase			):





  """simple docstring"""




  _lowerCamelCase		    =  (
      (
          MraModel,
          MraForMaskedLM,
          MraForMultipleChoice,
          MraForQuestionAnswering,
          MraForSequenceClassification,
          MraForTokenClassification,
      )
      if is_torch_available()
      else ()
  )
  _lowerCamelCase		    =  False
  _lowerCamelCase		    =  False
  _lowerCamelCase		    =  False
  _lowerCamelCase		    =  False

  _lowerCamelCase		    =  ()






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  MraModelTester(self     )
   lowerCamelCase_      =  ConfigTester(self			,  config_class=UpperCamelCase			,  hidden_size=37     )






  def   snake_case		(    self     ):







   """simple docstring"""
   self.config_tester.run_common_tests()






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*UpperCamelCase     )






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  self.model_tester.prepare_config_and_inputs()
   for type in ["absolute", "relative_key", "relative_key_query"]:
    lowerCamelCase_      =  type
    self.model_tester.create_and_check_model(*UpperCamelCase     )






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase     )






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_multiple_choice(*UpperCamelCase     )






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_question_answering(*UpperCamelCase     )






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_sequence_classification(*UpperCamelCase     )






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_token_classification(*UpperCamelCase     )






  @slow
  def   snake_case		(    self     ):







   """simple docstring"""
   for model_name in MRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
    lowerCamelCase_      =  MraModel.from_pretrained(UpperCamelCase     )
    self.assertIsNotNone(UpperCamelCase     )







  @unittest.skip(reason="MRA does not output attentions"     )
  def   snake_case		(    self     ):







   """simple docstring"""
   return





@require_torch
class      snake_case    (    unittest.TestCase			):





  """simple docstring"""




  @slow
  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  MraModel.from_pretrained("uw-madison/mra-base-512-4"     )
   lowerCamelCase_      =  torch.arange(256     ).unsqueeze(0     )

   with torch.no_grad():
    lowerCamelCase_      =  model(UpperCamelCase     )[0]

   lowerCamelCase_      =  torch.Size((1, 256, 768)     )
   self.assertEqual(output.shape			,  UpperCamelCase     )

   lowerCamelCase_      =  torch.tensor(
       [[[-0.0_140, 0.0_830, -0.0_381], [0.1_546, 0.1_402, 0.0_220], [0.1_162, 0.0_851, 0.0_165]]]     )

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






  @slow
  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  MraForMaskedLM.from_pretrained("uw-madison/mra-base-512-4"     )
   lowerCamelCase_      =  torch.arange(256     ).unsqueeze(0     )

   with torch.no_grad():
    lowerCamelCase_      =  model(UpperCamelCase     )[0]

   lowerCamelCase_      =  5_0265

   lowerCamelCase_      =  torch.Size((1, 256, vocab_size)     )
   self.assertEqual(output.shape			,  UpperCamelCase     )

   lowerCamelCase_      =  torch.tensor(
       [[[9.2_595, -3.6_038, 11.8_819], [9.3_869, -3.2_693, 11.0_956], [11.8_524, -3.4_938, 13.1_210]]]     )

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







  @slow
  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  MraForMaskedLM.from_pretrained("uw-madison/mra-base-4096-8-d3"     )
   lowerCamelCase_      =  torch.arange(4096     ).unsqueeze(0     )

   with torch.no_grad():
    lowerCamelCase_      =  model(UpperCamelCase     )[0]

   lowerCamelCase_      =  5_0265

   lowerCamelCase_      =  torch.Size((1, 4096, vocab_size)     )
   self.assertEqual(output.shape			,  UpperCamelCase     )

   lowerCamelCase_      =  torch.tensor(
       [[[5.4_789, -2.3_564, 7.5_064], [7.9_067, -1.3_369, 9.9_668], [9.0_712, -1.8_106, 7.0_380]]]     )

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







'''simple docstring'''

a_  :				Any						= """0.21.0"""

from .accelerator import Accelerator
from .big_modeling import (
    cpu_offload,
    cpu_offload_with_hook,
    disk_offload,
    dispatch_model,
    init_empty_weights,
    init_on_device,
    load_checkpoint_and_dispatch,
)
from .data_loader import skip_first_batches
from .launchers import debug_launcher, notebook_launcher
from .state import PartialState
from .utils import (
    DeepSpeedPlugin,
    DistributedDataParallelKwargs,
    DistributedType,
    FullyShardedDataParallelPlugin,
    GradScalerKwargs,
    InitProcessGroupKwargs,
    find_executable_batch_size,
    infer_auto_device_map,
    is_rich_available,
    load_checkpoint_in_model,
    synchronize_rng_states,
)


if is_rich_available():
     from .utils import rich



