import argparse
import json
import os

import numpy as np
import PIL
import requests
import tensorflow.keras.applications.efficientnet as efficientnet
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from tensorflow.keras.preprocessing import image

from transformers import (
    EfficientNetConfig,
    EfficientNetForImageClassification,
    EfficientNetImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
__lowercase											=      logging.get_logger(__name__)

__lowercase											=      {
    '''b0''': efficientnet.EfficientNetBa,
    '''b1''': efficientnet.EfficientNetBa,
    '''b2''': efficientnet.EfficientNetBa,
    '''b3''': efficientnet.EfficientNetBa,
    '''b4''': efficientnet.EfficientNetBa,
    '''b5''': efficientnet.EfficientNetBa,
    '''b6''': efficientnet.EfficientNetBa,
    '''b7''': efficientnet.EfficientNetBa,
}

__lowercase											=      {
    '''b0''': {
        '''hidden_dim''': 1280,
        '''width_coef''': 1.0,
        '''depth_coef''': 1.0,
        '''image_size''': 224,
        '''dropout_rate''': 0.2,
        '''dw_padding''': [],
    },
    '''b1''': {
        '''hidden_dim''': 1280,
        '''width_coef''': 1.0,
        '''depth_coef''': 1.1,
        '''image_size''': 240,
        '''dropout_rate''': 0.2,
        '''dw_padding''': [16],
    },
    '''b2''': {
        '''hidden_dim''': 1408,
        '''width_coef''': 1.1,
        '''depth_coef''': 1.2,
        '''image_size''': 260,
        '''dropout_rate''': 0.3,
        '''dw_padding''': [5, 8, 16],
    },
    '''b3''': {
        '''hidden_dim''': 1536,
        '''width_coef''': 1.2,
        '''depth_coef''': 1.4,
        '''image_size''': 300,
        '''dropout_rate''': 0.3,
        '''dw_padding''': [5, 18],
    },
    '''b4''': {
        '''hidden_dim''': 1792,
        '''width_coef''': 1.4,
        '''depth_coef''': 1.8,
        '''image_size''': 380,
        '''dropout_rate''': 0.4,
        '''dw_padding''': [6],
    },
    '''b5''': {
        '''hidden_dim''': 2048,
        '''width_coef''': 1.6,
        '''depth_coef''': 2.2,
        '''image_size''': 456,
        '''dropout_rate''': 0.4,
        '''dw_padding''': [13, 27],
    },
    '''b6''': {
        '''hidden_dim''': 2304,
        '''width_coef''': 1.8,
        '''depth_coef''': 2.6,
        '''image_size''': 528,
        '''dropout_rate''': 0.5,
        '''dw_padding''': [31],
    },
    '''b7''': {
        '''hidden_dim''': 2560,
        '''width_coef''': 2.0,
        '''depth_coef''': 3.1,
        '''image_size''': 600,
        '''dropout_rate''': 0.5,
        '''dw_padding''': [18],
    },
}



def lowerCamelCase						(     SCREAMING_SNAKE_CASE			):
    '''simple docstring'''







    __UpperCamelCase  :Dict    						=	EfficientNetConfig()
    __UpperCamelCase  :Any    						=	CONFIG_MAP[model_name]['''hidden_dim''']
    __UpperCamelCase  :List[Any]    						=	CONFIG_MAP[model_name]['''width_coef''']
    __UpperCamelCase  :int    						=	CONFIG_MAP[model_name]['''depth_coef''']
    __UpperCamelCase  :str    						=	CONFIG_MAP[model_name]['''image_size''']
    __UpperCamelCase  :Union[str, Any]    						=	CONFIG_MAP[model_name]['''dropout_rate''']
    __UpperCamelCase  :Any    						=	CONFIG_MAP[model_name]['''dw_padding''']

    __UpperCamelCase  :List[Any]    						=	'''huggingface/label-files'''
    __UpperCamelCase  :List[Any]    						=	'''imagenet-1k-id2label.json'''
    __UpperCamelCase  :Optional[Any]    						=	1_000
    __UpperCamelCase  :List[str]    						=	json.load(open(hf_hub_download(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    repo_type='''dataset'''			)						,    '''r'''			)			)
    __UpperCamelCase  :List[Any]    						=	{int(SCREAMING_SNAKE_CASE			): v for k, v in idalabel.items()}

    __UpperCamelCase  :Tuple    						=	idalabel
    __UpperCamelCase  :int    						=	{v: k for k, v in idalabel.items()}
    return config



def lowerCamelCase						(     ):
    '''simple docstring'''







    __UpperCamelCase  :str    						=	'''http://images.cocodataset.org/val2017/000000039769.jpg'''
    __UpperCamelCase  :List[str]    						=	Image.open(requests.get(SCREAMING_SNAKE_CASE						,    stream=SCREAMING_SNAKE_CASE			).raw			)
    return im



def lowerCamelCase						(     SCREAMING_SNAKE_CASE			):
    '''simple docstring'''







    __UpperCamelCase  :Union[str, Any]    						=	CONFIG_MAP[model_name]['''image_size''']
    __UpperCamelCase  :str    						=	EfficientNetImageProcessor(
        size={'''height''': size, '''width''': size}						,    image_mean=[0.485, 0.456, 0.406]						,    image_std=[0.47_853_944, 0.4_732_864, 0.47_434_163]						,    do_center_crop=SCREAMING_SNAKE_CASE						,    )
    return preprocessor



def lowerCamelCase						(     SCREAMING_SNAKE_CASE			):
    '''simple docstring'''







    __UpperCamelCase  :Dict    						=	[v.split('''_'''			)[0].split('''block'''			)[1] for v in original_param_names if v.startswith('''block'''			)]
    __UpperCamelCase  :List[str]    						=	sorted(set(SCREAMING_SNAKE_CASE			)			)
    __UpperCamelCase  :Any    						=	len(SCREAMING_SNAKE_CASE			)
    __UpperCamelCase  :List[str]    						=	{b: str(SCREAMING_SNAKE_CASE			) for b, i in zip(SCREAMING_SNAKE_CASE						,    range(SCREAMING_SNAKE_CASE			)			)}

    __UpperCamelCase  :Optional[int]    						=	[]
    rename_keys.append(('''stem_conv/kernel:0''', '''embeddings.convolution.weight''')			)
    rename_keys.append(('''stem_bn/gamma:0''', '''embeddings.batchnorm.weight''')			)
    rename_keys.append(('''stem_bn/beta:0''', '''embeddings.batchnorm.bias''')			)
    rename_keys.append(('''stem_bn/moving_mean:0''', '''embeddings.batchnorm.running_mean''')			)
    rename_keys.append(('''stem_bn/moving_variance:0''', '''embeddings.batchnorm.running_var''')			)

    for b in block_names:
        __UpperCamelCase  :Optional[Any]    						=	block_name_mapping[b]
        rename_keys.append((f"""block{b}_expand_conv/kernel:0""", f"""encoder.blocks.{hf_b}.expansion.expand_conv.weight""")			)
        rename_keys.append((f"""block{b}_expand_bn/gamma:0""", f"""encoder.blocks.{hf_b}.expansion.expand_bn.weight""")			)
        rename_keys.append((f"""block{b}_expand_bn/beta:0""", f"""encoder.blocks.{hf_b}.expansion.expand_bn.bias""")			)
        rename_keys.append(
            (f"""block{b}_expand_bn/moving_mean:0""", f"""encoder.blocks.{hf_b}.expansion.expand_bn.running_mean""")			)
        rename_keys.append(
            (f"""block{b}_expand_bn/moving_variance:0""", f"""encoder.blocks.{hf_b}.expansion.expand_bn.running_var""")			)
        rename_keys.append(
            (f"""block{b}_dwconv/depthwise_kernel:0""", f"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_conv.weight""")			)
        rename_keys.append((f"""block{b}_bn/gamma:0""", f"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.weight""")			)
        rename_keys.append((f"""block{b}_bn/beta:0""", f"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.bias""")			)
        rename_keys.append(
            (f"""block{b}_bn/moving_mean:0""", f"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_mean""")			)
        rename_keys.append(
            (f"""block{b}_bn/moving_variance:0""", f"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_var""")			)

        rename_keys.append((f"""block{b}_se_reduce/kernel:0""", f"""encoder.blocks.{hf_b}.squeeze_excite.reduce.weight""")			)
        rename_keys.append((f"""block{b}_se_reduce/bias:0""", f"""encoder.blocks.{hf_b}.squeeze_excite.reduce.bias""")			)
        rename_keys.append((f"""block{b}_se_expand/kernel:0""", f"""encoder.blocks.{hf_b}.squeeze_excite.expand.weight""")			)
        rename_keys.append((f"""block{b}_se_expand/bias:0""", f"""encoder.blocks.{hf_b}.squeeze_excite.expand.bias""")			)
        rename_keys.append(
            (f"""block{b}_project_conv/kernel:0""", f"""encoder.blocks.{hf_b}.projection.project_conv.weight""")			)
        rename_keys.append((f"""block{b}_project_bn/gamma:0""", f"""encoder.blocks.{hf_b}.projection.project_bn.weight""")			)
        rename_keys.append((f"""block{b}_project_bn/beta:0""", f"""encoder.blocks.{hf_b}.projection.project_bn.bias""")			)
        rename_keys.append(
            (f"""block{b}_project_bn/moving_mean:0""", f"""encoder.blocks.{hf_b}.projection.project_bn.running_mean""")			)
        rename_keys.append(
            (f"""block{b}_project_bn/moving_variance:0""", f"""encoder.blocks.{hf_b}.projection.project_bn.running_var""")			)

    rename_keys.append(('''top_conv/kernel:0''', '''encoder.top_conv.weight''')			)
    rename_keys.append(('''top_bn/gamma:0''', '''encoder.top_bn.weight''')			)
    rename_keys.append(('''top_bn/beta:0''', '''encoder.top_bn.bias''')			)
    rename_keys.append(('''top_bn/moving_mean:0''', '''encoder.top_bn.running_mean''')			)
    rename_keys.append(('''top_bn/moving_variance:0''', '''encoder.top_bn.running_var''')			)

    __UpperCamelCase  :List[Any]    						=	{}
    for item in rename_keys:
        if item[0] in original_param_names:
            __UpperCamelCase  :List[Any]    						=	'''efficientnet.''' + item[1]

    __UpperCamelCase  :Any    						=	'''classifier.weight'''
    __UpperCamelCase  :Optional[int]    						=	'''classifier.bias'''
    return key_mapping



def lowerCamelCase						(     SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE			):
    '''simple docstring'''







    for key, value in tf_params.items():
        if "normalization" in key:
            continue

        __UpperCamelCase  :Any    						=	key_mapping[key]
        if "_conv" in key and "kernel" in key:
            __UpperCamelCase  :Optional[int]    						=	torch.from_numpy(SCREAMING_SNAKE_CASE			).permute(3						,    2						,    0						,    1			)
        elif "depthwise_kernel" in key:
            __UpperCamelCase  :int    						=	torch.from_numpy(SCREAMING_SNAKE_CASE			).permute(2						,    3						,    0						,    1			)
        elif "kernel" in key:
            __UpperCamelCase  :Optional[Any]    						=	torch.from_numpy(np.transpose(SCREAMING_SNAKE_CASE			)			)
        else:
            __UpperCamelCase  :Optional[int]    						=	torch.from_numpy(SCREAMING_SNAKE_CASE			)

        # Replace HF parameters with original TF model parameters
        assert hf_params[hf_key].shape == new_hf_value.shape
        hf_params[hf_key].copy_(SCREAMING_SNAKE_CASE			)



@torch.no_grad()
def lowerCamelCase						(     SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE			):
    '''simple docstring'''







    __UpperCamelCase  :Optional[Any]    						=	model_classes[model_name](
        include_top=SCREAMING_SNAKE_CASE						,    weights='''imagenet'''						,    input_tensor=SCREAMING_SNAKE_CASE						,    input_shape=SCREAMING_SNAKE_CASE						,    pooling=SCREAMING_SNAKE_CASE						,    classes=1_000						,    classifier_activation='''softmax'''						,    )

    __UpperCamelCase  :Any    						=	original_model.trainable_variables
    __UpperCamelCase  :Any    						=	original_model.non_trainable_variables
    __UpperCamelCase  :Dict    						=	{param.name: param.numpy() for param in tf_params}
    for param in tf_non_train_params:
        __UpperCamelCase  :int    						=	param.numpy()
    __UpperCamelCase  :Optional[Any]    						=	list(tf_params.keys()			)

    # Load HuggingFace model
    __UpperCamelCase  :str    						=	get_efficientnet_config(SCREAMING_SNAKE_CASE			)
    __UpperCamelCase  :Dict    						=	EfficientNetForImageClassification(SCREAMING_SNAKE_CASE			).eval()
    __UpperCamelCase  :str    						=	hf_model.state_dict()

    # Create src-to-dst parameter name mapping dictionary
    print('''Converting parameters...'''			)
    __UpperCamelCase  :Union[str, Any]    						=	rename_keys(SCREAMING_SNAKE_CASE			)
    replace_params(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE			)

    # Initialize preprocessor and preprocess input image
    __UpperCamelCase  :Optional[int]    						=	convert_image_processor(SCREAMING_SNAKE_CASE			)
    __UpperCamelCase  :List[str]    						=	preprocessor(images=prepare_img()						,    return_tensors='''pt'''			)

    # HF model inference
    hf_model.eval()
    with torch.no_grad():
        __UpperCamelCase  :Dict    						=	hf_model(**SCREAMING_SNAKE_CASE			)
    __UpperCamelCase  :Union[str, Any]    						=	outputs.logits.detach().numpy()

    # Original model inference
    __UpperCamelCase  :List[str]    						=	False
    __UpperCamelCase  :List[str]    						=	CONFIG_MAP[model_name]['''image_size''']
    __UpperCamelCase  :List[Any]    						=	prepare_img().resize((image_size, image_size)						,    resample=PIL.Image.NEAREST			)
    __UpperCamelCase  :int    						=	image.img_to_array(SCREAMING_SNAKE_CASE			)
    __UpperCamelCase  :int    						=	np.expand_dims(SCREAMING_SNAKE_CASE						,    axis=0			)
    __UpperCamelCase  :Dict    						=	original_model.predict(SCREAMING_SNAKE_CASE			)

    # Check whether original and HF model outputs match  -> np.allclose
    assert np.allclose(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    atol=1e-3			), "The predicted logits are not the same."
    print('''Model outputs match!'''			)

    if save_model:
        # Create folder to save model
        if not os.path.isdir(SCREAMING_SNAKE_CASE			):
            os.mkdir(SCREAMING_SNAKE_CASE			)
        # Save converted model and image processor
        hf_model.save_pretrained(SCREAMING_SNAKE_CASE			)
        preprocessor.save_pretrained(SCREAMING_SNAKE_CASE			)

    if push_to_hub:
        # Push model and image processor to hub
        print(f"""Pushing converted {model_name} to the hub..."""			)
        __UpperCamelCase  :Optional[int]    						=	f"""efficientnet-{model_name}"""
        preprocessor.push_to_hub(SCREAMING_SNAKE_CASE			)
        hf_model.push_to_hub(SCREAMING_SNAKE_CASE			)


if __name__ == "__main__":
    __lowercase											=      argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        '''--model_name''',
        default='''b0''',
        type=str,
        help='''Version name of the EfficientNet model you want to convert, select from [b0, b1, b2, b3, b4, b5, b6, b7].''',
    )
    parser.add_argument(
        '''--pytorch_dump_folder_path''',
        default='''hf_model''',
        type=str,
        help='''Path to the output PyTorch model directory.''',
    )
    parser.add_argument('''--save_model''', action='''store_true''', help='''Save model to local''')
    parser.add_argument('''--push_to_hub''', action='''store_true''', help='''Push model and image processor to the hub''')

    __lowercase											=      parser.parse_args()
    convert_efficientnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)






