





'''simple docstring'''

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

from ...utils import logging


lowercase__						:      Union[str, Any]	=    logging.get_logger(__name__)







def  a__    (    lowercase :					Optional[int], lowercase :					int				)		->		Tuple:





     """simple docstring"""

     _UpperCamelCase       =	nn.functional.normalize(lowercase				)
     _UpperCamelCase       =	nn.functional.normalize(lowercase				)
     return torch.mm(lowercase, normalized_text_embeds.t()				)




class     __lowerCAmelCase					(   __magic_name__			):






       """simple docstring"""
       _snake_case     :					Union[str, Any]						    =   CLIPConfig

       _snake_case     :					Optional[Any]						    =   ['CLIPEncoderLayer']





       def __init__(							self       :		Union[str, Any]	,			lowerCAmelCase__       :		CLIPConfig					)   ->					Union[str, Any]:

            '''simple docstring'''




            super().__init__(lowerCAmelCase__					)

            _UpperCamelCase       =	CLIPVisionModel(config.vision_config					)
            _UpperCamelCase       =	nn.Linear(config.vision_config.hidden_size	,			config.projection_dim	,			bias=lowerCAmelCase__					)

            _UpperCamelCase       =	nn.Parameter(torch.ones(17	,			config.projection_dim					)	,			requires_grad=lowerCAmelCase__					)
            _UpperCamelCase       =	nn.Parameter(torch.ones(3	,			config.projection_dim					)	,			requires_grad=lowerCAmelCase__					)

            _UpperCamelCase       =	nn.Parameter(torch.ones(17					)	,			requires_grad=lowerCAmelCase__					)
            _UpperCamelCase       =	nn.Parameter(torch.ones(3					)	,			requires_grad=lowerCAmelCase__					)





       @torch.no_grad()
       def snake_case__						(							self       :		str	,			lowerCAmelCase__       :		Union[str, Any]	,			lowerCAmelCase__       :		int					)   ->					Tuple:

            '''simple docstring'''




            _UpperCamelCase       =	self.vision_model(lowerCAmelCase__					)[1]  # pooled_output
            _UpperCamelCase       =	self.visual_projection(lowerCAmelCase__					)

            # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
            _UpperCamelCase       =	cosine_distance(lowerCAmelCase__	,			self.special_care_embeds					).cpu().float().numpy()
            _UpperCamelCase       =	cosine_distance(lowerCAmelCase__	,			self.concept_embeds					).cpu().float().numpy()

            _UpperCamelCase       =	[]
            _UpperCamelCase       =	image_embeds.shape[0]
            for i in range(lowerCAmelCase__					):
                 _UpperCamelCase       =	{'''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
                 _UpperCamelCase       =	0.0

                 for concept_idx in range(len(special_cos_dist[0]					)					):
                      _UpperCamelCase       =	special_cos_dist[i][concept_idx]
                      _UpperCamelCase       =	self.special_care_embeds_weights[concept_idx].item()
                      _UpperCamelCase       =	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]}					)
                           _UpperCamelCase       =	0.01

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

                 result.append(lowerCAmelCase__					)

            _UpperCamelCase       =	[len(res['''bad_concepts''']					) > 0 for res in result]

            return images, has_nsfw_concepts







       @torch.no_grad()
       def snake_case__						(							self       :		List[Any]	,			lowerCAmelCase__       :		torch.FloatTensor	,			lowerCAmelCase__       :		torch.FloatTensor					)   ->					Union[str, Any]:

            '''simple docstring'''




            _UpperCamelCase       =	self.vision_model(lowerCAmelCase__					)[1]  # pooled_output
            _UpperCamelCase       =	self.visual_projection(lowerCAmelCase__					)

            _UpperCamelCase       =	cosine_distance(lowerCAmelCase__	,			self.special_care_embeds					)
            _UpperCamelCase       =	cosine_distance(lowerCAmelCase__	,			self.concept_embeds					)

            # increase this value to create a stronger `nsfw` filter
            # at the cost of increasing the possibility of filtering benign images
            _UpperCamelCase       =	0.0

            _UpperCamelCase       =	special_cos_dist - self.special_care_embeds_weights + adjustment
            # special_scores = special_scores.round(decimals=3)
            _UpperCamelCase       =	torch.any(special_scores > 0	,			dim=1					)
            _UpperCamelCase       =	special_care * 0.01
            _UpperCamelCase       =	special_adjustment.unsqueeze(1					).expand(-1	,			cos_dist.shape[1]					)

            _UpperCamelCase       =	(cos_dist - self.concept_embeds_weights) + special_adjustment
            # concept_scores = concept_scores.round(decimals=3)
            _UpperCamelCase       =	torch.any(concept_scores > 0	,			dim=1					)

            return images, has_nsfw_concepts










'''simple docstring'''

import unittest

import numpy as np

from transformers import RobertaConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
   from transformers.models.roberta.modeling_flax_roberta import (
       FlaxRobertaForCausalLM,
       FlaxRobertaForMaskedLM,
       FlaxRobertaForMultipleChoice,
       FlaxRobertaForQuestionAnswering,
       FlaxRobertaForSequenceClassification,
       FlaxRobertaForTokenClassification,
       FlaxRobertaModel,
   )
class     __lowerCAmelCase					(   unittest.TestCase			):






       """simple docstring"""





       def __init__(							self       :		Tuple	,			lowerCAmelCase__       :		Tuple	,			lowerCAmelCase__       :		Union[str, Any]=13	,			lowerCAmelCase__       :		Union[str, Any]=7	,			lowerCAmelCase__       :		str=True	,			lowerCAmelCase__       :		Tuple=True	,			lowerCAmelCase__       :		Dict=True	,			lowerCAmelCase__       :		Dict=True	,			lowerCAmelCase__       :		int=99	,			lowerCAmelCase__       :		str=32	,			lowerCAmelCase__       :		str=5	,			lowerCAmelCase__       :		str=4	,			lowerCAmelCase__       :		str=37	,			lowerCAmelCase__       :		int="gelu"	,			lowerCAmelCase__       :		Optional[Any]=0.1	,			lowerCAmelCase__       :		int=0.1	,			lowerCAmelCase__       :		Optional[int]=512	,			lowerCAmelCase__       :		Dict=16	,			lowerCAmelCase__       :		List[Any]=2	,			lowerCAmelCase__       :		Any=0.02	,			lowerCAmelCase__       :		Union[str, Any]=4	,			)   ->					Dict:

            '''simple docstring'''




            _UpperCamelCase       =	parent
            _UpperCamelCase       =	batch_size
            _UpperCamelCase       =	seq_length
            _UpperCamelCase       =	is_training
            _UpperCamelCase       =	use_attention_mask
            _UpperCamelCase       =	use_token_type_ids
            _UpperCamelCase       =	use_labels
            _UpperCamelCase       =	vocab_size
            _UpperCamelCase       =	hidden_size
            _UpperCamelCase       =	num_hidden_layers
            _UpperCamelCase       =	num_attention_heads
            _UpperCamelCase       =	intermediate_size
            _UpperCamelCase       =	hidden_act
            _UpperCamelCase       =	hidden_dropout_prob
            _UpperCamelCase       =	attention_probs_dropout_prob
            _UpperCamelCase       =	max_position_embeddings
            _UpperCamelCase       =	type_vocab_size
            _UpperCamelCase       =	type_sequence_label_size
            _UpperCamelCase       =	initializer_range
            _UpperCamelCase       =	num_choices





       def snake_case__						(							self       :		Optional[int]					)   ->					Optional[Any]:

            '''simple docstring'''




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

            _UpperCamelCase       =	None
            if self.use_attention_mask:
                 _UpperCamelCase       =	random_attention_mask([self.batch_size, self.seq_length]					)

            _UpperCamelCase       =	None
            if self.use_token_type_ids:
                 _UpperCamelCase       =	ids_tensor([self.batch_size, self.seq_length]	,			self.type_vocab_size					)

            _UpperCamelCase       =	RobertaConfig(
                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	,			)

            return config, input_ids, token_type_ids, attention_mask





       def snake_case__						(							self       :		Optional[int]					)   ->					Any:

            '''simple docstring'''




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







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

            '''simple docstring'''




            _UpperCamelCase       =	self.prepare_config_and_inputs()
            _UpperCamelCase   ,    _UpperCamelCase   ,    _UpperCamelCase   ,    _UpperCamelCase       =	config_and_inputs

            _UpperCamelCase       =	True
            _UpperCamelCase       =	floats_tensor([self.batch_size, self.seq_length, self.hidden_size]					)
            _UpperCamelCase       =	ids_tensor([self.batch_size, self.seq_length]	,			vocab_size=2					)

            return (
                config,
                input_ids,
                token_type_ids,
                encoder_hidden_states,
                encoder_attention_mask,
            )

@require_flax
class     __lowerCAmelCase					(   __magic_name__						,						unittest.TestCase			):






       """simple docstring"""
       _snake_case     :					Optional[int]						    =   True

       _snake_case     :					Optional[Any]						    =   (
           (
               FlaxRobertaModel,
               FlaxRobertaForCausalLM,
               FlaxRobertaForMaskedLM,
               FlaxRobertaForSequenceClassification,
               FlaxRobertaForTokenClassification,
               FlaxRobertaForMultipleChoice,
               FlaxRobertaForQuestionAnswering,
           )
           if is_flax_available()
           else ()
       )





       def snake_case__						(							self       :		Dict					)   ->					List[Any]:

            '''simple docstring'''




            _UpperCamelCase       =	FlaxRobertaModelTester(self					)







       @slow
       def snake_case__						(							self       :		List[str]					)   ->					Union[str, Any]:

            '''simple docstring'''




            for model_class_name in self.all_model_classes:
                 _UpperCamelCase       =	model_class_name.from_pretrained('''roberta-base'''	,			from_pt=lowerCAmelCase__					)
                 _UpperCamelCase       =	model(np.ones((1, 1)					)					)
                 self.assertIsNotNone(lowerCAmelCase__					)



