



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


					'''simple docstring'''





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



					def 		snake_case	(	self							:  List[str]					)-> str:
								lowerCamelCase__		:		List[Any]		     =ids_tensor([self.batch_size, self.seq_length],  self.vocab_size					)

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

								lowerCamelCase__		:		List[Any]		     =None
								if self.use_token_type_ids:
											lowerCamelCase__		:		Union[str, Any]		     =ids_tensor([self.batch_size, self.seq_length],  self.type_vocab_size					)

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

								lowerCamelCase__		:		int		     =self.get_config()

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



					def 		snake_case	(	self							:  Tuple					)-> int:
								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 		snake_case	(	self							:  int					)-> int:
								lowerCamelCase__		:		Any		     =self.get_config()
								lowerCamelCase__		:		Any		     =300
								return config



					def 		snake_case	(	self							:  str					)-> str:
								(
								    (
								    lowerCamelCase__
								)     ,       (
								    lowerCamelCase__
								)     ,       (
								    lowerCamelCase__
								)     ,       (
								    lowerCamelCase__
								)     ,       (
								    lowerCamelCase__
								)     ,       (
								    lowerCamelCase__
								)     ,       (
								    lowerCamelCase__
								)     ,       
								)		:		Any		     =self.prepare_config_and_inputs()

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



					def 		snake_case	(	self							:  int,  lowerCamelCase							:  Any,  lowerCamelCase							:  int,  lowerCamelCase							:  List[Any],  lowerCamelCase							:  Dict,  lowerCamelCase							:  int,  lowerCamelCase							:  Optional[int],  lowerCamelCase							:  Dict,  lowerCamelCase							:  Optional[Any],  lowerCamelCase							:  Dict,  )-> Tuple:
								lowerCamelCase__		:		Tuple		     =True
								lowerCamelCase__		:		Optional[Any]		     =MraModel(lowerCamelCase					)
								model.to(lowerCamelCase					)
								model.eval()
								lowerCamelCase__		:		List[Any]		     =model(
								    lowerCamelCase,  attention_mask=lowerCamelCase,  token_type_ids=lowerCamelCase,  encoder_hidden_states=lowerCamelCase,  encoder_attention_mask=lowerCamelCase,  )
								lowerCamelCase__		:		str		     =model(
								    lowerCamelCase,  attention_mask=lowerCamelCase,  token_type_ids=lowerCamelCase,  encoder_hidden_states=lowerCamelCase,  )
								lowerCamelCase__		:		Dict		     =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 		snake_case	(	self							:  str,  lowerCamelCase							:  Tuple,  lowerCamelCase							:  Dict,  lowerCamelCase							:  str,  lowerCamelCase							:  List[str],  lowerCamelCase							:  Optional[int],  lowerCamelCase							:  Optional[int],  lowerCamelCase							:  Optional[Any]					)-> Optional[int]:
								lowerCamelCase__		:		int		     =MraForMaskedLM(config=lowerCamelCase					)
								model.to(lowerCamelCase					)
								model.eval()
								lowerCamelCase__		:		List[str]		     =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 		snake_case	(	self							:  Optional[int],  lowerCamelCase							:  int,  lowerCamelCase							:  Any,  lowerCamelCase							:  Any,  lowerCamelCase							:  Union[str, Any],  lowerCamelCase							:  int,  lowerCamelCase							:  Dict,  lowerCamelCase							:  Union[str, Any]					)-> Tuple:
								lowerCamelCase__		:		Tuple		     =MraForQuestionAnswering(config=lowerCamelCase					)
								model.to(lowerCamelCase					)
								model.eval()
								lowerCamelCase__		:		List[Any]		     =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 		snake_case	(	self							:  Union[str, Any],  lowerCamelCase							:  int,  lowerCamelCase							:  Optional[int],  lowerCamelCase							:  str,  lowerCamelCase							:  int,  lowerCamelCase							:  int,  lowerCamelCase							:  Dict,  lowerCamelCase							:  Union[str, Any]					)-> Union[str, Any]:
								lowerCamelCase__		:		Optional[Any]		     =self.num_labels
								lowerCamelCase__		:		List[Any]		     =MraForSequenceClassification(lowerCamelCase					)
								model.to(lowerCamelCase					)
								model.eval()
								lowerCamelCase__		:		Any		     =model(lowerCamelCase,  attention_mask=lowerCamelCase,  token_type_ids=lowerCamelCase,  labels=lowerCamelCase					)
								self.parent.assertEqual(result.logits.shape,  (self.batch_size, self.num_labels)					)



					def 		snake_case	(	self							:  Any,  lowerCamelCase							:  Any,  lowerCamelCase							:  str,  lowerCamelCase							:  Optional[Any],  lowerCamelCase							:  str,  lowerCamelCase							:  Any,  lowerCamelCase							:  Union[str, Any],  lowerCamelCase							:  Optional[Any]					)-> str:
								lowerCamelCase__		:		Union[str, Any]		     =self.num_labels
								lowerCamelCase__		:		Union[str, Any]		     =MraForTokenClassification(config=lowerCamelCase					)
								model.to(lowerCamelCase					)
								model.eval()
								lowerCamelCase__		:		Optional[int]		     =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 		snake_case	(	self							:  Dict,  lowerCamelCase							:  Optional[int],  lowerCamelCase							:  List[Any],  lowerCamelCase							:  List[str],  lowerCamelCase							:  Tuple,  lowerCamelCase							:  Any,  lowerCamelCase							:  int,  lowerCamelCase							:  Tuple					)-> int:
								lowerCamelCase__		:		Tuple		     =self.num_choices
								lowerCamelCase__		:		Dict		     =MraForMultipleChoice(config=lowerCamelCase					)
								model.to(lowerCamelCase					)
								model.eval()
								lowerCamelCase__		:		Optional[int]		     =input_ids.unsqueeze(1					).expand(-1,  self.num_choices,  -1					).contiguous()
								lowerCamelCase__		:		Union[str, Any]		     =token_type_ids.unsqueeze(1					).expand(-1,  self.num_choices,  -1					).contiguous()
								lowerCamelCase__		:		Dict		     =input_mask.unsqueeze(1					).expand(-1,  self.num_choices,  -1					).contiguous()
								lowerCamelCase__		:		Dict		     =model(
								    lowerCamelCase,  attention_mask=lowerCamelCase,  token_type_ids=lowerCamelCase,  labels=lowerCamelCase,  )
								self.parent.assertEqual(result.logits.shape,  (self.batch_size, self.num_choices)					)




					def 		snake_case	(	self							:  Dict					)-> Union[str, Any]:
								lowerCamelCase__		:		str		     =self.prepare_config_and_inputs()
								(
								    (
								    lowerCamelCase__
								)     ,       (
								    lowerCamelCase__
								)     ,       (
								    lowerCamelCase__
								)     ,       (
								    lowerCamelCase__
								)     ,       (
								    lowerCamelCase__
								)     ,       (
								    lowerCamelCase__
								)     ,       (
								    lowerCamelCase__
								)     ,       
								)		:		List[str]		     =config_and_inputs
								lowerCamelCase__		:		List[Any]		     ={'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
								return config, inputs_dict





@require_torch
class 	__SCREAMING_SNAKE_CASE     (				lowerCAmelCase_		,       unittest.TestCase       ):


					'''simple docstring'''





					_a             =							(
					    (
					        MraModel,
					        MraForMaskedLM,
					        MraForMultipleChoice,
					        MraForQuestionAnswering,
					        MraForSequenceClassification,
					        MraForTokenClassification,
					    )
					    if is_torch_available()
					    else ()
					)
					_a             =							False
					_a             =							False
					_a             =							False
					_a             =							False

					_a             =							()



					def 		snake_case	(	self							:  Optional[Any]					)-> List[str]:
								lowerCamelCase__		:		List[str]		     =MraModelTester(self					)
								lowerCamelCase__		:		List[Any]		     =ConfigTester(self,  config_class=lowerCamelCase,  hidden_size=37					)



					def 		snake_case	(	self							:  Optional[int]					)-> Any:
								self.config_tester.run_common_tests()



					def 		snake_case	(	self							:  int					)-> str:
								lowerCamelCase__		:		Tuple		     =self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*lowerCamelCase					)



					def 		snake_case	(	self							:  Union[str, Any]					)-> List[str]:
								lowerCamelCase__		:		List[Any]		     =self.model_tester.prepare_config_and_inputs()
								for type in ["absolute", "relative_key", "relative_key_query"]:
											lowerCamelCase__		:		Tuple		     =type
											self.model_tester.create_and_check_model(*lowerCamelCase					)



					def 		snake_case	(	self							:  List[Any]					)-> Optional[int]:
								lowerCamelCase__		:		Tuple		     =self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_lm(*lowerCamelCase					)



					def 		snake_case	(	self							:  str					)-> int:
								lowerCamelCase__		:		List[str]		     =self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_multiple_choice(*lowerCamelCase					)



					def 		snake_case	(	self							:  List[str]					)-> Any:
								lowerCamelCase__		:		List[str]		     =self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_question_answering(*lowerCamelCase					)



					def 		snake_case	(	self							:  Tuple					)-> Optional[int]:
								lowerCamelCase__		:		Optional[int]		     =self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_sequence_classification(*lowerCamelCase					)



					def 		snake_case	(	self							:  Dict					)-> int:
								lowerCamelCase__		:		Optional[Any]		     =self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_token_classification(*lowerCamelCase					)



					@slow
					def 		snake_case	(	self							:  Tuple					)-> Dict:
								for model_name in MRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											lowerCamelCase__		:		Optional[Any]		     =MraModel.from_pretrained(lowerCamelCase					)
											self.assertIsNotNone(lowerCamelCase					)




					@unittest.skip(reason='''MRA does not output attentions'''					)
					def 		snake_case	(	self							:  Tuple					)-> Optional[Any]:
								return