import unittest

from transformers import MraConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_torch_available():
    import torch

    from transformers import (
        MraForMaskedLM,
        MraForMultipleChoice,
        MraForQuestionAnswering,
        MraForSequenceClassification,
        MraForTokenClassification,
        MraModel,
    )
    from transformers.models.mra.modeling_mra import MRA_PRETRAINED_MODEL_ARCHIVE_LIST


class 			lowerCamelCase_  :






    '''simple docstring'''






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







    def     UpperCamelCase__     (			self) ->			str:
        __UpperCamelCase  :int    						=	ids_tensor([self.batch_size, self.seq_length]				,				self.vocab_size)

        __UpperCamelCase  :List[Any]    						=	None
        if self.use_input_mask:
            __UpperCamelCase  :Any    						=	random_attention_mask([self.batch_size, self.seq_length])

        __UpperCamelCase  :Union[str, Any]    						=	None
        if self.use_token_type_ids:
            __UpperCamelCase  :Optional[int]    						=	ids_tensor([self.batch_size, self.seq_length]				,				self.type_vocab_size)

        __UpperCamelCase  :List[str]    						=	None
        __UpperCamelCase  :Tuple    						=	None
        __UpperCamelCase  :Union[str, Any]    						=	None
        if self.use_labels:
            __UpperCamelCase  :Dict    						=	ids_tensor([self.batch_size]				,				self.type_sequence_label_size)
            __UpperCamelCase  :Optional[int]    						=	ids_tensor([self.batch_size, self.seq_length]				,				self.num_labels)
            __UpperCamelCase  :List[str]    						=	ids_tensor([self.batch_size]				,				self.num_choices)

        __UpperCamelCase  :List[str]    						=	self.get_config()

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







    def     UpperCamelCase__     (			self) ->			Union[str, 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=__lowercase				,				initializer_range=self.initializer_range				,				)







    def     UpperCamelCase__     (			self) ->			List[Any]:
        __UpperCamelCase  :Dict    						=	self.get_config()
        __UpperCamelCase  :List[str]    						=	300
        return config







    def     UpperCamelCase__     (			self) ->			int:
        (
            (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, 
        )  :Tuple    						=	self.prepare_config_and_inputs()

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







    def     UpperCamelCase__     (			self				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				) ->			List[Any]:
        __UpperCamelCase  :Tuple    						=	True
        __UpperCamelCase  :Dict    						=	MraModel(__lowercase)
        model.to(__lowercase)
        model.eval()
        __UpperCamelCase  :int    						=	model(
            __lowercase				,				attention_mask=__lowercase				,				token_type_ids=__lowercase				,				encoder_hidden_states=__lowercase				,				encoder_attention_mask=__lowercase				,				)
        __UpperCamelCase  :Any    						=	model(
            __lowercase				,				attention_mask=__lowercase				,				token_type_ids=__lowercase				,				encoder_hidden_states=__lowercase				,				)
        __UpperCamelCase  :str    						=	model(__lowercase				,				attention_mask=__lowercase				,				token_type_ids=__lowercase)
        self.parent.assertEqual(result.last_hidden_state.shape				,				(self.batch_size, self.seq_length, self.hidden_size))







    def     UpperCamelCase__     (			self				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase) ->			Union[str, Any]:
        __UpperCamelCase  :Dict    						=	MraForMaskedLM(config=__lowercase)
        model.to(__lowercase)
        model.eval()
        __UpperCamelCase  :Optional[Any]    						=	model(__lowercase				,				attention_mask=__lowercase				,				token_type_ids=__lowercase				,				labels=__lowercase)
        self.parent.assertEqual(result.logits.shape				,				(self.batch_size, self.seq_length, self.vocab_size))







    def     UpperCamelCase__     (			self				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase) ->			int:
        __UpperCamelCase  :Tuple    						=	MraForQuestionAnswering(config=__lowercase)
        model.to(__lowercase)
        model.eval()
        __UpperCamelCase  :Any    						=	model(
            __lowercase				,				attention_mask=__lowercase				,				token_type_ids=__lowercase				,				start_positions=__lowercase				,				end_positions=__lowercase				,				)
        self.parent.assertEqual(result.start_logits.shape				,				(self.batch_size, self.seq_length))
        self.parent.assertEqual(result.end_logits.shape				,				(self.batch_size, self.seq_length))







    def     UpperCamelCase__     (			self				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase) ->			List[str]:
        __UpperCamelCase  :List[Any]    						=	self.num_labels
        __UpperCamelCase  :str    						=	MraForSequenceClassification(__lowercase)
        model.to(__lowercase)
        model.eval()
        __UpperCamelCase  :Dict    						=	model(__lowercase				,				attention_mask=__lowercase				,				token_type_ids=__lowercase				,				labels=__lowercase)
        self.parent.assertEqual(result.logits.shape				,				(self.batch_size, self.num_labels))







    def     UpperCamelCase__     (			self				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase) ->			Optional[int]:
        __UpperCamelCase  :str    						=	self.num_labels
        __UpperCamelCase  :Optional[int]    						=	MraForTokenClassification(config=__lowercase)
        model.to(__lowercase)
        model.eval()
        __UpperCamelCase  :Optional[Any]    						=	model(__lowercase				,				attention_mask=__lowercase				,				token_type_ids=__lowercase				,				labels=__lowercase)
        self.parent.assertEqual(result.logits.shape				,				(self.batch_size, self.seq_length, self.num_labels))







    def     UpperCamelCase__     (			self				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase) ->			Optional[Any]:
        __UpperCamelCase  :Optional[Any]    						=	self.num_choices
        __UpperCamelCase  :Dict    						=	MraForMultipleChoice(config=__lowercase)
        model.to(__lowercase)
        model.eval()
        __UpperCamelCase  :Tuple    						=	input_ids.unsqueeze(1).expand(-1				,				self.num_choices				,				-1).contiguous()
        __UpperCamelCase  :int    						=	token_type_ids.unsqueeze(1).expand(-1				,				self.num_choices				,				-1).contiguous()
        __UpperCamelCase  :Optional[int]    						=	input_mask.unsqueeze(1).expand(-1				,				self.num_choices				,				-1).contiguous()
        __UpperCamelCase  :Union[str, Any]    						=	model(
            __lowercase				,				attention_mask=__lowercase				,				token_type_ids=__lowercase				,				labels=__lowercase				,				)
        self.parent.assertEqual(result.logits.shape				,				(self.batch_size, self.num_choices))







    def     UpperCamelCase__     (			self) ->			Optional[int]:
        __UpperCamelCase  :str    						=	self.prepare_config_and_inputs()
        (
            (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, 
        )  :str    						=	config_and_inputs
        __UpperCamelCase  :List[str]    						=	{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
        return config, inputs_dict


@require_torch
class 			lowerCamelCase_  (     UpperCAmelCase_		,			unittest.TestCase				):






    '''simple docstring'''






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

    a__			:		Union[str, Any]      			=   ()







    def     UpperCamelCase__     (			self) ->			List[str]:
        __UpperCamelCase  :str    						=	MraModelTester(self)
        __UpperCamelCase  :Tuple    						=	ConfigTester(self				,				config_class=__lowercase				,				hidden_size=37)







    def     UpperCamelCase__     (			self) ->			Dict:
        self.config_tester.run_common_tests()







    def     UpperCamelCase__     (			self) ->			Union[str, Any]:
        __UpperCamelCase  :List[str]    						=	self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__lowercase)







    def     UpperCamelCase__     (			self) ->			List[Any]:
        __UpperCamelCase  :Union[str, Any]    						=	self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
            __UpperCamelCase  :Optional[int]    						=	type
            self.model_tester.create_and_check_model(*__lowercase)







    def     UpperCamelCase__     (			self) ->			List[Any]:
        __UpperCamelCase  :Union[str, Any]    						=	self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*__lowercase)







    def     UpperCamelCase__     (			self) ->			Any:
        __UpperCamelCase  :Dict    						=	self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_multiple_choice(*__lowercase)







    def     UpperCamelCase__     (			self) ->			Dict:
        __UpperCamelCase  :Union[str, Any]    						=	self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_question_answering(*__lowercase)







    def     UpperCamelCase__     (			self) ->			Tuple:
        __UpperCamelCase  :Optional[int]    						=	self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_sequence_classification(*__lowercase)







    def     UpperCamelCase__     (			self) ->			Any:
        __UpperCamelCase  :Dict    						=	self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*__lowercase)







    @slow
    def     UpperCamelCase__     (			self) ->			Dict:
        for model_name in MRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            __UpperCamelCase  :List[Any]    						=	MraModel.from_pretrained(__lowercase)
            self.assertIsNotNone(__lowercase)







    @unittest.skip(reason='''MRA does not output attentions''')
    def     UpperCamelCase__     (			self) ->			Any:
        return


