import unittest

from transformers import MPNetConfig, 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, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch

		from transformers import (
		    MPNetForMaskedLM,
		    MPNetForMultipleChoice,
		    MPNetForQuestionAnswering,
		    MPNetForSequenceClassification,
		    MPNetForTokenClassification,
		    MPNetModel,
		)




class   _UpperCamelCase :


			"""simple docstring"""






			def __init__(  self ,  lowerCAmelCase__ ,  lowerCAmelCase__=13 ,  lowerCAmelCase__=7 ,  lowerCAmelCase__=True ,  lowerCAmelCase__=True ,  lowerCAmelCase__=False ,  lowerCAmelCase__=True ,  lowerCAmelCase__=99 ,  lowerCAmelCase__=64 ,  lowerCAmelCase__=5 ,  lowerCAmelCase__=4 ,  lowerCAmelCase__=64 ,  lowerCAmelCase__="gelu" ,  lowerCAmelCase__=0.1 ,  lowerCAmelCase__=0.1 ,  lowerCAmelCase__=5_12 ,  lowerCAmelCase__=16 ,  lowerCAmelCase__=2 ,  lowerCAmelCase__=0.02 ,  lowerCAmelCase__=3 ,  lowerCAmelCase__=4 ,  lowerCAmelCase__=None ,  )	->   str:

							'''simple docstring'''


							__lowercase					     =				parent
							__lowercase					     =				batch_size
							__lowercase					     =				seq_length
							__lowercase					     =				is_training
							__lowercase					     =				use_input_mask
							__lowercase					     =				use_token_type_ids
							__lowercase					     =				use_labels
							__lowercase					     =				vocab_size
							__lowercase					     =				hidden_size
							__lowercase					     =				num_hidden_layers
							__lowercase					     =				num_attention_heads
							__lowercase					     =				intermediate_size
							__lowercase					     =				hidden_act
							__lowercase					     =				hidden_dropout_prob
							__lowercase					     =				attention_probs_dropout_prob
							__lowercase					     =				max_position_embeddings
							__lowercase					     =				type_vocab_size
							__lowercase					     =				type_sequence_label_size
							__lowercase					     =				initializer_range
							__lowercase					     =				num_labels
							__lowercase					     =				num_choices
							__lowercase					     =				scope





			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   str:

							'''simple docstring'''


							return MPNetConfig.from_pretrained('''microsoft/mpnet-base'''		)





			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   int:

							'''simple docstring'''


							__lowercase					     =				ids_tensor([self.batch_size, self.seq_length] ,  self.vocab_size		)

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

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

							__lowercase					     =				self.get_config()
							return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels





			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   Tuple:

							'''simple docstring'''


							return MPNetConfig(
							    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 ,  initializer_range=self.initializer_range ,  )





			def 			_SCREAMING_SNAKE_CASE     (  self ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__		)	->   Any:

							'''simple docstring'''


							__lowercase					     =				MPNetModel(config=lowerCAmelCase__		)
							model.to(lowerCAmelCase__		)
							model.eval()
							__lowercase					     =				model(lowerCAmelCase__ ,  lowerCAmelCase__		)
							__lowercase					     =				model(lowerCAmelCase__		)
							self.parent.assertEqual(result.last_hidden_state.shape ,  (self.batch_size, self.seq_length, self.hidden_size)		)
							self.parent.assertEqual(result.pooler_output.shape ,  (self.batch_size, self.hidden_size)		)





			def 			_SCREAMING_SNAKE_CASE     (  self ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__		)	->   str:

							'''simple docstring'''


							__lowercase					     =				MPNetForQuestionAnswering(config=lowerCAmelCase__		)
							model.to(lowerCAmelCase__		)
							model.eval()
							__lowercase					     =				model(
							    lowerCAmelCase__ ,  attention_mask=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 			_SCREAMING_SNAKE_CASE     (  self ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__		)	->   List[str]:

							'''simple docstring'''


							__lowercase					     =				self.num_labels
							__lowercase					     =				MPNetForSequenceClassification(lowerCAmelCase__		)
							model.to(lowerCAmelCase__		)
							model.eval()
							__lowercase					     =				model(lowerCAmelCase__ ,  attention_mask=lowerCAmelCase__ ,  labels=lowerCAmelCase__		)
							self.parent.assertEqual(result.logits.shape ,  (self.batch_size, self.num_labels)		)





			def 			_SCREAMING_SNAKE_CASE     (  self ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__		)	->   List[str]:

							'''simple docstring'''


							__lowercase					     =				self.num_choices
							__lowercase					     =				MPNetForMultipleChoice(config=lowerCAmelCase__		)
							model.to(lowerCAmelCase__		)
							model.eval()
							__lowercase					     =				input_ids.unsqueeze(1		).expand(-1 ,  self.num_choices ,  -1		).contiguous()
							__lowercase					     =				input_mask.unsqueeze(1		).expand(-1 ,  self.num_choices ,  -1		).contiguous()
							__lowercase					     =				model(
							    lowerCAmelCase__ ,  attention_mask=lowerCAmelCase__ ,  labels=lowerCAmelCase__ ,  )
							self.parent.assertEqual(result.logits.shape ,  (self.batch_size, self.num_choices)		)





			def 			_SCREAMING_SNAKE_CASE     (  self ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__		)	->   Optional[Any]:

							'''simple docstring'''


							__lowercase					     =				self.num_labels
							__lowercase					     =				MPNetForTokenClassification(config=lowerCAmelCase__		)
							model.to(lowerCAmelCase__		)
							model.eval()
							__lowercase					     =				model(lowerCAmelCase__ ,  attention_mask=lowerCAmelCase__ ,  labels=lowerCAmelCase__		)
							self.parent.assertEqual(result.logits.shape ,  (self.batch_size, self.seq_length, self.num_labels)		)





			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   Optional[int]:

							'''simple docstring'''


							__lowercase					     =				self.prepare_config_and_inputs()
							((__lowercase)     ,	(__lowercase)     ,	(__lowercase)     ,	(__lowercase)     ,	(__lowercase)     ,	(__lowercase))					     =				config_and_inputs
							__lowercase					     =				{'''input_ids''': input_ids, '''attention_mask''': input_mask}
							return config, inputs_dict




@require_torch
class   _UpperCamelCase (		_UpperCAmelCase		,_UpperCAmelCase		,unittest.TestCase					):


			"""simple docstring"""






			__a    :					Any     					=		(
			    (
			        MPNetForMaskedLM,
			        MPNetForMultipleChoice,
			        MPNetForQuestionAnswering,
			        MPNetForSequenceClassification,
			        MPNetForTokenClassification,
			        MPNetModel,
			    )
			    if is_torch_available()
			    else ()
			)
			__a    :					str     					=		(
			    {
			        '''feature-extraction''': MPNetModel,
			        '''fill-mask''': MPNetForMaskedLM,
			        '''question-answering''': MPNetForQuestionAnswering,
			        '''text-classification''': MPNetForSequenceClassification,
			        '''token-classification''': MPNetForTokenClassification,
			        '''zero-shot''': MPNetForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			__a    :					Union[str, Any]     					=		False
			__a    :					Optional[int]     					=		True





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

							'''simple docstring'''


							__lowercase					     =				MPNetModelTester(self		)
							__lowercase					     =				ConfigTester(self ,  config_class=lowerCAmelCase__ ,  hidden_size=37		)





			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   str:

							'''simple docstring'''


							self.config_tester.run_common_tests()





			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   Optional[int]:

							'''simple docstring'''


							__lowercase					     =				self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_mpnet_model(*lowerCAmelCase__		)





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

							'''simple docstring'''


							__lowercase					     =				self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_mpnet_for_sequence_classification(*lowerCAmelCase__		)





			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   Optional[Any]:

							'''simple docstring'''


							__lowercase					     =				self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_mpnet_for_multiple_choice(*lowerCAmelCase__		)





			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   Any:

							'''simple docstring'''


							__lowercase					     =				self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_mpnet_for_token_classification(*lowerCAmelCase__		)





			def 			_SCREAMING_SNAKE_CASE     (  self		)	->   Any:

							'''simple docstring'''


							__lowercase					     =				self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_mpnet_for_question_answering(*lowerCAmelCase__		)




@require_torch
class   _UpperCamelCase (		unittest.TestCase					):


			"""simple docstring"""






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

							'''simple docstring'''


							__lowercase					     =				MPNetModel.from_pretrained('''microsoft/mpnet-base'''		)
							__lowercase					     =				torch.tensor([[0, 3_45, 2_32, 3_28, 7_40, 1_40, 16_95, 69, 60_78, 15_88, 2]]		)
							__lowercase					     =				model(lowerCAmelCase__		)[0]
							__lowercase					     =				torch.Size((1, 11, 7_68)		)
							self.assertEqual(output.shape ,  lowerCAmelCase__		)
							__lowercase					     =				torch.tensor(
							    [[[-0.0550, 0.1943, -0.0740], [-0.0562, 0.2211, -0.0579], [-0.0437, 0.3337, -0.0641]]]		)
							# compare the actual values for a slice.
							self.assertTrue(torch.allclose(output[:, :3, :3] ,  lowerCAmelCase__ ,  atol=1E-4		)		)
import torch
from torch import nn

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin




class   _UpperCamelCase (		_UpperCAmelCase		,_UpperCAmelCase					):


			"""simple docstring"""






			@register_to_config
			def __init__(  self ,  *,
			    lowerCAmelCase__ = 4 ,  lowerCAmelCase__ = 7_68 ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  )	->   Dict:

							'''simple docstring'''


							super().__init__()

							__lowercase					     =				nn.Parameter(torch.zeros(lowerCAmelCase__		)		)

							# parameters for additional clip time embeddings
							__lowercase					     =				nn.Linear(lowerCAmelCase__ ,  lowerCAmelCase__		)
							__lowercase					     =				nn.Linear(lowerCAmelCase__ ,  lowerCAmelCase__		)

							# parameters for encoder hidden states
							__lowercase					     =				clip_extra_context_tokens
							__lowercase					     =				nn.Linear(
							    lowerCAmelCase__ ,  self.clip_extra_context_tokens * cross_attention_dim		)
							__lowercase					     =				nn.Linear(lowerCAmelCase__ ,  lowerCAmelCase__		)
							__lowercase					     =				nn.LayerNorm(lowerCAmelCase__		)





			def 			_SCREAMING_SNAKE_CASE     (  self ,  *, lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__		)	->   List[str]:

							'''simple docstring'''


							if do_classifier_free_guidance:
											# Add the classifier free guidance embeddings to the image embeddings
											__lowercase					     =				image_embeddings.shape[0]
											__lowercase					     =				self.learned_classifier_free_guidance_embeddings.unsqueeze(0		)
											__lowercase					     =				classifier_free_guidance_embeddings.expand(
											    lowerCAmelCase__ ,  -1		)
											__lowercase					     =				torch.cat([classifier_free_guidance_embeddings, image_embeddings] ,  dim=0		)

							# The image embeddings batch size and the text embeddings batch size are equal
							assert image_embeddings.shape[0] == prompt_embeds.shape[0]

							__lowercase					     =				prompt_embeds.shape[0]

							# "Specifically, we modify the architecture described in Nichol et al. (2021) by projecting and
							# adding CLIP embeddings to the existing timestep embedding, ...
							__lowercase					     =				self.embedding_proj(lowerCAmelCase__		)
							__lowercase					     =				self.clip_image_embeddings_project_to_time_embeddings(lowerCAmelCase__		)
							__lowercase					     =				time_projected_image_embeddings + time_projected_prompt_embeds

							# ... and by projecting CLIP embeddings into four
							# extra tokens of context that are concatenated to the sequence of outputs from the GLIDE text encoder"
							__lowercase					     =				self.clip_extra_context_tokens_proj(lowerCAmelCase__		)
							__lowercase					     =				clip_extra_context_tokens.reshape(lowerCAmelCase__ ,  -1 ,  self.clip_extra_context_tokens		)
							__lowercase					     =				clip_extra_context_tokens.permute(0 ,  2 ,  1		)

							__lowercase					     =				self.encoder_hidden_states_proj(lowerCAmelCase__		)
							__lowercase					     =				self.text_encoder_hidden_states_norm(lowerCAmelCase__		)
							__lowercase					     =				torch.cat([clip_extra_context_tokens, text_encoder_hidden_states] ,  dim=1		)

							return text_encoder_hidden_states, additive_clip_time_embeddings