




'''simple docstring'''






from __future__ import annotations

import unittest

from transformers import MPNetConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
      import tensorflow as tf

      from transformers.models.mpnet.modeling_tf_mpnet import (
          TFMPNetForMaskedLM,
          TFMPNetForMultipleChoice,
          TFMPNetForQuestionAnswering,
          TFMPNetForSequenceClassification,
          TFMPNetForTokenClassification,
          TFMPNetModel,
      )

class   SCREAMING_SNAKE_CASE_		:



 def __init__(  self   ,		A   ,		A=13   ,		A=7   ,		A=True   ,		A=True   ,		A=False   ,		A=True   ,		A=99   ,		A=64   ,		A=2   ,		A=4   ,		A=64   ,		A="gelu"   ,		A=0.1   ,		A=0.1   ,		A=512   ,		A=16   ,		A=2   ,		A=0.0_2   ,		A=3   ,		A=4   ,		A=None   ,		):






        '''simple docstring'''




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



 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




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

        lowercase  :		str						=      None
        if self.use_input_mask:
               lowercase  :		Dict						=      random_attention_mask([self.batch_size, self.seq_length] )

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

        lowercase  :		int						=      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   ,		)
        return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels



 def    SCREAMING_SNAKE_CASE(  self   ,		A   ,		A   ,		A   ,		A   ,		A   ,		A ):






        '''simple docstring'''




        lowercase  :		List[Any]						=      TFMPNetModel(config=A )
        lowercase  :		int						=      {'''input_ids''': input_ids, '''attention_mask''': input_mask}
        lowercase  :		Any						=      model(A )
        lowercase  :		Union[str, Any]						=      [input_ids, input_mask]
        lowercase  :		str						=      model(A )
        self.parent.assertEqual(result.last_hidden_state.shape   ,		(self.batch_size, self.seq_length, self.hidden_size) )



 def    SCREAMING_SNAKE_CASE(  self   ,		A   ,		A   ,		A   ,		A   ,		A   ,		A ):






        '''simple docstring'''




        lowercase  :		str						=      TFMPNetForMaskedLM(config=A )
        lowercase  :		str						=      {'''input_ids''': input_ids, '''attention_mask''': input_mask}
        lowercase  :		str						=      model(A )
        self.parent.assertEqual(result.logits.shape   ,		(self.batch_size, self.seq_length, self.vocab_size) )



 def    SCREAMING_SNAKE_CASE(  self   ,		A   ,		A   ,		A   ,		A   ,		A   ,		A ):






        '''simple docstring'''




        lowercase  :		List[Any]						=      TFMPNetForQuestionAnswering(config=A )
        lowercase  :		Dict						=      {
            '''input_ids''': input_ids,
            '''attention_mask''': input_mask,
        }
        lowercase  :		List[str]						=      model(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    SCREAMING_SNAKE_CASE(  self   ,		A   ,		A   ,		A   ,		A   ,		A   ,		A ):






        '''simple docstring'''




        lowercase  :		List[str]						=      self.num_labels
        lowercase  :		List[Any]						=      TFMPNetForSequenceClassification(A )
        lowercase  :		int						=      {'''input_ids''': input_ids, '''attention_mask''': input_mask}
        lowercase  :		List[Any]						=      model(A )
        self.parent.assertEqual(result.logits.shape   ,		(self.batch_size, self.num_labels) )



 def    SCREAMING_SNAKE_CASE(  self   ,		A   ,		A   ,		A   ,		A   ,		A   ,		A ):






        '''simple docstring'''




        lowercase  :		List[Any]						=      self.num_choices
        lowercase  :		Optional[Any]						=      TFMPNetForMultipleChoice(A )
        lowercase  :		str						=      tf.tile(tf.expand_dims(A   ,		1 )   ,		(1, self.num_choices, 1) )
        lowercase  :		Tuple						=      tf.tile(tf.expand_dims(A   ,		1 )   ,		(1, self.num_choices, 1) )
        lowercase  :		Optional[int]						=      {
            '''input_ids''': multiple_choice_inputs_ids,
            '''attention_mask''': multiple_choice_input_mask,
        }
        lowercase  :		Optional[int]						=      model(A )
        self.parent.assertEqual(result.logits.shape   ,		(self.batch_size, self.num_choices) )



 def    SCREAMING_SNAKE_CASE(  self   ,		A   ,		A   ,		A   ,		A   ,		A   ,		A ):






        '''simple docstring'''




        lowercase  :		Optional[int]						=      self.num_labels
        lowercase  :		Dict						=      TFMPNetForTokenClassification(A )
        lowercase  :		List[str]						=      {'''input_ids''': input_ids, '''attention_mask''': input_mask}
        lowercase  :		str						=      model(A )
        self.parent.assertEqual(result.logits.shape   ,		(self.batch_size, self.seq_length, self.num_labels) )





 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




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






@require_tf
class   SCREAMING_SNAKE_CASE_		(       lowerCamelCase_		,     lowerCamelCase_		,     unittest.TestCase				):
 __snake_case	:     Union[str, Any]			   =      (
     (
         TFMPNetForMaskedLM,
         TFMPNetForMultipleChoice,
         TFMPNetForQuestionAnswering,
         TFMPNetForSequenceClassification,
         TFMPNetForTokenClassification,
         TFMPNetModel,
     )
     if is_tf_available()
     else ()
 )
 __snake_case	:     str			   =      (
     {
         "feature-extraction": TFMPNetModel,
         "fill-mask": TFMPNetForMaskedLM,
         "question-answering": TFMPNetForQuestionAnswering,
         "text-classification": TFMPNetForSequenceClassification,
         "token-classification": TFMPNetForTokenClassification,
         "zero-shot": TFMPNetForSequenceClassification,
     }
     if is_tf_available()
     else {}
 )
 __snake_case	:     str			   =      False
 __snake_case	:     Optional[Any]			   =      False



 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        lowercase  :		Optional[int]						=      TFMPNetModelTester(self )
        lowercase  :		List[Any]						=      ConfigTester(self   ,		config_class=A   ,		hidden_size=37 )



 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        self.config_tester.run_common_tests()



 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        lowercase  :		Optional[int]						=      self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mpnet_model(*A )



 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        lowercase  :		Any						=      self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mpnet_for_masked_lm(*A )



 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        lowercase  :		Optional[int]						=      self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mpnet_for_question_answering(*A )



 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        lowercase  :		Any						=      self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mpnet_for_sequence_classification(*A )



 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        lowercase  :		Dict						=      self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mpnet_for_multiple_choice(*A )



 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        lowercase  :		Union[str, Any]						=      self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mpnet_for_token_classification(*A )





 @slow
 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        for model_name in ["microsoft/mpnet-base"]:
               lowercase  :		Union[str, Any]						=      TFMPNetModel.from_pretrained(A )
               self.assertIsNotNone(A )






@require_tf
class   SCREAMING_SNAKE_CASE_		(       unittest.TestCase				):



 @slow
 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        lowercase  :		Union[str, Any]						=      TFMPNetModel.from_pretrained('''microsoft/mpnet-base''' )
        lowercase  :		Optional[Any]						=      tf.constant([[0, 1, 2, 3, 4, 5]] )
        lowercase  :		Optional[Any]						=      model(A )[0]

        lowercase  :		Optional[Any]						=      [1, 6, 768]
        self.assertEqual(output.shape   ,		A )

        lowercase  :		Any						=      tf.constant(
            [
                [
                    [-0.1_0_6_7_1_7_2, 0.0_8_2_1_6_4_7_3, 0.0_0_2_4_5_4_3],
                    [-0.0_3_4_6_5_8_7_9, 0.8_3_5_4_1_1_8, -0.0_3_2_5_2_2_8_8],
                    [-0.0_6_5_6_9_4_7_6, -0.1_2_4_2_4_1_1_1, -0.0_4_9_4_4_3_6],
                ]
            ] )
        tf.debugging.assert_near(output[:, :3, :3]   ,		A   ,		atol=1E-4 )










'''simple docstring'''






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

from ...utils import logging


lowercase	:      Optional[Any]				=       logging.get_logger(__name__)




def 		a_( lowerCamelCase   ,				lowerCamelCase      )				->Optional[int]:






       '''simple docstring'''




       lowercase  :		List[str]						=      nn.functional.normalize(lowerCamelCase      )
       lowercase  :		Optional[int]						=      nn.functional.normalize(lowerCamelCase      )
       return torch.mm(lowerCamelCase   ,				normalized_text_embeds.t()      )





class   SCREAMING_SNAKE_CASE_		(       lowerCamelCase_				):
 __snake_case	:     List[str]			   =      CLIPConfig

 __snake_case	:     Optional[int]			   =      ["CLIPEncoderLayer"]



 def __init__(  self   ,		A ):






        '''simple docstring'''




        super().__init__(A )

        lowercase  :		Optional[Any]						=      CLIPVisionModel(config.vision_config )
        lowercase  :		Tuple						=      nn.Linear(config.vision_config.hidden_size   ,		config.projection_dim   ,		bias=A )

        lowercase  :		Any						=      nn.Parameter(torch.ones(17   ,		config.projection_dim )   ,		requires_grad=A )
        lowercase  :		int						=      nn.Parameter(torch.ones(3   ,		config.projection_dim )   ,		requires_grad=A )

        lowercase  :		List[str]						=      nn.Parameter(torch.ones(17 )   ,		requires_grad=A )
        lowercase  :		str						=      nn.Parameter(torch.ones(3 )   ,		requires_grad=A )



 @torch.no_grad()
 def    SCREAMING_SNAKE_CASE(  self   ,		A   ,		A ):






        '''simple docstring'''




        lowercase  :		Union[str, Any]						=      self.vision_model(A )[1]  # pooled_output
        lowercase  :		List[Any]						=      self.visual_projection(A )

        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
        lowercase  :		Optional[Any]						=      cosine_distance(A   ,		self.special_care_embeds ).cpu().float().numpy()
        lowercase  :		int						=      cosine_distance(A   ,		self.concept_embeds ).cpu().float().numpy()

        lowercase  :		Dict						=      []
        lowercase  :		str						=      image_embeds.shape[0]
        for i in range(A ):
               lowercase  :		List[str]						=      {'''special_scores''': {}, '''special_care''': [], '''concept_scores''': {}, '''bad_concepts''': []}

               # increase this value to create a stronger `nfsw` filter
               # at the cost of increasing the possibility of filtering benign images
               lowercase  :		str						=      0.0

               for concept_idx in range(len(special_cos_dist[0] ) ):
                      lowercase  :		List[str]						=      special_cos_dist[i][concept_idx]
                      lowercase  :		List[Any]						=      self.special_care_embeds_weights[concept_idx].item()
                      lowercase  :		Optional[Any]						=      round(concept_cos - concept_threshold + adjustment   ,		3 )
                      if result_img["special_scores"][concept_idx] > 0:
                             result_img["special_care"].append({concept_idx, result_img['''special_scores'''][concept_idx]} )
                             lowercase  :		List[str]						=      0.0_1

               for concept_idx in range(len(cos_dist[0] ) ):
                      lowercase  :		List[Any]						=      cos_dist[i][concept_idx]
                      lowercase  :		Optional[Any]						=      self.concept_embeds_weights[concept_idx].item()
                      lowercase  :		Any						=      round(concept_cos - concept_threshold + adjustment   ,		3 )
                      if result_img["concept_scores"][concept_idx] > 0:
                             result_img["bad_concepts"].append(A )

               result.append(A )

        lowercase  :		Dict						=      [len(res['''bad_concepts'''] ) > 0 for res in result]

        for idx, has_nsfw_concept in enumerate(A ):
               if has_nsfw_concept:
                      if torch.is_tensor(A ) or torch.is_tensor(images[0] ):
                             lowercase  :		Any						=      torch.zeros_like(images[idx] )  # black image
                      else:
                             lowercase  :		Tuple						=      np.zeros(images[idx].shape )  # black image

        if any(A ):
               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.''' )

        return images, has_nsfw_concepts





 @torch.no_grad()
 def    SCREAMING_SNAKE_CASE(  self   ,		A   ,		A ):






        '''simple docstring'''




        lowercase  :		Dict						=      self.vision_model(A )[1]  # pooled_output
        lowercase  :		Optional[int]						=      self.visual_projection(A )

        lowercase  :		Tuple						=      cosine_distance(A   ,		self.special_care_embeds )
        lowercase  :		List[str]						=      cosine_distance(A   ,		self.concept_embeds )

        # increase this value to create a stronger `nsfw` filter
        # at the cost of increasing the possibility of filtering benign images
        lowercase  :		List[str]						=      0.0

        lowercase  :		Optional[Any]						=      special_cos_dist - self.special_care_embeds_weights + adjustment
        # special_scores = special_scores.round(decimals=3)
        lowercase  :		Union[str, Any]						=      torch.any(special_scores > 0   ,		dim=1 )
        lowercase  :		Union[str, Any]						=      special_care * 0.0_1
        lowercase  :		Any						=      special_adjustment.unsqueeze(1 ).expand(-1   ,		cos_dist.shape[1] )

        lowercase  :		Union[str, Any]						=      (cos_dist - self.concept_embeds_weights) + special_adjustment
        # concept_scores = concept_scores.round(decimals=3)
        lowercase  :		Tuple						=      torch.any(concept_scores > 0   ,		dim=1 )

        lowercase  :		Dict						=      0.0  # black image

        return images, has_nsfw_concepts