@require_torch
class 			lowerCamelCase_  (     unittest.TestCase				):






    '''simple docstring'''






    @slow
    def     UpperCamelCase__     (			self) ->			Optional[int]:
        __UpperCamelCase  :int    						=	MraModel.from_pretrained('''uw-madison/mra-base-512-4''')
        __UpperCamelCase  :Union[str, Any]    						=	torch.arange(256).unsqueeze(0)

        with torch.no_grad():
            __UpperCamelCase  :List[Any]    						=	model(__lowercase)[0]

        __UpperCamelCase  :Dict    						=	torch.Size((1, 256, 768))
        self.assertEqual(output.shape				,				__lowercase)

        __UpperCamelCase  :Optional[Any]    						=	torch.tensor(
            [[[-0.01_40, 0.08_30, -0.03_81], [0.15_46, 0.14_02, 0.02_20], [0.11_62, 0.08_51, 0.01_65]]])

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







    @slow
    def     UpperCamelCase__     (			self) ->			Any:
        __UpperCamelCase  :str    						=	MraForMaskedLM.from_pretrained('''uw-madison/mra-base-512-4''')
        __UpperCamelCase  :str    						=	torch.arange(256).unsqueeze(0)

        with torch.no_grad():
            __UpperCamelCase  :Optional[Any]    						=	model(__lowercase)[0]

        __UpperCamelCase  :List[str]    						=	50_265

        __UpperCamelCase  :List[Any]    						=	torch.Size((1, 256, vocab_size))
        self.assertEqual(output.shape				,				__lowercase)

        __UpperCamelCase  :Optional[Any]    						=	torch.tensor(
            [[[9.25_95, -3.60_38, 11.88_19], [9.38_69, -3.26_93, 11.09_56], [11.85_24, -3.49_38, 13.12_10]]])

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







    @slow
    def     UpperCamelCase__     (			self) ->			Optional[int]:
        __UpperCamelCase  :Dict    						=	MraForMaskedLM.from_pretrained('''uw-madison/mra-base-4096-8-d3''')
        __UpperCamelCase  :Optional[int]    						=	torch.arange(4_096).unsqueeze(0)

        with torch.no_grad():
            __UpperCamelCase  :Tuple    						=	model(__lowercase)[0]

        __UpperCamelCase  :Optional[int]    						=	50_265

        __UpperCamelCase  :Optional[Any]    						=	torch.Size((1, 4_096, vocab_size))
        self.assertEqual(output.shape				,				__lowercase)

        __UpperCamelCase  :List[str]    						=	torch.tensor(
            [[[5.47_89, -2.35_64, 7.50_64], [7.90_67, -1.33_69, 9.96_68], [9.07_12, -1.81_06, 7.03_80]]])

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





