
"""simple docstring"""





import enum
import warnings

from ..tokenization_utils import TruncationStrategy
from ..utils import add_end_docstrings, is_tf_available, is_torch_available, logging
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_tf_available():
				import tensorflow as tf

				from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING

if is_torch_available():
				from ..models.auto.modeling_auto import MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING

_a								=						logging.get_logger(__name__)






class        _UpperCAmelCase(       enum.Enum      ):
		lowercase__         = 0
		lowercase__         = 1






@add_end_docstrings(lowerCamelCase      )
class        _UpperCAmelCase(       lowerCamelCase      ):
		lowercase__         = 'generated'



		def __init__(   self							,						*__a							,						**__a)			->			List[Any]:






							'''simple docstring'''
							super().__init__(*__a							,						**__a)

							self.check_model_type(
							    TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING
							    if self.framework == '''tf'''
							    else MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING)



		def 		UpperCAmelCase				(   self							,						__a=None							,						__a=None							,						__a=None							,						__a=None							,						__a=None							,						__a=None							,						**__a							,						)			->			Tuple:






							'''simple docstring'''
							_UpperCamelCase									=    {}
							if truncation is not None:
												_UpperCamelCase									=    truncation

							_UpperCamelCase									=    generate_kwargs

							_UpperCamelCase									=    {}
							if return_tensors is not None and return_type is None:
												_UpperCamelCase									=    ReturnType.TENSORS if return_tensors else ReturnType.TEXT
							if return_type is not None:
												_UpperCamelCase									=    return_type

							if clean_up_tokenization_spaces is not None:
												_UpperCamelCase									=    clean_up_tokenization_spaces

							if stop_sequence is not None:
												_UpperCamelCase									=    self.tokenizer.encode(__a							,						add_special_tokens=__a)
												if len(__a) > 1:
																	warnings.warn(
																	    '''Stopping on a multiple token sequence is not yet supported on transformers. The first token of'''
																	    ''' the stop sequence will be used as the stop sequence string in the interim.''')
												_UpperCamelCase									=    stop_sequence_ids[0]

							return preprocess_params, forward_params, postprocess_params



		def 		UpperCAmelCase				(   self							,						__a							,						__a							,						__a)			->			Any:






							'''simple docstring'''
							return True



		def 		UpperCAmelCase				(   self							,						*__a							,						__a)			->			Any:






							'''simple docstring'''
							_UpperCamelCase									=    self.model.config.prefix if self.model.config.prefix is not None else ''''''
							if isinstance(args[0]							,						__a):
												if self.tokenizer.pad_token_id is None:
																	raise ValueError('''Please make sure that the tokenizer has a pad_token_id when using a batch input''')
												_UpperCamelCase									=    ([prefix + arg for arg in args[0]],)
												_UpperCamelCase									=    True

							elif isinstance(args[0]							,						__a):
												_UpperCamelCase									=    (prefix + args[0],)
												_UpperCamelCase									=    False
							else:
												raise ValueError(
												    F''' `args[0]`: {args[0]} have the wrong format. The should be either of type `str` or type `list`''')
							_UpperCamelCase									=    self.tokenizer(*__a							,						padding=__a							,						truncation=__a							,						return_tensors=self.framework)
							# This is produced by tokenizers but is an invalid generate kwargs
							if "token_type_ids" in inputs:
												del inputs["token_type_ids"]
							return inputs



		def __call__(   self							,						*__a							,						**__a)			->			Optional[int]:






							'''simple docstring'''
							_UpperCamelCase									=    super().__call__(*__a							,						**__a)
							if (
							    isinstance(args[0]							,						__a)
							    and all(isinstance(__a							,						__a) for el in args[0])
							    and all(len(__a) == 1 for res in result)
							):
												return [res[0] for res in result]
							return result



		def 		UpperCAmelCase				(   self							,						__a							,						__a=TruncationStrategy.DO_NOT_TRUNCATE							,						**__a)			->			Dict:






							'''simple docstring'''
							_UpperCamelCase									=    self._parse_and_tokenize(__a							,						truncation=__a							,						**__a)
							return inputs



		def 		UpperCAmelCase				(   self							,						__a							,						**__a)			->			Optional[int]:






							'''simple docstring'''
							if self.framework == "pt":
												_UpperCamelCase					,    _UpperCamelCase									=    model_inputs['''input_ids'''].shape
							elif self.framework == "tf":
												_UpperCamelCase					,    _UpperCamelCase									=    tf.shape(model_inputs['''input_ids''']).numpy()

							_UpperCamelCase									=    generate_kwargs.get('''min_length'''							,						self.model.config.min_length)
							_UpperCamelCase									=    generate_kwargs.get('''max_length'''							,						self.model.config.max_length)
							self.check_inputs(__a							,						generate_kwargs['''min_length''']							,						generate_kwargs['''max_length'''])
							_UpperCamelCase									=    self.model.generate(**__a							,						**__a)
							_UpperCamelCase									=    output_ids.shape[0]
							if self.framework == "pt":
												_UpperCamelCase									=    output_ids.reshape(__a							,						out_b // in_b							,						*output_ids.shape[1:])
							elif self.framework == "tf":
												_UpperCamelCase									=    tf.reshape(__a							,						(in_b, out_b // in_b, *output_ids.shape[1:]))
							return {"output_ids": output_ids}




		def 		UpperCAmelCase				(   self							,						__a							,						__a=ReturnType.TEXT							,						__a=False)			->			Optional[Any]:






							'''simple docstring'''
							_UpperCamelCase									=    []
							for output_ids in model_outputs["output_ids"][0]:
												if return_type == ReturnType.TENSORS:
																	_UpperCamelCase									=    {F'''{self.return_name}_token_ids''': output_ids}
												elif return_type == ReturnType.TEXT:
																	_UpperCamelCase									=    {
																	    F'''{self.return_name}_text''': self.tokenizer.decode(
																	        __a							,						skip_special_tokens=__a							,						clean_up_tokenization_spaces=__a							,						)
																	}
												records.append(__a)
							return records






@add_end_docstrings(lowerCamelCase      )
class        _UpperCAmelCase(       lowerCamelCase      ):
		lowercase__         = 'summary'



		def __call__(   self							,						*__a							,						**__a)			->			Dict:






							'''simple docstring'''
							return super().__call__(*__a							,						**__a)




		def 		UpperCAmelCase				(   self							,						__a							,						__a							,						__a)			->			bool:






							'''simple docstring'''
							if max_length < min_length:
												logger.warning(F'''Your min_length={min_length} must be inferior than your max_length={max_length}.''')

							if input_length < max_length:
												logger.warning(
												    F'''Your max_length is set to {max_length}, but your input_length is only {input_length}. Since this is '''
												    '''a summarization task, where outputs shorter than the input are typically wanted, you might '''
												    F'''consider decreasing max_length manually, e.g. summarizer(\'...\', max_length={input_length//2})''')






@add_end_docstrings(lowerCamelCase      )
class        _UpperCAmelCase(       lowerCamelCase      ):
		lowercase__         = 'translation'



		def 		UpperCAmelCase				(   self							,						__a							,						__a							,						__a)			->			List[str]:






							'''simple docstring'''
							if input_length > 0.9 * max_length:
												logger.warning(
												    F'''Your input_length: {input_length} is bigger than 0.9 * max_length: {max_length}. You might consider '''
												    '''increasing your max_length manually, e.g. translator(\'...\', max_length=400)''')
							return True



		def 		UpperCAmelCase				(   self							,						*__a							,						__a=TruncationStrategy.DO_NOT_TRUNCATE							,						__a=None							,						__a=None)			->			Dict:






							'''simple docstring'''
							if getattr(self.tokenizer							,						'''_build_translation_inputs'''							,						__a):
												return self.tokenizer._build_translation_inputs(
												    *__a							,						return_tensors=self.framework							,						truncation=__a							,						src_lang=__a							,						tgt_lang=__a)
							else:
												return super()._parse_and_tokenize(*__a							,						truncation=__a)



		def 		UpperCAmelCase				(   self							,						__a=None							,						__a=None							,						**__a)			->			Union[str, Any]:






							'''simple docstring'''
							_UpperCamelCase					,    _UpperCamelCase					,    _UpperCamelCase									=    super()._sanitize_parameters(**__a)
							if src_lang is not None:
												_UpperCamelCase									=    src_lang
							if tgt_lang is not None:
												_UpperCamelCase									=    tgt_lang
							if src_lang is None and tgt_lang is None:
												# Backward compatibility, direct arguments use is preferred.
												_UpperCamelCase									=    kwargs.get('''task'''							,						self.task)
												_UpperCamelCase									=    task.split('''_''')
												if task and len(__a) == 4:
																	# translation, XX, to YY
																	_UpperCamelCase									=    items[1]
																	_UpperCamelCase									=    items[3]
							return preprocess_params, forward_params, postprocess_params




		def __call__(   self							,						*__a							,						**__a)			->			Dict:






							'''simple docstring'''
							return super().__call__(*__a							,						**__a)














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


