

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						,					A_						,					A_=2						,					A_=8						,					A_=True						,					A_=True						,					A_=True						,					A_=True						,					A_=99						,					A_=16						,					A_=5						,					A_=2						,					A_=36						,					A_="gelu"						,					A_=0.0						,					A_=0.0						,					A_=512						,					A_=16						,					A_=2						,					A_=0.02						,					A_=3						,					A_=4						,					A_=None						,					):







										'''simple docstring'''


										_UpperCAmelCase       :		Optional[int] 		=    parent
										_UpperCAmelCase       :		Optional[int] 		=    batch_size
										_UpperCAmelCase       :		Dict 		=    seq_length
										_UpperCAmelCase       :		int 		=    is_training
										_UpperCAmelCase       :		int 		=    use_input_mask
										_UpperCAmelCase       :		Optional[int] 		=    use_token_type_ids
										_UpperCAmelCase       :		List[Any] 		=    use_labels
										_UpperCAmelCase       :		Union[str, Any] 		=    vocab_size
										_UpperCAmelCase       :		Optional[int] 		=    hidden_size
										_UpperCAmelCase       :		List[Any] 		=    num_hidden_layers
										_UpperCAmelCase       :		int 		=    num_attention_heads
										_UpperCAmelCase       :		List[str] 		=    intermediate_size
										_UpperCAmelCase       :		Dict 		=    hidden_act
										_UpperCAmelCase       :		str 		=    hidden_dropout_prob
										_UpperCAmelCase       :		Optional[Any] 		=    attention_probs_dropout_prob
										_UpperCAmelCase       :		List[str] 		=    max_position_embeddings
										_UpperCAmelCase       :		Dict 		=    type_vocab_size
										_UpperCAmelCase       :		Tuple 		=    type_sequence_label_size
										_UpperCAmelCase       :		str 		=    initializer_range
										_UpperCAmelCase       :		Union[str, Any] 		=    num_labels
										_UpperCAmelCase       :		Dict 		=    num_choices
										_UpperCAmelCase       :		Dict 		=    scope


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Dict 		=    ids_tensor([self.batch_size, self.seq_length]						,					self.vocab_size	)

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

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

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

										_UpperCAmelCase       :		Any 		=    self.get_config()

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


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										return MraConfig(
										    vocab_size=self.vocab_size						,					hidden_size=self.hidden_size						,					num_hidden_layers=self.num_hidden_layers						,					num_attention_heads=self.num_attention_heads						,					intermediate_size=self.intermediate_size						,					hidden_act=self.hidden_act						,					hidden_dropout_prob=self.hidden_dropout_prob						,					attention_probs_dropout_prob=self.attention_probs_dropout_prob						,					max_position_embeddings=self.max_position_embeddings						,					type_vocab_size=self.type_vocab_size						,					is_decoder=A_						,					initializer_range=self.initializer_range						,					)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Dict 		=    self.get_config()
										_UpperCAmelCase       :		Union[str, Any] 		=    300
										return config


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										(
										    _UpperCAmelCase							
										)       :		Dict 		=    self.prepare_config_and_inputs()

										_UpperCAmelCase       :		Optional[Any] 		=    True
										_UpperCAmelCase       :		List[str] 		=    floats_tensor([self.batch_size, self.seq_length, self.hidden_size]	)
										_UpperCAmelCase       :		Optional[int] 		=    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						,					A_						,					A_						,					A_						,					A_						,					A_						,					A_						,					A_	):







										'''simple docstring'''


										_UpperCAmelCase       :		Any 		=    MraModel(config=A_	)
										model.to(A_	)
										model.eval()
										_UpperCAmelCase       :		int 		=    model(A_						,					attention_mask=A_						,					token_type_ids=A_	)
										_UpperCAmelCase       :		List[str] 		=    model(A_						,					token_type_ids=A_	)
										_UpperCAmelCase       :		Tuple 		=    model(A_	)
										self.parent.assertEqual(result.last_hidden_state.shape						,					(self.batch_size, self.seq_length, self.hidden_size)	)


						def        _UpperCAmelCase				(       self						,					A_						,					A_						,					A_						,					A_						,					A_						,					A_						,					A_						,					A_						,					A_						,					):







										'''simple docstring'''


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


						def        _UpperCAmelCase				(       self						,					A_						,					A_						,					A_						,					A_						,					A_						,					A_						,					A_	):







										'''simple docstring'''


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


						def        _UpperCAmelCase				(       self						,					A_						,					A_						,					A_						,					A_						,					A_						,					A_						,					A_	):







										'''simple docstring'''


										_UpperCAmelCase       :		List[str] 		=    MraForQuestionAnswering(config=A_	)
										model.to(A_	)
										model.eval()
										_UpperCAmelCase       :		Union[str, Any] 		=    model(
										    A_						,					attention_mask=A_						,					token_type_ids=A_						,					start_positions=A_						,					end_positions=A_						,					)
										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						,					A_						,					A_						,					A_						,					A_						,					A_						,					A_						,					A_	):







										'''simple docstring'''


										_UpperCAmelCase       :		List[str] 		=    self.num_labels
										_UpperCAmelCase       :		Optional[Any] 		=    MraForSequenceClassification(A_	)
										model.to(A_	)
										model.eval()
										_UpperCAmelCase       :		str 		=    model(A_						,					attention_mask=A_						,					token_type_ids=A_						,					labels=A_	)
										self.parent.assertEqual(result.logits.shape						,					(self.batch_size, self.num_labels)	)


						def        _UpperCAmelCase				(       self						,					A_						,					A_						,					A_						,					A_						,					A_						,					A_						,					A_	):







										'''simple docstring'''


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


						def        _UpperCAmelCase				(       self						,					A_						,					A_						,					A_						,					A_						,					A_						,					A_						,					A_	):







										'''simple docstring'''


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





						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		int 		=    self.prepare_config_and_inputs()
										(
										    _UpperCAmelCase							
										)       :		int 		=    config_and_inputs
										_UpperCAmelCase       :		List[Any] 		=    {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
										return config, inputs_dict







@require_torch
class 				a  (	UpperCAmelCase							, unittest.TestCase				):
						_lowercase    						=			(
						    (
						        MraModel,
						        MraForMaskedLM,
						        MraForMultipleChoice,
						        MraForQuestionAnswering,
						        MraForSequenceClassification,
						        MraForTokenClassification,
						    )
						    if is_torch_available()
						    else ()
						)
						_lowercase    						=			False
						_lowercase    						=			False
						_lowercase    						=			False
						_lowercase    						=			False

						_lowercase    						=			()


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Tuple 		=    MraModelTester(self	)
										_UpperCAmelCase       :		Dict 		=    ConfigTester(self						,					config_class=A_						,					hidden_size=37	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										self.config_tester.run_common_tests()


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Union[str, Any] 		=    self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*A_	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		List[Any] 		=    self.model_tester.prepare_config_and_inputs()
										for type in ["absolute", "relative_key", "relative_key_query"]:
														_UpperCAmelCase       :		Union[str, Any] 		=    type
														self.model_tester.create_and_check_model(*A_	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Union[str, Any] 		=    self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_masked_lm(*A_	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Optional[int] 		=    self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_multiple_choice(*A_	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Optional[int] 		=    self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_question_answering(*A_	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Any 		=    self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_sequence_classification(*A_	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		str 		=    self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_token_classification(*A_	)


						@slow
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										for model_name in MRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
														_UpperCAmelCase       :		Dict 		=    MraModel.from_pretrained(A_	)
														self.assertIsNotNone(A_	)





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







										'''simple docstring'''


										return







@require_torch
class 				a  (	unittest.TestCase				):


						@slow
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		List[str] 		=    MraModel.from_pretrained("uw-madison/mra-base-512-4"	)
										_UpperCAmelCase       :		Any 		=    torch.arange(256	).unsqueeze(0	)

										with torch.no_grad():
														_UpperCAmelCase       :		Optional[Any] 		=    model(A_	)[0]

										_UpperCAmelCase       :		Union[str, Any] 		=    torch.Size((1, 256, 768)	)
										self.assertEqual(output.shape						,					A_	)

										_UpperCAmelCase       :		Union[str, 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]						,					A_						,					atol=1e-4	)	)


						@slow
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Optional[Any] 		=    MraForMaskedLM.from_pretrained("uw-madison/mra-base-512-4"	)
										_UpperCAmelCase       :		Tuple 		=    torch.arange(256	).unsqueeze(0	)

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

										_UpperCAmelCase       :		Optional[Any] 		=    50265

										_UpperCAmelCase       :		Tuple 		=    torch.Size((1, 256, vocab_size)	)
										self.assertEqual(output.shape						,					A_	)

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





						@slow
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Optional[Any] 		=    MraForMaskedLM.from_pretrained("uw-madison/mra-base-4096-8-d3"	)
										_UpperCAmelCase       :		List[Any] 		=    torch.arange(4096	).unsqueeze(0	)

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

										_UpperCAmelCase       :		List[Any] 		=    50265

										_UpperCAmelCase       :		Any 		=    torch.Size((1, 4096, vocab_size)	)
										self.assertEqual(output.shape						,					A_	)

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





from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import rescale, resize, to_channel_dimension_format
from ...image_utils import (
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
					import PIL


SCREAMING_SNAKE_CASE_        =      logging.get_logger(__name__)

def 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  Dict  ,      lowerCAmelCase:  Union[str, Any] )   ->       Union[str, Any]:
				_UpperCAmelCase       :		Dict 		=    b.T
				_UpperCAmelCase       :		Dict 		=    np.sum(np.square(lowerCAmelCase )  ,      axis=1 )
				_UpperCAmelCase       :		Optional[Any] 		=    np.sum(np.square(lowerCAmelCase )  ,      axis=0 )
				_UpperCAmelCase       :		str 		=    np.matmul(lowerCAmelCase  ,      lowerCAmelCase )
				_UpperCAmelCase       :		Any 		=    aa[:, None] - 2 * ab + ba[None, :]
				return d

def 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  Union[str, Any]  ,      lowerCAmelCase:  Dict )   ->       int:
				_UpperCAmelCase       :		Any 		=    x.reshape(-1  ,      3 )
				_UpperCAmelCase       :		List[str] 		=    squared_euclidean_distance(lowerCAmelCase  ,      lowerCAmelCase )
				return np.argmin(lowerCAmelCase  ,      axis=1 )





class 				a  (	UpperCAmelCase				):
						_lowercase    						=			["pixel_values"]


						def __init__(       self						,					A_ = None						,					A_ = True						,					A_ = None						,					A_ = PILImageResampling.BILINEAR						,					A_ = True						,					A_ = True						,					**A_						,					):







										'''simple docstring'''


										super().__init__(**A_	)
										_UpperCAmelCase       :		Optional[Any] 		=    size if size is not None else {"height": 256, "width": 256}
										_UpperCAmelCase       :		Optional[int] 		=    get_size_dict(A_	)
										_UpperCAmelCase       :		Union[str, Any] 		=    np.array(A_	) if clusters is not None else None
										_UpperCAmelCase       :		int 		=    do_resize
										_UpperCAmelCase       :		Union[str, Any] 		=    size
										_UpperCAmelCase       :		Optional[Any] 		=    resample
										_UpperCAmelCase       :		str 		=    do_normalize
										_UpperCAmelCase       :		List[str] 		=    do_color_quantize


						def        _UpperCAmelCase				(       self						,					A_						,					A_						,					A_ = PILImageResampling.BILINEAR						,					A_ = None						,					**A_						,					):







										'''simple docstring'''


										_UpperCAmelCase       :		int 		=    get_size_dict(A_	)
										if "height" not in size or "width" not in size:
														raise ValueError(f'Size dictionary must contain both height and width keys. Got {size.keys()}'	)
										return resize(
										    A_						,					size=(size["height"], size["width"])						,					resample=A_						,					data_format=A_						,					**A_	)


						def        _UpperCAmelCase				(       self						,					A_						,					A_ = None						,					):







										'''simple docstring'''


										_UpperCAmelCase       :		Dict 		=    rescale(image=A_						,					scale=1 / 1_27.5						,					data_format=A_	)
										_UpperCAmelCase       :		List[Any] 		=    image - 1
										return image





						def        _UpperCAmelCase				(       self						,					A_						,					A_ = None						,					A_ = None						,					A_ = None						,					A_ = None						,					A_ = None						,					A_ = None						,					A_ = None						,					A_ = ChannelDimension.FIRST						,					**A_						,					):







										'''simple docstring'''


										_UpperCAmelCase       :		Optional[int] 		=    do_resize if do_resize is not None else self.do_resize
										_UpperCAmelCase       :		Any 		=    size if size is not None else self.size
										_UpperCAmelCase       :		Dict 		=    get_size_dict(A_	)
										_UpperCAmelCase       :		List[Any] 		=    resample if resample is not None else self.resample
										_UpperCAmelCase       :		Any 		=    do_normalize if do_normalize is not None else self.do_normalize
										_UpperCAmelCase       :		Optional[Any] 		=    do_color_quantize if do_color_quantize is not None else self.do_color_quantize
										_UpperCAmelCase       :		Any 		=    clusters if clusters is not None else self.clusters
										_UpperCAmelCase       :		Optional[int] 		=    np.array(A_	)

										_UpperCAmelCase       :		List[str] 		=    make_list_of_images(A_	)

										if not valid_images(A_	):
														raise ValueError(
														    "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
														    "torch.Tensor, tf.Tensor or jax.ndarray."	)

										if do_resize and size is None or resample is None:
														raise ValueError("Size and resample must be specified if do_resize is True."	)

										if do_color_quantize and clusters is None:
														raise ValueError("Clusters must be specified if do_color_quantize is True."	)

										# All transformations expect numpy arrays.
										_UpperCAmelCase       :		List[str] 		=    [to_numpy_array(A_	) for image in images]

										if do_resize:
														_UpperCAmelCase       :		int 		=    [self.resize(image=A_						,					size=A_						,					resample=A_	) for image in images]

										if do_normalize:
														_UpperCAmelCase       :		List[str] 		=    [self.normalize(image=A_	) for image in images]

										if do_color_quantize:
														_UpperCAmelCase       :		Tuple 		=    [to_channel_dimension_format(A_						,					ChannelDimension.LAST	) for image in images]
														# color quantize from (batch_size, height, width, 3) to (batch_size, height, width)
														_UpperCAmelCase       :		List[str] 		=    np.array(A_	)
														_UpperCAmelCase       :		List[Any] 		=    color_quantize(A_						,					A_	).reshape(images.shape[:-1]	)

														# flatten to (batch_size, height*width)
														_UpperCAmelCase       :		Any 		=    images.shape[0]
														_UpperCAmelCase       :		List[Any] 		=    images.reshape(A_						,					-1	)

														# We need to convert back to a list of images to keep consistent behaviour across processors.
														_UpperCAmelCase       :		Union[str, Any] 		=    list(A_	)
										else:
														_UpperCAmelCase       :		Optional[Any] 		=    [to_channel_dimension_format(A_						,					A_	) for image in images]

										_UpperCAmelCase       :		List[Any] 		=    {"input_ids": images}
										return BatchFeature(data=A_						,					tensor_type=A_	)


