





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

from ...utils import logging


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







def 			A__							(    SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__)   ->    List[str]:
  __snake_case:						str   =						nn.functional.normalize(SCREAMING_SNAKE_CASE__)
  __snake_case:						Dict   =						nn.functional.normalize(SCREAMING_SNAKE_CASE__)
  return torch.mm(SCREAMING_SNAKE_CASE__					,					normalized_text_embeds.t())







class     __snake_case       (							__lowerCamelCase			):





       '''simple docstring'''







       lowerCAmelCase__      		=     CLIPConfig

       lowerCAmelCase__      		=     ["""CLIPEncoderLayer"""]




       def __init__(  self :    Tuple							,				A :    CLIPConfig		):
         super().__init__(A		)

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

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

         __snake_case:						Tuple   =						nn.Parameter(torch.ones(17		)							,				requires_grad=A		)
         __snake_case:						Union[str, Any]   =						nn.Parameter(torch.ones(3		)							,				requires_grad=A		)




       @torch.no_grad()
       def 							UpperCAmelCase__ (  self :    List[str]							,				A :    Tuple							,				A :    List[str]		):
         __snake_case:						Optional[Any]   =						self.vision_model(A		)[1]  # pooled_output
         __snake_case:						List[Any]   =						self.visual_projection(A		)

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

         __snake_case:						Tuple   =						[]
         __snake_case:						str   =						image_embeds.shape[0]
         for i in range(A		):
           __snake_case:						Union[str, 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
           __snake_case:						List[str]   =						0.0

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

           for concept_idx in range(len(cos_dist[0]		)		):
             __snake_case:						int   =						cos_dist[i][concept_idx]
             __snake_case:						Optional[Any]   =						self.concept_embeds_weights[concept_idx].item()
             __snake_case:						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		)

         __snake_case:						Optional[int]   =						[len(res["""bad_concepts"""]		) > 0 for res in result]

         return images, has_nsfw_concepts






       @torch.no_grad()
       def 							UpperCAmelCase__ (  self :    Optional[Any]							,				A :    torch.FloatTensor							,				A :    torch.FloatTensor		):
         __snake_case:						Optional[Any]   =						self.vision_model(A		)[1]  # pooled_output
         __snake_case:						str   =						self.visual_projection(A		)

         __snake_case:						Tuple   =						cosine_distance(A							,				self.special_care_embeds		)
         __snake_case:						Union[str, Any]   =						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
         __snake_case:						Dict   =						0.0

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

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

         return images, has_nsfw_concepts













import unittest

from transformers import load_tool

from .test_tools_common import ToolTesterMixin







class     __snake_case       (							unittest.TestCase   ,			__lowerCamelCase			):





       '''simple docstring'''







       def 							UpperCAmelCase__ (  self :    Union[str, Any]		):
         __snake_case:						Optional[int]   =						load_tool("""text-classification"""		)
         self.tool.setup()
         __snake_case:						Dict   =						load_tool("""text-classification"""							,				remote=A		)




       def 							UpperCAmelCase__ (  self :    Optional[int]		):
         __snake_case:						List[str]   =						self.tool("""That's quite cool"""							,				["""positive""", """negative"""]		)
         self.assertEqual(A							,				"""positive"""		)




       def 							UpperCAmelCase__ (  self :    Dict		):
         __snake_case:						List[str]   =						self.remote_tool("""That's quite cool"""							,				["""positive""", """negative"""]		)
         self.assertEqual(A							,				"""positive"""		)




       def 							UpperCAmelCase__ (  self :    Optional[Any]		):
         __snake_case:						Optional[Any]   =						self.tool(text="""That's quite cool"""							,				labels=["""positive""", """negative"""]		)
         self.assertEqual(A							,				"""positive"""		)






       def 							UpperCAmelCase__ (  self :    Tuple		):
         __snake_case:						int   =						self.remote_tool(text="""That's quite cool"""							,				labels=["""positive""", """negative"""]		)
         self.assertEqual(A							,				"""positive"""		)






