

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				:			Union[str, Any]   ,   a				:			Tuple   ,   a				:			List[str]=13   ,   a				:			List[str]=7   ,   a				:			Union[str, Any]=True   ,   a				:			List[Any]=True   ,   a				:			Any=False   ,   a				:			List[str]=True   ,   a				:			str=99   ,   a				:			Union[str, Any]=64   ,   a				:			Any=5   ,   a				:			Dict=4   ,   a				:			List[Any]=64   ,   a				:			Optional[Any]="gelu"   ,   a				:			Tuple=0.1   ,   a				:			Union[str, Any]=0.1   ,   a				:			Dict=512   ,   a				:			List[str]=16   ,   a				:			Tuple=2   ,   a				:			Optional[int]=0.02   ,   a				:			Tuple=3   ,   a				:			List[Any]=4   ,   a				:			Any=None   ,   )				->							Optional[Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Dict					  =							parent
				SCREAMING_SNAKE_CASE :   List[Any]					  =							batch_size
				SCREAMING_SNAKE_CASE :   Any					  =							seq_length
				SCREAMING_SNAKE_CASE :   Tuple					  =							is_training
				SCREAMING_SNAKE_CASE :   List[Any]					  =							use_input_mask
				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							use_token_type_ids
				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							use_labels
				SCREAMING_SNAKE_CASE :   str					  =							vocab_size
				SCREAMING_SNAKE_CASE :   int					  =							hidden_size
				SCREAMING_SNAKE_CASE :   Tuple					  =							num_hidden_layers
				SCREAMING_SNAKE_CASE :   List[Any]					  =							num_attention_heads
				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							intermediate_size
				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							hidden_act
				SCREAMING_SNAKE_CASE :   Dict					  =							hidden_dropout_prob
				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							attention_probs_dropout_prob
				SCREAMING_SNAKE_CASE :   Dict					  =							max_position_embeddings
				SCREAMING_SNAKE_CASE :   List[str]					  =							type_vocab_size
				SCREAMING_SNAKE_CASE :   str					  =							type_sequence_label_size
				SCREAMING_SNAKE_CASE :   List[Any]					  =							initializer_range
				SCREAMING_SNAKE_CASE :   Dict					  =							num_labels
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							num_choices
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							scope
			def 				__UpperCamelCase			(      self				:			int						)				->							List[str]:

				"""simple docstring"""

				return MPNetConfig.from_pretrained("microsoft/mpnet-base"						)
			def 				__UpperCamelCase			(      self				:			List[Any]						)				->							int:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							ids_tensor([self.batch_size, self.seq_length]   ,   self.vocab_size						)

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

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

				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							self.get_config()
				return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
			def 				__UpperCamelCase			(      self				:			List[Any]						)				->							Dict:

				"""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 				__UpperCamelCase			(      self				:			Optional[int]   ,   a				:			Union[str, Any]   ,   a				:			str   ,   a				:			int   ,   a				:			List[str]   ,   a				:			Union[str, Any]   ,   a				:			List[str]						)				->							Optional[int]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Optional[int]					  =							MPNetModel(config=a						)
				model.to(a						)
				model.eval()
				SCREAMING_SNAKE_CASE :   Tuple					  =							model(a   ,   a						)
				SCREAMING_SNAKE_CASE :   int					  =							model(a						)
				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 				__UpperCamelCase			(      self				:			List[str]   ,   a				:			Any   ,   a				:			Union[str, Any]   ,   a				:			Dict   ,   a				:			int   ,   a				:			List[Any]   ,   a				:			Dict						)				->							int:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							MPNetForQuestionAnswering(config=a						)
				model.to(a						)
				model.eval()
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							model(
				    a   ,   attention_mask=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				:			Tuple   ,   a				:			Union[str, Any]   ,   a				:			int   ,   a				:			Optional[int]   ,   a				:			Dict   ,   a				:			Any   ,   a				:			List[str]						)				->							int:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							self.num_labels
				SCREAMING_SNAKE_CASE :   Dict					  =							MPNetForSequenceClassification(a						)
				model.to(a						)
				model.eval()
				SCREAMING_SNAKE_CASE :   List[Any]					  =							model(a   ,   attention_mask=a   ,   labels=a						)
				self.parent.assertEqual(result.logits.shape   ,   (self.batch_size, self.num_labels)						)
			def 				__UpperCamelCase			(      self				:			List[str]   ,   a				:			Any   ,   a				:			Optional[Any]   ,   a				:			Dict   ,   a				:			Optional[Any]   ,   a				:			int   ,   a				:			Tuple						)				->							List[str]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							self.num_choices
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							MPNetForMultipleChoice(config=a						)
				model.to(a						)
				model.eval()
				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							input_ids.unsqueeze(1						).expand(-1   ,   self.num_choices   ,   -1						).contiguous()
				SCREAMING_SNAKE_CASE :   Tuple					  =							input_mask.unsqueeze(1						).expand(-1   ,   self.num_choices   ,   -1						).contiguous()
				SCREAMING_SNAKE_CASE :   List[Any]					  =							model(
				    a   ,   attention_mask=a   ,   labels=a   ,   )
				self.parent.assertEqual(result.logits.shape   ,   (self.batch_size, self.num_choices)						)
			def 				__UpperCamelCase			(      self				:			Optional[Any]   ,   a				:			str   ,   a				:			Tuple   ,   a				:			Optional[int]   ,   a				:			List[Any]   ,   a				:			Tuple   ,   a				:			int						)				->							List[str]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Tuple					  =							self.num_labels
				SCREAMING_SNAKE_CASE :   List[str]					  =							MPNetForTokenClassification(config=a						)
				model.to(a						)
				model.eval()
				SCREAMING_SNAKE_CASE :   List[Any]					  =							model(a   ,   attention_mask=a   ,   labels=a						)
				self.parent.assertEqual(result.logits.shape   ,   (self.batch_size, self.seq_length, self.num_labels)						)






			def 				__UpperCamelCase			(      self				:			Dict						)				->							Optional[Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Any					  =							self.prepare_config_and_inputs()
				((SCREAMING_SNAKE_CASE)  ,(SCREAMING_SNAKE_CASE)  ,(SCREAMING_SNAKE_CASE)  ,(SCREAMING_SNAKE_CASE)  ,(SCREAMING_SNAKE_CASE)  ,(SCREAMING_SNAKE_CASE)) :   Optional[Any]					  =							config_and_inputs
				SCREAMING_SNAKE_CASE :   List[str]					  =							{"input_ids": input_ids, "attention_mask": input_mask}
				return config, inputs_dict






@require_torch
class _UpperCamelCase				(				__A				,						__A				,						unittest.TestCase ):
			'''simple docstring'''







			lowerCamelCase__							     =(
			    (
			        MPNetForMaskedLM,
			        MPNetForMultipleChoice,
			        MPNetForQuestionAnswering,
			        MPNetForSequenceClassification,
			        MPNetForTokenClassification,
			        MPNetModel,
			    )
			    if is_torch_available()
			    else ()
			)
			lowerCamelCase__							     =(
			    {
			        'feature-extraction': MPNetModel,
			        'fill-mask': MPNetForMaskedLM,
			        'question-answering': MPNetForQuestionAnswering,
			        'text-classification': MPNetForSequenceClassification,
			        'token-classification': MPNetForTokenClassification,
			        'zero-shot': MPNetForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			lowerCamelCase__							     =False
			lowerCamelCase__							     =True
			def 				__UpperCamelCase			(      self				:			Tuple						)				->							Tuple:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   List[str]					  =							MPNetModelTester(self						)
				SCREAMING_SNAKE_CASE :   str					  =							ConfigTester(self   ,   config_class=a   ,   hidden_size=37						)
			def 				__UpperCamelCase			(      self				:			int						)				->							List[str]:

				"""simple docstring"""

				self.config_tester.run_common_tests()
			def 				__UpperCamelCase			(      self				:			Union[str, Any]						)				->							Optional[int]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   str					  =							self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_mpnet_model(*a						)
			def 				__UpperCamelCase			(      self				:			Union[str, Any]						)				->							str:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Tuple					  =							self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_mpnet_for_sequence_classification(*a						)
			def 				__UpperCamelCase			(      self				:			Dict						)				->							List[Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   List[Any]					  =							self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_mpnet_for_multiple_choice(*a						)
			def 				__UpperCamelCase			(      self				:			Any						)				->							int:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Optional[int]					  =							self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_mpnet_for_token_classification(*a						)






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

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Tuple					  =							self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_mpnet_for_question_answering(*a						)






@require_torch
class _UpperCamelCase				(				unittest.TestCase ):
			'''simple docstring'''







			@slow
			def 				__UpperCamelCase			(      self				:			List[Any]						)				->							Optional[Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   List[str]					  =							MPNetModel.from_pretrained("microsoft/mpnet-base"						)
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]						)
				SCREAMING_SNAKE_CASE :   List[str]					  =							model(a						)[0]
				SCREAMING_SNAKE_CASE :   Tuple					  =							torch.Size((1, 11, 768)						)
				self.assertEqual(output.shape   ,   a						)
				SCREAMING_SNAKE_CASE :   List[Any]					  =							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]   ,   a   ,   atol=1e-4						)						)




'''simple docstring'''





from typing import Optional, Tuple

import jax
import jax.numpy as jnp
from flax import linen as nn
from flax.core.frozen_dict import FrozenDict
from transformers import CLIPConfig, FlaxPreTrainedModel
from transformers.models.clip.modeling_flax_clip import FlaxCLIPVisionModule

def 			a_	(				_lowerCAmelCase			,_lowerCAmelCase			,_lowerCAmelCase=1E-12			) -> List[Any]:
  __lowerCamelCase			:			int       			=					jnp.divide(emb_a.T			,jnp.clip(jnp.linalg.norm(_lowerCAmelCase			,axis=1			)			,a_min=_lowerCAmelCase			)			).T
  __lowerCamelCase			:			Tuple       			=					jnp.divide(emb_a.T			,jnp.clip(jnp.linalg.norm(_lowerCAmelCase			,axis=1			)			,a_min=_lowerCAmelCase			)			).T
  return jnp.matmul(_lowerCAmelCase			,norm_emb_a.T			)






class    lowerCamelCase_       (							nn.Module	):
  """simple docstring"""




  a_				   =42
  a_				   =jnp.floataa




  def 		_lowercase	(					self :      Optional[Any]  )   ->   Tuple:
    __lowerCamelCase			:			Optional[int]       			=					FlaxCLIPVisionModule(self.config.vision_config  )
    __lowerCamelCase			:			Any       			=					nn.Dense(self.config.projection_dim     ,      use_bias=_a     ,      dtype=self.dtype  )

    __lowerCamelCase			:			str       			=					self.param('concept_embeds'     ,      jax.nn.initializers.ones     ,      (17, self.config.projection_dim)  )
    __lowerCamelCase			:			str       			=					self.param(
        'special_care_embeds'     ,      jax.nn.initializers.ones     ,      (3, self.config.projection_dim)  )

    __lowerCamelCase			:			Optional[int]       			=					self.param('concept_embeds_weights'     ,      jax.nn.initializers.ones     ,      (17,)  )
    __lowerCamelCase			:			Dict       			=					self.param('special_care_embeds_weights'     ,      jax.nn.initializers.ones     ,      (3,)  )




  def __call__(					self :      Tuple     ,      _a :      Any  )   ->   Optional[int]:
    __lowerCamelCase			:			int       			=					self.vision_model(_a  )[1]
    __lowerCamelCase			:			Optional[int]       			=					self.visual_projection(_a  )

    __lowerCamelCase			:			Optional[Any]       			=					jax_cosine_distance(_a     ,      self.special_care_embeds  )
    __lowerCamelCase			:			Tuple       			=					jax_cosine_distance(_a     ,      self.concept_embeds  )

    # increase this value to create a stronger `nfsw` filter
    # at the cost of increasing the possibility of filtering benign image inputs
    __lowerCamelCase			:			Dict       			=					0.0

    __lowerCamelCase			:			str       			=					special_cos_dist - self.special_care_embeds_weights[None, :] + adjustment
    __lowerCamelCase			:			Optional[int]       			=					jnp.round(_a     ,      3  )
    __lowerCamelCase			:			Dict       			=					jnp.any(special_scores > 0     ,      axis=1     ,      keepdims=_a  )
    # Use a lower threshold if an image has any special care concept
    __lowerCamelCase			:			Dict       			=					is_special_care * 0.01

    __lowerCamelCase			:			Optional[int]       			=					cos_dist - self.concept_embeds_weights[None, :] + special_adjustment
    __lowerCamelCase			:			List[str]       			=					jnp.round(_a     ,      3  )
    __lowerCamelCase			:			List[Any]       			=					jnp.any(concept_scores > 0     ,      axis=1  )

    return has_nsfw_concepts






class    lowerCamelCase_       (							SCREAMING_SNAKE_CASE_	):
  """simple docstring"""




  a_				   =CLIPConfig
  a_				   ="""clip_input"""
  a_				   =FlaxStableDiffusionSafetyCheckerModule




  def __init__(					self :      List[str]     ,      _a :      CLIPConfig     ,      _a :      Optional[Tuple] = None     ,      _a :      int = 0     ,      _a :      jnp.dtype = jnp.floataa     ,      _a :      bool = True     ,      **_a :      List[Any]     ,      )   ->   List[Any]:
    if input_shape is None:
      __lowerCamelCase			:			str       			=					(1, 224, 224, 3)
    __lowerCamelCase			:			List[str]       			=					self.module_class(config=_a     ,      dtype=_a     ,      **_a  )
    super().__init__(_a     ,      _a     ,      input_shape=_a     ,      seed=_a     ,      dtype=_a     ,      _do_init=_do_init  )




  def 		_lowercase	(					self :      int     ,      _a :      jax.random.KeyArray     ,      _a :      Tuple     ,      _a :      FrozenDict = None  )   ->   FrozenDict:
    # init input tensor
    __lowerCamelCase			:			Optional[int]       			=					jax.random.normal(_a     ,      _a  )

    __lowerCamelCase			,__lowerCamelCase			:			Optional[Any]       			=					jax.random.split(_a  )
    __lowerCamelCase			:			int       			=					{'params': params_rng, 'dropout': dropout_rng}

    __lowerCamelCase			:			Tuple       			=					self.module.init(_a     ,      _a  )['params']

    return random_params




  def __call__(					self :      Optional[int]     ,      _a :      int     ,      _a :      dict = None     ,      )   ->   Dict:
    __lowerCamelCase			:			str       			=					jnp.transpose(_a     ,      (0, 2, 3, 1)  )

    return self.module.apply(
        {'params': params or self.params}     ,      jnp.array(_a     ,      dtype=jnp.floataa  )     ,      rngs={}     ,      )