@require_torch
class 	__SCREAMING_SNAKE_CASE     (				unittest.TestCase       ):


					'''simple docstring'''





					@slow
					def 		snake_case	(	self							:  List[Any]					)-> Dict:
								lowerCamelCase__		:		Any		     =MraModel.from_pretrained('''uw-madison/mra-base-512-4'''					)
								lowerCamelCase__		:		Any		     =torch.arange(256					).unsqueeze(0					)

								with torch.no_grad():
											lowerCamelCase__		:		Dict		     =model(lowerCamelCase					)[0]

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

								lowerCamelCase__		:		List[Any]		     =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 		snake_case	(	self							:  Any					)-> List[Any]:
								lowerCamelCase__		:		List[Any]		     =MraForMaskedLM.from_pretrained('''uw-madison/mra-base-512-4'''					)
								lowerCamelCase__		:		Union[str, Any]		     =torch.arange(256					).unsqueeze(0					)

								with torch.no_grad():
											lowerCamelCase__		:		Optional[int]		     =model(lowerCamelCase					)[0]

								lowerCamelCase__		:		int		     =5_0265

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

								lowerCamelCase__		:		List[Any]		     =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 		snake_case	(	self							:  Tuple					)-> str:
								lowerCamelCase__		:		int		     =MraForMaskedLM.from_pretrained('''uw-madison/mra-base-4096-8-d3'''					)
								lowerCamelCase__		:		Optional[Any]		     =torch.arange(4096					).unsqueeze(0					)

								with torch.no_grad():
											lowerCamelCase__		:		Tuple		     =model(lowerCamelCase					)[0]

								lowerCamelCase__		:		Dict		     =5_0265

								lowerCamelCase__		:		Dict		     =torch.Size((1, 4096, vocab_size)					)
								self.assertEqual(output.shape,  lowerCamelCase					)

								lowerCamelCase__		:		Union[str, Any]		     =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					)					)









