






import argparse
import gc
import json
import os
import shutil
import warnings

import torch

from transformers import LlamaConfig, LlamaForCausalLM, LlamaTokenizer


try:
    from transformers import LlamaTokenizerFast
except ImportError as e:
    warnings.warn(e)
    warnings.warn(
        "The converted tokenizer will be the `slow` tokenizer. To use the fast, update your `tokenizers` library and re-run the tokenizer conversion"
    )
    _SCREAMING_SNAKE_CASE   =  None
_SCREAMING_SNAKE_CASE   =  {
    "7B": 1_10_08,
    "13B": 1_38_24,
    "30B": 1_79_20,
    "65B": 2_20_16,
    "70B": 2_86_72,
}
_SCREAMING_SNAKE_CASE   =  {
    "7B": 1,
    "7Bf": 1,
    "13B": 2,
    "13Bf": 2,
    "30B": 4,
    "65B": 8,
    "70B": 8,
    "70Bf": 8,
}





def        _snake_case   (_snake_case				:       Union[str, Any]      ,    _snake_case				:       str=1      ,    _snake_case				:       Tuple=256)					->	List[Any]:
   return multiple_of * ((int(ffn_dim_multiplier * int(8 * n / 3)) + multiple_of - 1) // multiple_of)





def        _snake_case   (_snake_case				:       Optional[Any])					->	Optional[int]:
   with open(_snake_case      ,    'r') as f:
      return json.load(_snake_case)





def        _snake_case   (_snake_case				:       Dict      ,    _snake_case				:       List[str])					->	Any:
   with open(_snake_case      ,    'w') as f:
      json.dump(_snake_case      ,    _snake_case)





def        _snake_case   (_snake_case				:       Union[str, Any]      ,    _snake_case				:       Any      ,    _snake_case				:       List[Any]      ,    _snake_case				:       str=True)					->	int:
   os.makedirs(_snake_case      ,    exist_ok=_snake_case)
   _lowercase              =os.path.join(_snake_case      ,    'tmp')
   os.makedirs(_snake_case      ,    exist_ok=_snake_case)

   _lowercase              =read_json(os.path.join(_snake_case      ,    'params.json'))
   _lowercase              =NUM_SHARDS[model_size]
   _lowercase              =params['n_layers']
   _lowercase              =params['n_heads']
   _lowercase              =n_heads // num_shards
   _lowercase              =params['dim']
   _lowercase              =dim // n_heads
   _lowercase              =10_000.0
   _lowercase              =1.0 / (base ** (torch.arange(0      ,    _snake_case      ,    2).float() / dims_per_head))

   if "n_kv_heads" in params:
      _lowercase              =params['n_kv_heads']  # for GQA / MQA
      _lowercase              =n_heads_per_shard // num_key_value_heads
      _lowercase              =dim // num_key_value_heads
   else:  # compatibility with other checkpoints
      _lowercase              =n_heads
      _lowercase              =n_heads_per_shard
      _lowercase              =dim

   # permute for sliced rotary
   def permute(_snake_case				:       str      ,    _snake_case				:       int=n_heads      ,    _snake_case				:       Dict=dim      ,    _snake_case				:       Optional[int]=dim):
      return w.view(_snake_case      ,    dima // n_heads // 2      ,    2      ,    _snake_case).transpose(1      ,    2).reshape(_snake_case      ,    _snake_case)

   print(f'''Fetching all parameters from the checkpoint at {input_base_path}.''')
   # Load weights
   if model_size == "7B":
      # Not sharded
      # (The sharded implementation would also work, but this is simpler.)
      _lowercase              =torch.load(os.path.join(_snake_case      ,    'consolidated.00.pth')      ,    map_location='cpu')
   else:
      # Sharded
      _lowercase              =[
          torch.load(os.path.join(_snake_case      ,    f'''consolidated.{i:02d}.pth''')      ,    map_location='cpu')
          for i in range(_snake_case)
      ]
   _lowercase              =0
   _lowercase              ={'weight_map': {}}
   for layer_i in range(_snake_case):
      _lowercase              =f'''pytorch_model-{layer_i + 1}-of-{n_layers + 1}.bin'''
      if model_size == "7B":
         # Unsharded
         _lowercase              ={
             f'''model.layers.{layer_i}.self_attn.q_proj.weight''': permute(
                 loaded[f'''layers.{layer_i}.attention.wq.weight''']),
             f'''model.layers.{layer_i}.self_attn.k_proj.weight''': permute(
                 loaded[f'''layers.{layer_i}.attention.wk.weight''']),
             f'''model.layers.{layer_i}.self_attn.v_proj.weight''': loaded[f'''layers.{layer_i}.attention.wv.weight'''],
             f'''model.layers.{layer_i}.self_attn.o_proj.weight''': loaded[f'''layers.{layer_i}.attention.wo.weight'''],
             f'''model.layers.{layer_i}.mlp.gate_proj.weight''': loaded[f'''layers.{layer_i}.feed_forward.w1.weight'''],
             f'''model.layers.{layer_i}.mlp.down_proj.weight''': loaded[f'''layers.{layer_i}.feed_forward.w2.weight'''],
             f'''model.layers.{layer_i}.mlp.up_proj.weight''': loaded[f'''layers.{layer_i}.feed_forward.w3.weight'''],
             f'''model.layers.{layer_i}.input_layernorm.weight''': loaded[f'''layers.{layer_i}.attention_norm.weight'''],
             f'''model.layers.{layer_i}.post_attention_layernorm.weight''': loaded[f'''layers.{layer_i}.ffn_norm.weight'''],
         }
      else:
         # Sharded
         # Note that attention.w{q,k,v,o}, feed_fordward.w[1,2,3], attention_norm.weight and ffn_norm.weight share
         # the same storage object, saving attention_norm and ffn_norm will save other weights too, which is
         # redundant as other weights will be stitched from multiple shards. To avoid that, they are cloned.

         _lowercase              ={
             f'''model.layers.{layer_i}.input_layernorm.weight''': loaded[0][
                 f'''layers.{layer_i}.attention_norm.weight'''
             ].clone(),
             f'''model.layers.{layer_i}.post_attention_layernorm.weight''': loaded[0][
                 f'''layers.{layer_i}.ffn_norm.weight'''
             ].clone(),
         }
         _lowercase              =permute(
             torch.cat(
                 [
                     loaded[i][f'''layers.{layer_i}.attention.wq.weight'''].view(_snake_case      ,    _snake_case      ,    _snake_case)
                     for i in range(_snake_case)
                 ]      ,    dim=0      ,    ).reshape(_snake_case      ,    _snake_case))
         _lowercase              =permute(
             torch.cat(
                 [
                     loaded[i][f'''layers.{layer_i}.attention.wk.weight'''].view(
                         _snake_case      ,    _snake_case      ,    _snake_case)
                     for i in range(_snake_case)
                 ]      ,    dim=0      ,    ).reshape(_snake_case      ,    _snake_case)      ,    _snake_case      ,    _snake_case      ,    _snake_case      ,    )
         _lowercase              =torch.cat(
             [
                 loaded[i][f'''layers.{layer_i}.attention.wv.weight'''].view(
                     _snake_case      ,    _snake_case      ,    _snake_case)
                 for i in range(_snake_case)
             ]      ,    dim=0      ,    ).reshape(_snake_case      ,    _snake_case)

         _lowercase              =torch.cat(
             [loaded[i][f'''layers.{layer_i}.attention.wo.weight'''] for i in range(_snake_case)]      ,    dim=1)
         _lowercase              =torch.cat(
             [loaded[i][f'''layers.{layer_i}.feed_forward.w1.weight'''] for i in range(_snake_case)]      ,    dim=0)
         _lowercase              =torch.cat(
             [loaded[i][f'''layers.{layer_i}.feed_forward.w2.weight'''] for i in range(_snake_case)]      ,    dim=1)
         _lowercase              =torch.cat(
             [loaded[i][f'''layers.{layer_i}.feed_forward.w3.weight'''] for i in range(_snake_case)]      ,    dim=0)

      _lowercase              =inv_freq
      for k, v in state_dict.items():
         _lowercase              =filename
         param_count += v.numel()
      torch.save(_snake_case      ,    os.path.join(_snake_case      ,    _snake_case))

   _lowercase              =f'''pytorch_model-{n_layers + 1}-of-{n_layers + 1}.bin'''
   if model_size == "7B":
      # Unsharded
      _lowercase              ={
          'model.embed_tokens.weight': loaded['tok_embeddings.weight'],
          'model.norm.weight': loaded['norm.weight'],
          'lm_head.weight': loaded['output.weight'],
      }
   else:
      _lowercase              ={
          'model.norm.weight': loaded[0]['norm.weight'],
          'model.embed_tokens.weight': torch.cat(
              [loaded[i]['tok_embeddings.weight'] for i in range(_snake_case)]      ,    dim=1),
          'lm_head.weight': torch.cat([loaded[i]['output.weight'] for i in range(_snake_case)]      ,    dim=0),
      }

   for k, v in state_dict.items():
      _lowercase              =filename
      param_count += v.numel()
   torch.save(_snake_case      ,    os.path.join(_snake_case      ,    _snake_case))

   # Write configs
   _lowercase              ={'total_size': param_count * 2}
   write_json(_snake_case      ,    os.path.join(_snake_case      ,    'pytorch_model.bin.index.json'))
   _lowercase              =params['ffn_dim_multiplier'] if 'ffn_dim_multiplier' in params else 1
   _lowercase              =params['multiple_of'] if 'multiple_of' in params else 256
   _lowercase              =LlamaConfig(
       hidden_size=_snake_case      ,    intermediate_size=compute_intermediate_size(_snake_case      ,    _snake_case      ,    _snake_case)      ,    num_attention_heads=params['n_heads']      ,    num_hidden_layers=params['n_layers']      ,    rms_norm_eps=params['norm_eps']      ,    num_key_value_heads=_snake_case      ,    )
   config.save_pretrained(_snake_case)

   # Make space so we can load the model properly now.
   del state_dict
   del loaded
   gc.collect()

   print('Loading the checkpoint in a Llama model.')
   _lowercase              =LlamaForCausalLM.from_pretrained(_snake_case      ,    torch_dtype=torch.floataa      ,    low_cpu_mem_usage=_snake_case)
   # Avoid saving this as part of the config.
   del model.config._name_or_path

   print('Saving in the Transformers format.')
   model.save_pretrained(_snake_case      ,    safe_serialization=_snake_case)
   shutil.rmtree(_snake_case)





def        _snake_case   (_snake_case				:       Optional[Any]      ,    _snake_case				:       int)					->	List[Any]:
   # Initialize the tokenizer based on the `spm` model
   _lowercase              =LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
   print(f'''Saving a {tokenizer_class.__name__} to {tokenizer_path}.''')
   _lowercase              =tokenizer_class(_snake_case)
   tokenizer.save_pretrained(_snake_case)





def        _snake_case   ()					->	Optional[Any]:
   _lowercase              =argparse.ArgumentParser()
   parser.add_argument(
       '--input_dir'      ,    help='Location of LLaMA weights, which contains tokenizer.model and model folders'      ,    )
   parser.add_argument(
       '--model_size'      ,    choices=['7B', '7Bf', '13B', '13Bf', '30B', '65B', '70B', '70Bf', 'tokenizer_only']      ,    )
   parser.add_argument(
       '--output_dir'      ,    help='Location to write HF model and tokenizer'      ,    )
   parser.add_argument('--safe_serialization'      ,    type=_snake_case      ,    help='Whether or not to save using `safetensors`.')
   _lowercase              =parser.parse_args()
   if args.model_size != "tokenizer_only":
      write_model(
          model_path=args.output_dir      ,    input_base_path=os.path.join(args.input_dir      ,    args.model_size)      ,    model_size=args.model_size      ,    safe_serialization=args.safe_serialization      ,    )
   _lowercase              =os.path.join(args.input_dir      ,    'tokenizer.model')
   write_tokenizer(args.output_dir      ,    _snake_case)


if __name__ == "__main__":
    main()










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 					SCREAMING_SNAKE_CASE_  :
       """simple docstring"""


       def __init__(				self	:str,      snake_case	:Any,      snake_case	:str=2,      snake_case	:int=8,      snake_case	:List[Any]=True,      snake_case	:str=True,      snake_case	:Optional[int]=True,      snake_case	:str=True,      snake_case	:Optional[Any]=99,      snake_case	:Tuple=16,      snake_case	:Optional[Any]=5,      snake_case	:Optional[Any]=2,      snake_case	:Optional[Any]=36,      snake_case	:List[str]="gelu",      snake_case	:str=0.0,      snake_case	:List[Any]=0.0,      snake_case	:List[str]=512,      snake_case	:Dict=16,      snake_case	:int=2,      snake_case	:Tuple=0.0_2,      snake_case	:List[str]=3,      snake_case	:Any=4,      snake_case	:Union[str, Any]=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	:int):



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





       def 				UpperCamelCase__						(				self	:Optional[Any]):



          """simple docstring"""

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





       def 				UpperCamelCase__						(				self	:Any):



          """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	:Optional[int],      snake_case	:Optional[int],      snake_case	:List[str],      snake_case	:int,      snake_case	:Tuple,      snake_case	:Optional[int],      snake_case	:Dict,      snake_case	:str):



          """simple docstring"""

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





       def 				UpperCamelCase__						(				self	:Any,      snake_case	:Any,      snake_case	:str,      snake_case	:Optional[int],      snake_case	:int,      snake_case	:int,      snake_case	:int,      snake_case	:Optional[Any],      snake_case	:Tuple,      snake_case	:Tuple,      ):



          """simple docstring"""

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





       def 				UpperCamelCase__						(				self	:str,      snake_case	:Optional[int],      snake_case	:Optional[int],      snake_case	:int,      snake_case	:Optional[int],      snake_case	:List[str],      snake_case	:Optional[int],      snake_case	:Optional[Any]):



          """simple docstring"""

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





       def 				UpperCamelCase__						(				self	:Tuple,      snake_case	:Union[str, Any],      snake_case	:Dict,      snake_case	:Any,      snake_case	:Union[str, Any],      snake_case	:List[Any],      snake_case	:Tuple,      snake_case	:Any):



          """simple docstring"""

          _lowercase              =MraForQuestionAnswering(config=snake_case)
          model.to(snake_case)
          model.eval()
          _lowercase              =model(
              snake_case,      attention_mask=snake_case,      token_type_ids=snake_case,      start_positions=snake_case,      end_positions=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 				UpperCamelCase__						(				self	:int,      snake_case	:List[str],      snake_case	:int,      snake_case	:Optional[Any],      snake_case	:Union[str, Any],      snake_case	:str,      snake_case	:List[Any],      snake_case	:Optional[Any]):



          """simple docstring"""

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





       def 				UpperCamelCase__						(				self	:Any,      snake_case	:Optional[int],      snake_case	:List[Any],      snake_case	:Optional[Any],      snake_case	:int,      snake_case	:int,      snake_case	:Dict,      snake_case	:List[Any]):



          """simple docstring"""

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





       def 				UpperCamelCase__						(				self	:Optional[int],      snake_case	:List[str],      snake_case	:Optional[int],      snake_case	:Dict,      snake_case	:str,      snake_case	:Optional[Any],      snake_case	:Tuple,      snake_case	:Dict):



          """simple docstring"""

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





       def 				UpperCamelCase__						(				self	:Optional[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 					SCREAMING_SNAKE_CASE_  (       _a  ,						unittest.TestCase ):
       """simple docstring"""


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

       __lowerCAmelCase    :		Dict					=()





       def 				UpperCamelCase__						(				self	:Any):



          """simple docstring"""

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





       def 				UpperCamelCase__						(				self	:str):



          """simple docstring"""

          self.config_tester.run_common_tests()





       def 				UpperCamelCase__						(				self	:Optional[int]):



          """simple docstring"""

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





       def 				UpperCamelCase__						(				self	:str):



          """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(*snake_case)





       def 				UpperCamelCase__						(				self	:Dict):



          """simple docstring"""

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





       def 				UpperCamelCase__						(				self	:Optional[Any]):



          """simple docstring"""

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





       def 				UpperCamelCase__						(				self	:int):



          """simple docstring"""

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





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





       def 				UpperCamelCase__						(				self	:Optional[int]):



          """simple docstring"""

          _lowercase              =self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_token_classification(*snake_case)





       @slow
       def 				UpperCamelCase__						(				self	:int):



          """simple docstring"""

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





       @unittest.skip(reason='MRA does not output attentions')
       def 				UpperCamelCase__						(				self	:Optional[Any]):



          """simple docstring"""

          return






@require_torch
class 					SCREAMING_SNAKE_CASE_  (       unittest.TestCase ):
       """simple docstring"""


       @slow
       def 				UpperCamelCase__						(				self	:Dict):



          """simple docstring"""

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

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

          _lowercase              =torch.Size((1, 256, 768))
          self.assertEqual(output.shape,      snake_case)

          _lowercase              =torch.tensor(
              [[[-0.0_1_4_0, 0.0_8_3_0, -0.0_3_8_1], [0.1_5_4_6, 0.1_4_0_2, 0.0_2_2_0], [0.1_1_6_2, 0.0_8_5_1, 0.0_1_6_5]]])

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





       @slow
       def 				UpperCamelCase__						(				self	:List[Any]):



          """simple docstring"""

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

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

          _lowercase              =5_0265

          _lowercase              =torch.Size((1, 256, vocab_size))
          self.assertEqual(output.shape,      snake_case)

          _lowercase              =torch.tensor(
              [[[9.2_5_9_5, -3.6_0_3_8, 1_1.8_8_1_9], [9.3_8_6_9, -3.2_6_9_3, 1_1.0_9_5_6], [1_1.8_5_2_4, -3.4_9_3_8, 1_3.1_2_1_0]]])

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





       @slow
       def 				UpperCamelCase__						(				self	:Dict):



          """simple docstring"""

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

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

          _lowercase              =5_0265

          _lowercase              =torch.Size((1, 4096, vocab_size))
          self.assertEqual(output.shape,      snake_case)

          _lowercase              =torch.tensor(
              [[[5.4_7_8_9, -2.3_5_6_4, 7.5_0_6_4], [7.9_0_6_7, -1.3_3_6_9, 9.9_6_6_8], [9.0_7_1_2, -1.8_1_0_6, 7.0_3_8_0]]])

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


