





'''simple docstring'''




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

from ...utils import logging


lowercase__   :				Optional[int]      =   logging.get_logger(__name__)






def 	_lowerCAmelCase      (  __snake_case  :			str			,							__snake_case  :			Optional[Any]      )					->					str:
     __A						:			str								=	nn.functional.normalize(__snake_case      )
     __A						:			Optional[int]								=	nn.functional.normalize(__snake_case      )
     return torch.mm(__snake_case			,							normalized_text_embeds.t()      )







class 			SCREAMING_SNAKE_CASE   (a__							):
     lowerCAmelCase		      =	CLIPConfig

     lowerCAmelCase		      =	['''CLIPEncoderLayer''']


     def __init__(						self						,     _UpperCAmelCase):

          '''simple docstring'''






          super().__init__(_UpperCAmelCase)

          __A						:			List[str]								=	CLIPVisionModel(config.vision_config)
          __A						:			Optional[Any]								=	nn.Linear(config.vision_config.hidden_size						,     config.projection_dim						,     bias=_UpperCAmelCase)

          __A						:			Optional[int]								=	nn.Parameter(torch.ones(17						,     config.projection_dim)						,     requires_grad=_UpperCAmelCase)
          __A						:			str								=	nn.Parameter(torch.ones(3						,     config.projection_dim)						,     requires_grad=_UpperCAmelCase)

          __A						:			Dict								=	nn.Parameter(torch.ones(17)						,     requires_grad=_UpperCAmelCase)
          __A						:			str								=	nn.Parameter(torch.ones(3)						,     requires_grad=_UpperCAmelCase)


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

          '''simple docstring'''






          __A						:			Dict								=	self.vision_model(_UpperCAmelCase)[1]  # pooled_output
          __A						:			Optional[Any]								=	self.visual_projection(_UpperCAmelCase)

          # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
          __A						:			Union[str, Any]								=	cosine_distance(_UpperCAmelCase						,     self.special_care_embeds).cpu().float().numpy()
          __A						:			List[Any]								=	cosine_distance(_UpperCAmelCase						,     self.concept_embeds).cpu().float().numpy()

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

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

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

               result.append(_UpperCAmelCase)

          __A						:			Optional[Any]								=	[len(res['bad_concepts']) > 0 for res in result]

          return images, has_nsfw_concepts


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

          '''simple docstring'''






          __A						:			Optional[Any]								=	self.vision_model(_UpperCAmelCase)[1]  # pooled_output
          __A						:			Tuple								=	self.visual_projection(_UpperCAmelCase)

          __A						:			Union[str, Any]								=	cosine_distance(_UpperCAmelCase						,     self.special_care_embeds)
          __A						:			Tuple								=	cosine_distance(_UpperCAmelCase						,     self.concept_embeds)

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

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

          __A						:			Optional[Any]								=	(cos_dist - self.concept_embeds_weights) + special_adjustment
          # concept_scores = concept_scores.round(decimals=3)
          __A						:			Optional[int]								=	torch.any(concept_scores > 0						,     dim=1)

          return images, has_nsfw_concepts






'''simple docstring'''




from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding


class 			SCREAMING_SNAKE_CASE   (a__							):
     lowerCAmelCase		      =	['''image_processor''', '''tokenizer''']
     lowerCAmelCase		      =	'''AutoImageProcessor'''
     lowerCAmelCase		      =	'''AutoTokenizer'''


     def __init__(						self						,     _UpperCAmelCase						,     _UpperCAmelCase):

          '''simple docstring'''






          super().__init__(_UpperCAmelCase						,     _UpperCAmelCase)
          __A						:			Tuple								=	self.image_processor


     def __call__(						self						,     _UpperCAmelCase=None						,     _UpperCAmelCase=None						,     _UpperCAmelCase=None						,     **_UpperCAmelCase):

          '''simple docstring'''






          if text is None and images is None:
               raise ValueError('You have to specify either text or images. Both cannot be none.')

          if text is not None:
               __A						:			Any								=	self.tokenizer(_UpperCAmelCase						,     return_tensors=_UpperCAmelCase						,     **_UpperCAmelCase)

          if images is not None:
               __A						:			Tuple								=	self.image_processor(_UpperCAmelCase						,     return_tensors=_UpperCAmelCase						,     **_UpperCAmelCase)

          if text is not None and images is not None:
               __A						:			Optional[int]								=	image_features.pixel_values
               return encoding
          elif text is not None:
               return encoding
          else:
               return BatchEncoding(data=dict(**_UpperCAmelCase)						,     tensor_type=_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self						,     *_UpperCAmelCase						,     **_UpperCAmelCase):

          '''simple docstring'''






          return self.tokenizer.batch_decode(*_UpperCAmelCase						,     **_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self						,     *_UpperCAmelCase						,     **_UpperCAmelCase):

          '''simple docstring'''






          return self.tokenizer.decode(*_UpperCAmelCase						,     **_UpperCAmelCase)


     @property
     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          return ["input_ids", "attention_mask", "pixel_values"]