"""simple docstring"""






import numpy as np
import torch
import torch.nn as nn
from transformers import CLIPConfig, CLIPVisionModelWithProjection, PreTrainedModel

from ...utils import logging


_lowercase			:			int										=						logging.get_logger(__name__)

class 	__SCREAMING_SNAKE_CASE     (				lowerCAmelCase_       ):


					'''simple docstring'''





					_a             =							CLIPConfig

					_a             =							['CLIPEncoderLayer']



					def __init__(	self							:  int,  lowerCamelCase							:  CLIPConfig					)-> List[Any]:
								super().__init__(lowerCamelCase					)

								lowerCamelCase__		:		Dict		     =CLIPVisionModelWithProjection(config.vision_config					)

								lowerCamelCase__		:		Dict		     =nn.Linear(config.vision_config.projection_dim,  1					)
								lowerCamelCase__		:		Union[str, Any]		     =nn.Linear(config.vision_config.projection_dim,  1					)




					@torch.no_grad()
					def 		snake_case	(	self							:  Dict,  lowerCamelCase							:  Optional[Any],  lowerCamelCase							:  List[str],  lowerCamelCase							:  List[Any]=0.5,  lowerCamelCase							:  Optional[Any]=0.5					)-> Optional[int]:
								lowerCamelCase__		:		Dict		     =self.vision_model(lowerCamelCase					)[0]

								lowerCamelCase__		:		List[Any]		     =self.p_head(lowerCamelCase					)
								lowerCamelCase__		:		Optional[Any]		     =nsfw_detected.flatten()
								lowerCamelCase__		:		Any		     =nsfw_detected > p_threshold
								lowerCamelCase__		:		Optional[Any]		     =nsfw_detected.tolist()

								if any(lowerCamelCase					):
											logger.warning(
											    '''Potential NSFW content was detected in one or more images. A black image will be returned instead.'''
											    ''' Try again with a different prompt and/or seed.'''					)

								for idx, nsfw_detected_ in enumerate(lowerCamelCase					):
											if nsfw_detected_:
														lowerCamelCase__		:		Tuple		     =np.zeros(images[idx].shape					)

								lowerCamelCase__		:		str		     =self.w_head(lowerCamelCase					)
								lowerCamelCase__		:		List[str]		     =watermark_detected.flatten()
								lowerCamelCase__		:		Optional[Any]		     =watermark_detected > w_threshold
								lowerCamelCase__		:		int		     =watermark_detected.tolist()

								if any(lowerCamelCase					):
											logger.warning(
											    '''Potential watermarked content was detected in one or more images. A black image will be returned instead.'''
											    ''' Try again with a different prompt and/or seed.'''					)

								for idx, watermark_detected_ in enumerate(lowerCamelCase					):
											if watermark_detected_:
														lowerCamelCase__		:		Optional[int]		     =np.zeros(images[idx].shape					)

								return images, nsfw_detected, watermark_detected




