



"""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       :







		'''simple docstring'''





		def __init__(	self						:  Any			,_snake_case						:  str			,_snake_case						:  List[Any]=2			,_snake_case						:  List[Any]=8			,_snake_case						:  List[Any]=True			,_snake_case						:  Optional[int]=True			,_snake_case						:  str=True			,_snake_case						:  int=True			,_snake_case						:  str=99			,_snake_case						:  Tuple=16			,_snake_case						:  List[str]=5			,_snake_case						:  Any=2			,_snake_case						:  Dict=36			,_snake_case						:  Optional[int]="gelu"			,_snake_case						:  Any=0.0			,_snake_case						:  int=0.0			,_snake_case						:  Union[str, Any]=512			,_snake_case						:  Optional[Any]=16			,_snake_case						:  Optional[Any]=2			,_snake_case						:  Optional[Any]=0.02			,_snake_case						:  Tuple=3			,_snake_case						:  List[str]=4			,_snake_case						:  Tuple=None			,)					->		int:






				"""simple docstring"""



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







		def 					UpperCAmelCase    (	self						:  str						)					->		str:






				"""simple docstring"""



				lowercase__   :    Tuple				=			ids_tensor([self.batch_size, self.seq_length]			,self.vocab_size						)

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

				lowercase__   :    Optional[int]				=			None
				if self.use_token_type_ids:
						lowercase__   :    Dict				=			ids_tensor([self.batch_size, self.seq_length]			,self.type_vocab_size						)

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

				lowercase__   :    Union[str, Any]				=			self.get_config()

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







		def 					UpperCAmelCase    (	self						:  Union[str, Any]						)					->		Tuple:






				"""simple docstring"""



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







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






				"""simple docstring"""



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







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






				"""simple docstring"""



				(
				    (
				    lowercase__
				)		,			(
				    lowercase__
				)		,			(
				    lowercase__
				)		,			(
				    lowercase__
				)		,			(
				    lowercase__
				)		,			(
				    lowercase__
				)		,			(
				    lowercase__
				)		,			
				)   :    str				=			self.prepare_config_and_inputs()

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






				"""simple docstring"""



				lowercase__   :    List[Any]				=			MraModel(config=_snake_case						)
				model.to(_snake_case						)
				model.eval()
				lowercase__   :    List[Any]				=			model(_snake_case			,attention_mask=_snake_case			,token_type_ids=_snake_case						)
				lowercase__   :    Optional[Any]				=			model(_snake_case			,token_type_ids=_snake_case						)
				lowercase__   :    Union[str, Any]				=			model(_snake_case						)
				self.parent.assertEqual(result.last_hidden_state.shape			,(self.batch_size, self.seq_length, self.hidden_size)						)







		def 					UpperCAmelCase    (	self						:  int			,_snake_case						:  Optional[Any]			,_snake_case						:  Optional[Any]			,_snake_case						:  List[str]			,_snake_case						:  Any			,_snake_case						:  Tuple			,_snake_case						:  int			,_snake_case						:  Optional[Any]			,_snake_case						:  Dict			,_snake_case						:  Tuple			,)					->		Tuple:






				"""simple docstring"""



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







		def 					UpperCAmelCase    (	self						:  Union[str, Any]			,_snake_case						:  List[str]			,_snake_case						:  str			,_snake_case						:  int			,_snake_case						:  Optional[int]			,_snake_case						:  Tuple			,_snake_case						:  Tuple			,_snake_case						:  str						)					->		Optional[Any]:






				"""simple docstring"""



				lowercase__   :    int				=			MraForMaskedLM(config=_snake_case						)
				model.to(_snake_case						)
				model.eval()
				lowercase__   :    Optional[Any]				=			model(_snake_case			,attention_mask=_snake_case			,token_type_ids=_snake_case			,labels=_snake_case						)
				self.parent.assertEqual(result.logits.shape			,(self.batch_size, self.seq_length, self.vocab_size)						)







		def 					UpperCAmelCase    (	self						:  Optional[Any]			,_snake_case						:  Dict			,_snake_case						:  Any			,_snake_case						:  List[Any]			,_snake_case						:  int			,_snake_case						:  Union[str, Any]			,_snake_case						:  Union[str, Any]			,_snake_case						:  List[str]						)					->		Any:






				"""simple docstring"""



				lowercase__   :    List[str]				=			MraForQuestionAnswering(config=_snake_case						)
				model.to(_snake_case						)
				model.eval()
				lowercase__   :    Union[str, Any]				=			model(
				    _snake_case			,attention_mask=_snake_case			,token_type_ids=_snake_case			,start_positions=_snake_case			,end_positions=_snake_case			,)
				self.parent.assertEqual(result.start_logits.shape			,(self.batch_size, self.seq_length)						)
				self.parent.assertEqual(result.end_logits.shape			,(self.batch_size, self.seq_length)						)







		def 					UpperCAmelCase    (	self						:  str			,_snake_case						:  List[Any]			,_snake_case						:  Dict			,_snake_case						:  str			,_snake_case						:  int			,_snake_case						:  List[str]			,_snake_case						:  Tuple			,_snake_case						:  Optional[Any]						)					->		Union[str, Any]:






				"""simple docstring"""



				lowercase__   :    Tuple				=			self.num_labels
				lowercase__   :    Tuple				=			MraForSequenceClassification(_snake_case						)
				model.to(_snake_case						)
				model.eval()
				lowercase__   :    Optional[Any]				=			model(_snake_case			,attention_mask=_snake_case			,token_type_ids=_snake_case			,labels=_snake_case						)
				self.parent.assertEqual(result.logits.shape			,(self.batch_size, self.num_labels)						)







		def 					UpperCAmelCase    (	self						:  Optional[int]			,_snake_case						:  Tuple			,_snake_case						:  Dict			,_snake_case						:  Any			,_snake_case						:  Optional[Any]			,_snake_case						:  Union[str, Any]			,_snake_case						:  List[Any]			,_snake_case						:  List[Any]						)					->		Union[str, Any]:






				"""simple docstring"""



				lowercase__   :    List[Any]				=			self.num_labels
				lowercase__   :    int				=			MraForTokenClassification(config=_snake_case						)
				model.to(_snake_case						)
				model.eval()
				lowercase__   :    int				=			model(_snake_case			,attention_mask=_snake_case			,token_type_ids=_snake_case			,labels=_snake_case						)
				self.parent.assertEqual(result.logits.shape			,(self.batch_size, self.seq_length, self.num_labels)						)







		def 					UpperCAmelCase    (	self						:  Optional[int]			,_snake_case						:  Union[str, Any]			,_snake_case						:  Optional[int]			,_snake_case						:  Union[str, Any]			,_snake_case						:  Any			,_snake_case						:  str			,_snake_case						:  Optional[int]			,_snake_case						:  Optional[int]						)					->		Dict:






				"""simple docstring"""



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







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






				"""simple docstring"""



				lowercase__   :    Tuple				=			self.prepare_config_and_inputs()
				(
				    (
				    lowercase__
				)		,			(
				    lowercase__
				)		,			(
				    lowercase__
				)		,			(
				    lowercase__
				)		,			(
				    lowercase__
				)		,			(
				    lowercase__
				)		,			(
				    lowercase__
				)		,			
				)   :    Dict				=			config_and_inputs
				lowercase__   :    str				=			{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
				return config, inputs_dict





@require_torch
class        __A       (	A_     ,unittest.TestCase    ):







		'''simple docstring'''





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

		lowerCAmelCase			:					Dict   	=				()







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






				"""simple docstring"""



				lowercase__   :    Dict				=			MraModelTester(self						)
				lowercase__   :    Optional[int]				=			ConfigTester(self			,config_class=_snake_case			,hidden_size=37						)







		def 					UpperCAmelCase    (	self						:  Optional[Any]						)					->		Union[str, Any]:






				"""simple docstring"""



				self.config_tester.run_common_tests()







		def 					UpperCAmelCase    (	self						:  int						)					->		Tuple:






				"""simple docstring"""



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







		def 					UpperCAmelCase    (	self						:  Any						)					->		Dict:






				"""simple docstring"""



				lowercase__   :    int				=			self.model_tester.prepare_config_and_inputs()
				for type in ["absolute", "relative_key", "relative_key_query"]:
						lowercase__   :    Any				=			type
						self.model_tester.create_and_check_model(*_snake_case						)







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






				"""simple docstring"""



				lowercase__   :    Dict				=			self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_masked_lm(*_snake_case						)







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






				"""simple docstring"""



				lowercase__   :    Optional[Any]				=			self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_multiple_choice(*_snake_case						)







		def 					UpperCAmelCase    (	self						:  int						)					->		Union[str, Any]:






				"""simple docstring"""



				lowercase__   :    Optional[int]				=			self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_question_answering(*_snake_case						)







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






				"""simple docstring"""



				lowercase__   :    Tuple				=			self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_sequence_classification(*_snake_case						)







		def 					UpperCAmelCase    (	self						:  str						)					->		str:






				"""simple docstring"""



				lowercase__   :    List[Any]				=			self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_token_classification(*_snake_case						)







		@slow
		def 					UpperCAmelCase    (	self						:  Optional[int]						)					->		Union[str, Any]:






				"""simple docstring"""



				for model_name in MRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
						lowercase__   :    str				=			MraModel.from_pretrained(_snake_case						)
						self.assertIsNotNone(_snake_case						)







		@unittest.skip(reason='''MRA does not output attentions'''						)
		def 					UpperCAmelCase    (	self						:  Optional[int]						)					->		str:






				"""simple docstring"""



				return





@require_torch
class        __A       (	unittest.TestCase    ):







		'''simple docstring'''





		@slow
		def 					UpperCAmelCase    (	self						:  Optional[Any]						)					->		int:






				"""simple docstring"""



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

				with torch.no_grad():
						lowercase__   :    int				=			model(_snake_case						)[0]

				lowercase__   :    Optional[Any]				=			torch.Size((1, 256, 768)						)
				self.assertEqual(output.shape			,_snake_case						)

				lowercase__   :    str				=			torch.tensor(
				    [[[-0.0140, 0.0830, -0.0381], [0.1546, 0.1402, 0.0220], [0.1162, 0.0851, 0.0165]]]						)

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







		@slow
		def 					UpperCAmelCase    (	self						:  str						)					->		List[str]:






				"""simple docstring"""



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

				with torch.no_grad():
						lowercase__   :    int				=			model(_snake_case						)[0]

				lowercase__   :    Dict				=			50_265

				lowercase__   :    List[str]				=			torch.Size((1, 256, vocab_size)						)
				self.assertEqual(output.shape			,_snake_case						)

				lowercase__   :    Union[str, Any]				=			torch.tensor(
				    [[[9.2595, -3.6038, 11.8819], [9.3869, -3.2693, 11.0956], [11.8524, -3.4938, 13.1210]]]						)

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







		@slow
		def 					UpperCAmelCase    (	self						:  Dict						)					->		Tuple:






				"""simple docstring"""



				lowercase__   :    int				=			MraForMaskedLM.from_pretrained('''uw-madison/mra-base-4096-8-d3'''						)
				lowercase__   :    List[str]				=			torch.arange(4_096						).unsqueeze(0						)

				with torch.no_grad():
						lowercase__   :    List[str]				=			model(_snake_case						)[0]

				lowercase__   :    Any				=			50_265

				lowercase__   :    List[str]				=			torch.Size((1, 4_096, vocab_size)						)
				self.assertEqual(output.shape			,_snake_case						)

				lowercase__   :    Union[str, Any]				=			torch.tensor(
				    [[[5.4789, -2.3564, 7.5064], [7.9067, -1.3369, 9.9668], [9.0712, -1.8106, 7.0380]]]						)

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





"""simple docstring"""







from collections import UserDict
from typing import List, Union

from ..utils import (
    add_end_docstrings,
    is_tf_available,
    is_torch_available,
    is_vision_available,
    logging,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
							from PIL import Image

							from ..image_utils import load_image

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

if is_tf_available():
							from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING
							from ..tf_utils import stable_softmax

lowerCAmelCase_   						=   logging.get_logger(__name__)
@add_end_docstrings(A_    )
class        __A       (	A_    ):







		'''simple docstring'''





		def __init__(	self						:  List[str]			,**_snake_case						:  Dict						)					->		List[Any]:






				"""simple docstring"""



				super().__init__(**_snake_case						)

				requires_backends(self			,'''vision'''						)
				self.check_model_type(
				    TF_MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING
				    if self.framework == '''tf'''
				    else MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING						)







		def __call__(	self						:  Optional[int]			,_snake_case						:  Union[str, List[str], "Image", List["Image"]]			,**_snake_case						:  int						)					->		Optional[Any]:






				"""simple docstring"""



				return super().__call__(_snake_case			,**_snake_case						)







		def 					UpperCAmelCase    (	self						:  Dict			,**_snake_case						:  Optional[int]						)					->		List[Any]:






				"""simple docstring"""



				lowercase__   :    List[str]				=			{}
				if "candidate_labels" in kwargs:
						lowercase__   :    Any				=			kwargs['''candidate_labels''']
				if "hypothesis_template" in kwargs:
						lowercase__   :    Optional[Any]				=			kwargs['''hypothesis_template''']

				return preprocess_params, {}, {}







		def 					UpperCAmelCase    (	self						:  Optional[Any]			,_snake_case						:  Optional[int]			,_snake_case						:  Dict=None			,_snake_case						:  Union[str, Any]="This is a photo of {}."						)					->		List[str]:






				"""simple docstring"""



				lowercase__   :    List[Any]				=			load_image(_snake_case						)
				lowercase__   :    int				=			self.image_processor(images=[image]			,return_tensors=self.framework						)
				lowercase__   :    str				=			candidate_labels
				lowercase__   :    Dict				=			[hypothesis_template.format(_snake_case						) for x in candidate_labels]
				lowercase__   :    Any				=			self.tokenizer(_snake_case			,return_tensors=self.framework			,padding=_snake_case						)
				lowercase__   :    Optional[int]				=			[text_inputs]
				return inputs







		def 					UpperCAmelCase    (	self						:  Optional[Any]			,_snake_case						:  Optional[int]						)					->		List[Any]:






				"""simple docstring"""



				lowercase__   :    Optional[int]				=			model_inputs.pop('''candidate_labels'''						)
				lowercase__   :    Union[str, Any]				=			model_inputs.pop('''text_inputs'''						)
				if isinstance(text_inputs[0]			,_snake_case						):
						lowercase__   :    List[str]				=			text_inputs[0]
				else:
						# Batching case.
						lowercase__   :    int				=			text_inputs[0][0]

				lowercase__   :    Tuple				=			self.model(**_snake_case			,**_snake_case						)

				lowercase__   :    Union[str, Any]				=			{
				    '''candidate_labels''': candidate_labels,
				    '''logits''': outputs.logits_per_image,
				}
				return model_outputs







		def 					UpperCAmelCase    (	self						:  Any			,_snake_case						:  Tuple						)					->		Any:






				"""simple docstring"""



				lowercase__   :    Dict				=			model_outputs.pop('''candidate_labels'''						)
				lowercase__   :    Optional[Any]				=			model_outputs['''logits'''][0]
				if self.framework == "pt":
						lowercase__   :    Optional[int]				=			logits.softmax(dim=-1						).squeeze(-1						)
						lowercase__   :    Tuple				=			probs.tolist()
						if not isinstance(_snake_case			,_snake_case						):
								lowercase__   :    Any				=			[scores]
				elif self.framework == "tf":
						lowercase__   :    List[str]				=			stable_softmax(_snake_case			,axis=-1						)
						lowercase__   :    Optional[Any]				=			probs.numpy().tolist()
				else:
						raise ValueError(f"""Unsupported framework: {self.framework}"""						)

				lowercase__   :    Union[str, Any]				=			[
				    {'''score''': score, '''label''': candidate_label}
				    for score, candidate_label in sorted(zip(_snake_case			,_snake_case						)			,key=lambda _snake_case						: -x[0]						)
				]
				return result
