



import inspect
import unittest

import numpy as np

from transformers import BeitConfig
from transformers.testing_utils import require_flax, require_vision, slow
from transformers.utils import cached_property, is_flax_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor


if is_flax_available():
			import jax

			from transformers import FlaxBeitForImageClassification, FlaxBeitForMaskedImageModeling, FlaxBeitModel

if is_vision_available():
			from PIL import Image

			from transformers import BeitImageProcessor


class 						lowercase    (							unittest.TestCase):
						'''simple docstring'''


						def __init__(  self    :				Optional[Any]       ,	snake_case    :				Tuple       ,	snake_case    :				Optional[int]=100       ,	snake_case    :				Optional[Any]=13       ,	snake_case    :				int=30       ,	snake_case    :				int=2       ,	snake_case    :				Union[str, Any]=3       ,	snake_case    :				Tuple=True       ,	snake_case    :				List[Any]=True       ,	snake_case    :				List[Any]=32       ,	snake_case    :				Union[str, Any]=5       ,	snake_case    :				Optional[Any]=4       ,	snake_case    :				Any=37       ,	snake_case    :				Optional[int]="gelu"       ,	snake_case    :				List[str]=0.1       ,	snake_case    :				Union[str, Any]=0.1       ,	snake_case    :				Tuple=10       ,	snake_case    :				Optional[Any]=0.02       ,	snake_case    :				int=3       ,	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     Dict						=				parent
											SCREAMING_SNAKE_CASE							:     List[Any]						=				vocab_size
											SCREAMING_SNAKE_CASE							:     Optional[Any]						=				batch_size
											SCREAMING_SNAKE_CASE							:     List[str]						=				image_size
											SCREAMING_SNAKE_CASE							:     int						=				patch_size
											SCREAMING_SNAKE_CASE							:     Optional[Any]						=				num_channels
											SCREAMING_SNAKE_CASE							:     Tuple						=				is_training
											SCREAMING_SNAKE_CASE							:     Optional[int]						=				use_labels
											SCREAMING_SNAKE_CASE							:     List[Any]						=				hidden_size
											SCREAMING_SNAKE_CASE							:     Any						=				num_hidden_layers
											SCREAMING_SNAKE_CASE							:     List[str]						=				num_attention_heads
											SCREAMING_SNAKE_CASE							:     Optional[Any]						=				intermediate_size
											SCREAMING_SNAKE_CASE							:     str						=				hidden_act
											SCREAMING_SNAKE_CASE							:     List[Any]						=				hidden_dropout_prob
											SCREAMING_SNAKE_CASE							:     Tuple						=				attention_probs_dropout_prob
											SCREAMING_SNAKE_CASE							:     Optional[int]						=				type_sequence_label_size
											SCREAMING_SNAKE_CASE							:     Any						=				initializer_range

											# in BeiT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
											SCREAMING_SNAKE_CASE							:     int						=				(image_size // patch_size) ** 2
											SCREAMING_SNAKE_CASE							:     Tuple						=				num_patches + 1







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



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     Any						=				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

											SCREAMING_SNAKE_CASE							:     Optional[Any]						=				None
											if self.use_labels:
																SCREAMING_SNAKE_CASE							:     Dict						=				ids_tensor([self.batch_size]       ,	self.type_sequence_label_size	)

											SCREAMING_SNAKE_CASE							:     str						=				BeitConfig(
											    vocab_size=self.vocab_size       ,	image_size=self.image_size       ,	patch_size=self.patch_size       ,	num_channels=self.num_channels       ,	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       ,	is_decoder=__lowerCAmelCase       ,	initializer_range=self.initializer_range       ,	)

											return config, pixel_values, labels







						def       lowerCamelCase_				(  self    :				Any       ,	snake_case    :				Any       ,	snake_case    :				Tuple       ,	snake_case    :				int	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     int						=				FlaxBeitModel(config=__lowerCAmelCase	)
											SCREAMING_SNAKE_CASE							:     Any						=				model(__lowerCAmelCase	)
											self.parent.assertEqual(result.last_hidden_state.shape       ,	(self.batch_size, self.seq_length, self.hidden_size)	)







						def       lowerCamelCase_				(  self    :				Dict       ,	snake_case    :				Optional[Any]       ,	snake_case    :				List[str]       ,	snake_case    :				Any	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     List[Any]						=				FlaxBeitForMaskedImageModeling(config=__lowerCAmelCase	)
											SCREAMING_SNAKE_CASE							:     Union[str, Any]						=				model(__lowerCAmelCase	)
											self.parent.assertEqual(result.logits.shape       ,	(self.batch_size, self.seq_length - 1, self.vocab_size)	)







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



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     int						=				self.type_sequence_label_size
											SCREAMING_SNAKE_CASE							:     Tuple						=				FlaxBeitForImageClassification(config=__lowerCAmelCase	)
											SCREAMING_SNAKE_CASE							:     Tuple						=				model(__lowerCAmelCase	)
											self.parent.assertEqual(result.logits.shape       ,	(self.batch_size, self.type_sequence_label_size)	)

											# test greyscale images
											SCREAMING_SNAKE_CASE							:     str						=				1
											SCREAMING_SNAKE_CASE							:     Tuple						=				FlaxBeitForImageClassification(__lowerCAmelCase	)

											SCREAMING_SNAKE_CASE							:     Union[str, Any]						=				floats_tensor([self.batch_size, 1, self.image_size, self.image_size]	)
											SCREAMING_SNAKE_CASE							:     Union[str, Any]						=				model(__lowerCAmelCase	)







						def       lowerCamelCase_				(  self    :				Dict	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     List[str]						=				self.prepare_config_and_inputs()
											(
											    (
											    SCREAMING_SNAKE_CASE					
											)					,  (
											    SCREAMING_SNAKE_CASE					
											)					,  (
											    SCREAMING_SNAKE_CASE					
											)					,  					
											)							:     Any						=				config_and_inputs
											SCREAMING_SNAKE_CASE							:     List[str]						=				{'pixel_values': pixel_values}
											return config, inputs_dict


@require_flax
class 						lowercase    (							lowerCAmelCase_      ,    unittest.TestCase):
						'''simple docstring'''


						UpperCAmelCase			:		Union[str, Any]		=		(
						    (FlaxBeitModel, FlaxBeitForImageClassification, FlaxBeitForMaskedImageModeling) if is_flax_available() else ()
						)







						def       lowerCamelCase_				(  self    :				Dict	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     List[str]						=				FlaxBeitModelTester(self	)
											SCREAMING_SNAKE_CASE							:     int						=				ConfigTester(self       ,	config_class=__lowerCAmelCase       ,	has_text_modality=__lowerCAmelCase       ,	hidden_size=37	)







						def       lowerCamelCase_				(  self    :				Tuple	):



											'''simple docstring'''


											self.config_tester.run_common_tests()







						def       lowerCamelCase_				(  self    :				List[str]	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE					,  SCREAMING_SNAKE_CASE							:     List[str]						=				self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																SCREAMING_SNAKE_CASE							:     List[str]						=				model_class(__lowerCAmelCase	)
																SCREAMING_SNAKE_CASE							:     Optional[int]						=				inspect.signature(model.__call__	)
																# signature.parameters is an OrderedDict => so arg_names order is deterministic
																SCREAMING_SNAKE_CASE							:     Optional[int]						=				[*signature.parameters.keys()]

																SCREAMING_SNAKE_CASE							:     str						=				['pixel_values']
																self.assertListEqual(arg_names[:1]       ,	__lowerCAmelCase	)







						def       lowerCamelCase_				(  self    :				Optional[Any]	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE					,  SCREAMING_SNAKE_CASE							:     int						=				self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																with self.subTest(model_class.__name__	):
																					SCREAMING_SNAKE_CASE							:     Tuple						=				self._prepare_for_class(__lowerCAmelCase       ,	__lowerCAmelCase	)
																					SCREAMING_SNAKE_CASE							:     List[Any]						=				model_class(__lowerCAmelCase	)

																					@jax.jit
																					def model_jitted(snake_case    :				Optional[int]       ,	**snake_case    :				Union[str, Any]	):
																										return model(pixel_values=__lowerCAmelCase       ,	**__lowerCAmelCase	)

																					with self.subTest('JIT Enabled'	):
																										SCREAMING_SNAKE_CASE							:     Dict						=				model_jitted(**__lowerCAmelCase	).to_tuple()

																					with self.subTest('JIT Disabled'	):
																										with jax.disable_jit():
																															SCREAMING_SNAKE_CASE							:     str						=				model_jitted(**__lowerCAmelCase	).to_tuple()

																					self.assertEqual(len(__lowerCAmelCase	)       ,	len(__lowerCAmelCase	)	)
																					for jitted_output, output in zip(__lowerCAmelCase       ,	__lowerCAmelCase	):
																										self.assertEqual(jitted_output.shape       ,	output.shape	)







						def       lowerCamelCase_				(  self    :				int	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     Optional[int]						=				self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*__lowerCAmelCase	)







						def       lowerCamelCase_				(  self    :				str	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     Tuple						=				self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_lm(*__lowerCAmelCase	)







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



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     Any						=				self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_image_classification(*__lowerCAmelCase	)







						@slow
						def       lowerCamelCase_				(  self    :				Optional[Any]	):



											'''simple docstring'''


											for model_class_name in self.all_model_classes:
																SCREAMING_SNAKE_CASE							:     Any						=				model_class_name.from_pretrained('microsoft/beit-base-patch16-224'	)
																SCREAMING_SNAKE_CASE							:     List[Any]						=				model(np.ones((1, 3, 224, 224)	)	)
																self.assertIsNotNone(__lowerCAmelCase	)


def        __a			(   )     ->     Union[str, Any]:
					SCREAMING_SNAKE_CASE							:     Tuple						=				Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'   )
					return image



@require_vision
@require_flax
class 						lowercase    (							unittest.TestCase):
						'''simple docstring'''


						@cached_property
						def       lowerCamelCase_				(  self    :				Dict	):



											'''simple docstring'''


											return BeitImageProcessor.from_pretrained('microsoft/beit-base-patch16-224'	) if is_vision_available() else None







						@slow
						def       lowerCamelCase_				(  self    :				Optional[int]	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     Dict						=				FlaxBeitForMaskedImageModeling.from_pretrained('microsoft/beit-base-patch16-224-pt22k'	)

											SCREAMING_SNAKE_CASE							:     int						=				self.default_image_processor
											SCREAMING_SNAKE_CASE							:     List[str]						=				prepare_img()
											SCREAMING_SNAKE_CASE							:     int						=				image_processor(images=__lowerCAmelCase       ,	return_tensors='np'	).pixel_values

											# prepare bool_masked_pos
											SCREAMING_SNAKE_CASE							:     Optional[Any]						=				np.ones((1, 196)       ,	dtype=__lowerCAmelCase	)

											# forward pass
											SCREAMING_SNAKE_CASE							:     Optional[Any]						=				model(pixel_values=__lowerCAmelCase       ,	bool_masked_pos=__lowerCAmelCase	)
											SCREAMING_SNAKE_CASE							:     Dict						=				outputs.logits

											# verify the logits
											SCREAMING_SNAKE_CASE							:     List[Any]						=				(1, 196, 8192)
											self.assertEqual(logits.shape       ,	__lowerCAmelCase	)

											SCREAMING_SNAKE_CASE							:     Optional[Any]						=				np.array(
											    [[-3.2437, 0.5072, -13.9174], [-3.2456, 0.4948, -13.9401], [-3.2033, 0.5121, -13.8550]]	)

											self.assertTrue(np.allclose(logits[bool_masked_pos][:3, :3]       ,	__lowerCAmelCase       ,	atol=1E-2	)	)







						@slow
						def       lowerCamelCase_				(  self    :				List[str]	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     str						=				FlaxBeitForImageClassification.from_pretrained('microsoft/beit-base-patch16-224'	)

											SCREAMING_SNAKE_CASE							:     List[str]						=				self.default_image_processor
											SCREAMING_SNAKE_CASE							:     List[Any]						=				prepare_img()
											SCREAMING_SNAKE_CASE							:     Dict						=				image_processor(images=__lowerCAmelCase       ,	return_tensors='np'	)

											# forward pass
											SCREAMING_SNAKE_CASE							:     str						=				model(**__lowerCAmelCase	)
											SCREAMING_SNAKE_CASE							:     int						=				outputs.logits

											# verify the logits
											SCREAMING_SNAKE_CASE							:     Optional[Any]						=				(1, 1000)
											self.assertEqual(logits.shape       ,	__lowerCAmelCase	)

											SCREAMING_SNAKE_CASE							:     str						=				np.array([-1.2385, -1.0987, -1.0108]	)

											self.assertTrue(np.allclose(logits[0, :3]       ,	__lowerCAmelCase       ,	atol=1E-4	)	)

											SCREAMING_SNAKE_CASE							:     List[Any]						=				281
											self.assertEqual(logits.argmax(-1	).item()       ,	__lowerCAmelCase	)







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



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     Tuple						=				FlaxBeitForImageClassification.from_pretrained('microsoft/beit-large-patch16-224-pt22k-ft22k'	)

											SCREAMING_SNAKE_CASE							:     Optional[Any]						=				self.default_image_processor
											SCREAMING_SNAKE_CASE							:     Union[str, Any]						=				prepare_img()
											SCREAMING_SNAKE_CASE							:     Dict						=				image_processor(images=__lowerCAmelCase       ,	return_tensors='np'	)

											# forward pass
											SCREAMING_SNAKE_CASE							:     Union[str, Any]						=				model(**__lowerCAmelCase	)
											SCREAMING_SNAKE_CASE							:     Tuple						=				outputs.logits

											# verify the logits
											SCREAMING_SNAKE_CASE							:     int						=				(1, 21841)
											self.assertEqual(logits.shape       ,	__lowerCAmelCase	)

											SCREAMING_SNAKE_CASE							:     Dict						=				np.array([1.6881, -0.2787, 0.5901]	)

											self.assertTrue(np.allclose(logits[0, :3]       ,	__lowerCAmelCase       ,	atol=1E-4	)	)

											SCREAMING_SNAKE_CASE							:     int						=				2396
											self.assertEqual(logits.argmax(-1	).item()       ,	__lowerCAmelCase	)
"""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 							a					:


    def __init__(       self						:	Optional[int] ,					__lowerCAmelCase						:	Optional[Any] ,					__lowerCAmelCase						:	str=2 ,					__lowerCAmelCase						:	List[Any]=8 ,					__lowerCAmelCase						:	Dict=True ,					__lowerCAmelCase						:	List[Any]=True ,					__lowerCAmelCase						:	Optional[int]=True ,					__lowerCAmelCase						:	str=True ,					__lowerCAmelCase						:	Optional[Any]=99 ,					__lowerCAmelCase						:	Optional[Any]=16 ,					__lowerCAmelCase						:	Optional[int]=5 ,					__lowerCAmelCase						:	Optional[int]=2 ,					__lowerCAmelCase						:	Optional[int]=36 ,					__lowerCAmelCase						:	Any="gelu" ,					__lowerCAmelCase						:	Any=0.0 ,					__lowerCAmelCase						:	int=0.0 ,					__lowerCAmelCase						:	List[str]=512 ,					__lowerCAmelCase						:	List[str]=16 ,					__lowerCAmelCase						:	int=2 ,					__lowerCAmelCase						:	Any=0.02 ,					__lowerCAmelCase						:	List[str]=3 ,					__lowerCAmelCase						:	str=4 ,					__lowerCAmelCase						:	Dict=None ,					):
      _UpperCAmelCase    			=   parent
      _UpperCAmelCase    			=   batch_size
      _UpperCAmelCase    			=   seq_length
      _UpperCAmelCase    			=   is_training
      _UpperCAmelCase    			=   use_input_mask
      _UpperCAmelCase    			=   use_token_type_ids
      _UpperCAmelCase    			=   use_labels
      _UpperCAmelCase    			=   vocab_size
      _UpperCAmelCase    			=   hidden_size
      _UpperCAmelCase    			=   num_hidden_layers
      _UpperCAmelCase    			=   num_attention_heads
      _UpperCAmelCase    			=   intermediate_size
      _UpperCAmelCase    			=   hidden_act
      _UpperCAmelCase    			=   hidden_dropout_prob
      _UpperCAmelCase    			=   attention_probs_dropout_prob
      _UpperCAmelCase    			=   max_position_embeddings
      _UpperCAmelCase    			=   type_vocab_size
      _UpperCAmelCase    			=   type_sequence_label_size
      _UpperCAmelCase    			=   initializer_range
      _UpperCAmelCase    			=   num_labels
      _UpperCAmelCase    			=   num_choices
      _UpperCAmelCase    			=   scope


    def   lowerCAmelCase_	(       self						:	Optional[int]       ):
      _UpperCAmelCase    			=   ids_tensor([self.batch_size, self.seq_length] ,					self.vocab_size       )

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

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

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

      _UpperCAmelCase    			=   self.get_config()

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


    def   lowerCAmelCase_	(       self						:	Any       ):
      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=__lowerCAmelCase ,					initializer_range=self.initializer_range ,					)


    def   lowerCAmelCase_	(       self						:	int       ):
      _UpperCAmelCase    			=   self.get_config()
      _UpperCAmelCase    			=   300
      return config


    def   lowerCAmelCase_	(       self						:	Any       ):
      (
          (
          _UpperCAmelCase
      )			,					(
          _UpperCAmelCase
      )			,					(
          _UpperCAmelCase
      )			,					(
          _UpperCAmelCase
      )			,					(
          _UpperCAmelCase
      )			,					(
          _UpperCAmelCase
      )			,					(
          _UpperCAmelCase
      )			,					
      )    			=   self.prepare_config_and_inputs()

      _UpperCAmelCase    			=   True
      _UpperCAmelCase    			=   floats_tensor([self.batch_size, self.seq_length, self.hidden_size]       )
      _UpperCAmelCase    			=   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   lowerCAmelCase_	(       self						:	List[str] ,					__lowerCAmelCase						:	Dict ,					__lowerCAmelCase						:	Dict ,					__lowerCAmelCase						:	Any ,					__lowerCAmelCase						:	int ,					__lowerCAmelCase						:	Dict ,					__lowerCAmelCase						:	List[Any] ,					__lowerCAmelCase						:	Tuple       ):
      _UpperCAmelCase    			=   MraModel(config=__lowerCAmelCase       )
      model.to(__lowerCAmelCase       )
      model.eval()
      _UpperCAmelCase    			=   model(__lowerCAmelCase ,					attention_mask=__lowerCAmelCase ,					token_type_ids=__lowerCAmelCase       )
      _UpperCAmelCase    			=   model(__lowerCAmelCase ,					token_type_ids=__lowerCAmelCase       )
      _UpperCAmelCase    			=   model(__lowerCAmelCase       )
      self.parent.assertEqual(result.last_hidden_state.shape ,					(self.batch_size, self.seq_length, self.hidden_size)       )


    def   lowerCAmelCase_	(       self						:	Tuple ,					__lowerCAmelCase						:	Dict ,					__lowerCAmelCase						:	Any ,					__lowerCAmelCase						:	Optional[Any] ,					__lowerCAmelCase						:	int ,					__lowerCAmelCase						:	Optional[int] ,					__lowerCAmelCase						:	Dict ,					__lowerCAmelCase						:	Dict ,					__lowerCAmelCase						:	Optional[int] ,					__lowerCAmelCase						:	Tuple ,					):
      _UpperCAmelCase    			=   True
      _UpperCAmelCase    			=   MraModel(__lowerCAmelCase       )
      model.to(__lowerCAmelCase       )
      model.eval()
      _UpperCAmelCase    			=   model(
          __lowerCAmelCase ,					attention_mask=__lowerCAmelCase ,					token_type_ids=__lowerCAmelCase ,					encoder_hidden_states=__lowerCAmelCase ,					encoder_attention_mask=__lowerCAmelCase ,					)
      _UpperCAmelCase    			=   model(
          __lowerCAmelCase ,					attention_mask=__lowerCAmelCase ,					token_type_ids=__lowerCAmelCase ,					encoder_hidden_states=__lowerCAmelCase ,					)
      _UpperCAmelCase    			=   model(__lowerCAmelCase ,					attention_mask=__lowerCAmelCase ,					token_type_ids=__lowerCAmelCase       )
      self.parent.assertEqual(result.last_hidden_state.shape ,					(self.batch_size, self.seq_length, self.hidden_size)       )


    def   lowerCAmelCase_	(       self						:	Dict ,					__lowerCAmelCase						:	Optional[int] ,					__lowerCAmelCase						:	List[Any] ,					__lowerCAmelCase						:	Optional[int] ,					__lowerCAmelCase						:	str ,					__lowerCAmelCase						:	Dict ,					__lowerCAmelCase						:	int ,					__lowerCAmelCase						:	str       ):
      _UpperCAmelCase    			=   MraForMaskedLM(config=__lowerCAmelCase       )
      model.to(__lowerCAmelCase       )
      model.eval()
      _UpperCAmelCase    			=   model(__lowerCAmelCase ,					attention_mask=__lowerCAmelCase ,					token_type_ids=__lowerCAmelCase ,					labels=__lowerCAmelCase       )
      self.parent.assertEqual(result.logits.shape ,					(self.batch_size, self.seq_length, self.vocab_size)       )


    def   lowerCAmelCase_	(       self						:	int ,					__lowerCAmelCase						:	Dict ,					__lowerCAmelCase						:	Optional[Any] ,					__lowerCAmelCase						:	int ,					__lowerCAmelCase						:	Tuple ,					__lowerCAmelCase						:	List[str] ,					__lowerCAmelCase						:	Any ,					__lowerCAmelCase						:	int       ):
      _UpperCAmelCase    			=   MraForQuestionAnswering(config=__lowerCAmelCase       )
      model.to(__lowerCAmelCase       )
      model.eval()
      _UpperCAmelCase    			=   model(
          __lowerCAmelCase ,					attention_mask=__lowerCAmelCase ,					token_type_ids=__lowerCAmelCase ,					start_positions=__lowerCAmelCase ,					end_positions=__lowerCAmelCase ,					)
      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   lowerCAmelCase_	(       self						:	List[Any] ,					__lowerCAmelCase						:	int ,					__lowerCAmelCase						:	List[str] ,					__lowerCAmelCase						:	Any ,					__lowerCAmelCase						:	Dict ,					__lowerCAmelCase						:	int ,					__lowerCAmelCase						:	Union[str, Any] ,					__lowerCAmelCase						:	int       ):
      _UpperCAmelCase    			=   self.num_labels
      _UpperCAmelCase    			=   MraForSequenceClassification(__lowerCAmelCase       )
      model.to(__lowerCAmelCase       )
      model.eval()
      _UpperCAmelCase    			=   model(__lowerCAmelCase ,					attention_mask=__lowerCAmelCase ,					token_type_ids=__lowerCAmelCase ,					labels=__lowerCAmelCase       )
      self.parent.assertEqual(result.logits.shape ,					(self.batch_size, self.num_labels)       )


    def   lowerCAmelCase_	(       self						:	List[str] ,					__lowerCAmelCase						:	Optional[int] ,					__lowerCAmelCase						:	Union[str, Any] ,					__lowerCAmelCase						:	List[Any] ,					__lowerCAmelCase						:	Optional[int] ,					__lowerCAmelCase						:	Tuple ,					__lowerCAmelCase						:	Optional[int] ,					__lowerCAmelCase						:	Optional[int]       ):
      _UpperCAmelCase    			=   self.num_labels
      _UpperCAmelCase    			=   MraForTokenClassification(config=__lowerCAmelCase       )
      model.to(__lowerCAmelCase       )
      model.eval()
      _UpperCAmelCase    			=   model(__lowerCAmelCase ,					attention_mask=__lowerCAmelCase ,					token_type_ids=__lowerCAmelCase ,					labels=__lowerCAmelCase       )
      self.parent.assertEqual(result.logits.shape ,					(self.batch_size, self.seq_length, self.num_labels)       )


    def   lowerCAmelCase_	(       self						:	Optional[Any] ,					__lowerCAmelCase						:	Any ,					__lowerCAmelCase						:	int ,					__lowerCAmelCase						:	List[Any] ,					__lowerCAmelCase						:	Any ,					__lowerCAmelCase						:	str ,					__lowerCAmelCase						:	List[str] ,					__lowerCAmelCase						:	Optional[Any]       ):
      _UpperCAmelCase    			=   self.num_choices
      _UpperCAmelCase    			=   MraForMultipleChoice(config=__lowerCAmelCase       )
      model.to(__lowerCAmelCase       )
      model.eval()
      _UpperCAmelCase    			=   input_ids.unsqueeze(1       ).expand(-1 ,					self.num_choices ,					-1       ).contiguous()
      _UpperCAmelCase    			=   token_type_ids.unsqueeze(1       ).expand(-1 ,					self.num_choices ,					-1       ).contiguous()
      _UpperCAmelCase    			=   input_mask.unsqueeze(1       ).expand(-1 ,					self.num_choices ,					-1       ).contiguous()
      _UpperCAmelCase    			=   model(
          __lowerCAmelCase ,					attention_mask=__lowerCAmelCase ,					token_type_ids=__lowerCAmelCase ,					labels=__lowerCAmelCase ,					)
      self.parent.assertEqual(result.logits.shape ,					(self.batch_size, self.num_choices)       )


    def   lowerCAmelCase_	(       self						:	Tuple       ):
      _UpperCAmelCase    			=   self.prepare_config_and_inputs()
      (
          (
          _UpperCAmelCase
      )			,					(
          _UpperCAmelCase
      )			,					(
          _UpperCAmelCase
      )			,					(
          _UpperCAmelCase
      )			,					(
          _UpperCAmelCase
      )			,					(
          _UpperCAmelCase
      )			,					(
          _UpperCAmelCase
      )			,					
      )    			=   config_and_inputs
      _UpperCAmelCase    			=   {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
      return config, inputs_dict



@require_torch
class 							a					(		lowerCAmelCase_						,				unittest.TestCase ):
    _snake_case  :	str  				=				(
        (
            MraModel,
            MraForMaskedLM,
            MraForMultipleChoice,
            MraForQuestionAnswering,
            MraForSequenceClassification,
            MraForTokenClassification,
        )
        if is_torch_available()
        else ()
    )
    _snake_case  :	Dict  				=				False
    _snake_case  :	Any  				=				False
    _snake_case  :	List[str]  				=				False
    _snake_case  :	Tuple  				=				False

    _snake_case  :	int  				=				()


    def   lowerCAmelCase_	(       self						:	Union[str, Any]       ):
      _UpperCAmelCase    			=   MraModelTester(self       )
      _UpperCAmelCase    			=   ConfigTester(self ,					config_class=__lowerCAmelCase ,					hidden_size=37       )


    def   lowerCAmelCase_	(       self						:	Optional[Any]       ):
      self.config_tester.run_common_tests()


    def   lowerCAmelCase_	(       self						:	Tuple       ):
      _UpperCAmelCase    			=   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*__lowerCAmelCase       )


    def   lowerCAmelCase_	(       self						:	Any       ):
      _UpperCAmelCase    			=   self.model_tester.prepare_config_and_inputs()
      for type in ["absolute", "relative_key", "relative_key_query"]:
        _UpperCAmelCase    			=   type
        self.model_tester.create_and_check_model(*__lowerCAmelCase       )


    def   lowerCAmelCase_	(       self						:	List[Any]       ):
      _UpperCAmelCase    			=   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_masked_lm(*__lowerCAmelCase       )


    def   lowerCAmelCase_	(       self						:	str       ):
      _UpperCAmelCase    			=   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_multiple_choice(*__lowerCAmelCase       )


    def   lowerCAmelCase_	(       self						:	str       ):
      _UpperCAmelCase    			=   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_question_answering(*__lowerCAmelCase       )


    def   lowerCAmelCase_	(       self						:	List[str]       ):
      _UpperCAmelCase    			=   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_sequence_classification(*__lowerCAmelCase       )


    def   lowerCAmelCase_	(       self						:	List[Any]       ):
      _UpperCAmelCase    			=   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_token_classification(*__lowerCAmelCase       )


    @slow
    def   lowerCAmelCase_	(       self						:	Optional[Any]       ):
      for model_name in MRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        _UpperCAmelCase    			=   MraModel.from_pretrained(__lowerCAmelCase       )
        self.assertIsNotNone(__lowerCAmelCase       )


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



@require_torch
class 							a					(		unittest.TestCase ):


    @slow
    def   lowerCAmelCase_	(       self						:	Dict       ):
      _UpperCAmelCase    			=   MraModel.from_pretrained("""uw-madison/mra-base-512-4"""       )
      _UpperCAmelCase    			=   torch.arange(256       ).unsqueeze(0       )

      with torch.no_grad():
        _UpperCAmelCase    			=   model(__lowerCAmelCase       )[0]

      _UpperCAmelCase    			=   torch.Size((1, 256, 768)       )
      self.assertEqual(output.shape ,					__lowerCAmelCase       )

      _UpperCAmelCase    			=   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] ,					__lowerCAmelCase ,					atol=1e-4       )       )


    @slow
    def   lowerCAmelCase_	(       self						:	str       ):
      _UpperCAmelCase    			=   MraForMaskedLM.from_pretrained("""uw-madison/mra-base-512-4"""       )
      _UpperCAmelCase    			=   torch.arange(256       ).unsqueeze(0       )

      with torch.no_grad():
        _UpperCAmelCase    			=   model(__lowerCAmelCase       )[0]

      _UpperCAmelCase    			=   5_0265

      _UpperCAmelCase    			=   torch.Size((1, 256, vocab_size)       )
      self.assertEqual(output.shape ,					__lowerCAmelCase       )

      _UpperCAmelCase    			=   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] ,					__lowerCAmelCase ,					atol=1e-4       )       )


    @slow
    def   lowerCAmelCase_	(       self						:	Dict       ):
      _UpperCAmelCase    			=   MraForMaskedLM.from_pretrained("""uw-madison/mra-base-4096-8-d3"""       )
      _UpperCAmelCase    			=   torch.arange(4096       ).unsqueeze(0       )

      with torch.no_grad():
        _UpperCAmelCase    			=   model(__lowerCAmelCase       )[0]

      _UpperCAmelCase    			=   5_0265

      _UpperCAmelCase    			=   torch.Size((1, 4096, vocab_size)       )
      self.assertEqual(output.shape ,					__lowerCAmelCase       )

      _UpperCAmelCase    			=   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] ,					__lowerCAmelCase ,					atol=1e-4       )       )
