




'''simple docstring'''





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

from ...utils import logging


_SCREAMING_SNAKE_CASE				:  str 						=					logging.get_logger(__name__)


def 	UpperCamelCase_(      snake_case				:			Any				,							snake_case				:			Dict	):

							'''simple docstring'''






							snake_case_         = nn.functional.normalize(snake_case	)
							snake_case_         = nn.functional.normalize(snake_case	)
							return torch.mm(snake_case				,							normalized_text_embeds.t()	)




class 					_snake_case						(      _UpperCAmelCase				):
			lowerCAmelCase_	: List[Any]					=       CLIPConfig

			lowerCAmelCase_	: Optional[Any]					=       ["CLIPEncoderLayer"]




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

										'''simple docstring'''




										super().__init__(a__				)

										snake_case_         = CLIPVisionModel(config.vision_config				)
										snake_case_         = nn.Linear(config.vision_config.hidden_size							,				config.projection_dim							,				bias=a__				)

										snake_case_         = nn.Parameter(torch.ones(17							,				config.projection_dim				)							,				requires_grad=a__				)
										snake_case_         = nn.Parameter(torch.ones(3							,				config.projection_dim				)							,				requires_grad=a__				)

										snake_case_         = nn.Parameter(torch.ones(17				)							,				requires_grad=a__				)
										snake_case_         = nn.Parameter(torch.ones(3				)							,				requires_grad=a__				)




			@torch.no_grad()
			def   lowerCAmelCase__							(   self							,				a__							,				a__				)							->							List[str]:

										'''simple docstring'''




										snake_case_         = self.vision_model(a__				)[1]  # pooled_output
										snake_case_         = self.visual_projection(a__				)

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

										snake_case_         = []
										snake_case_         = image_embeds.shape[0]
										for i in range(a__				):
																	snake_case_         = {"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_         = 0.0

																	for concept_idx in range(len(special_cos_dist[0]				)				):
																								snake_case_         = special_cos_dist[i][concept_idx]
																								snake_case_         = self.special_care_embeds_weights[concept_idx].item()
																								snake_case_         = 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_         = 0.0_1

																	for concept_idx in range(len(cos_dist[0]				)				):
																								snake_case_         = cos_dist[i][concept_idx]
																								snake_case_         = self.concept_embeds_weights[concept_idx].item()
																								snake_case_         = 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_         = [len(res["bad_concepts"]				) > 0 for res in result]

										return images, has_nsfw_concepts




			@torch.no_grad()
			def   lowerCAmelCase__							(   self							,				a__							,				a__				)							->							int:

										'''simple docstring'''




										snake_case_         = self.vision_model(a__				)[1]  # pooled_output
										snake_case_         = self.visual_projection(a__				)

										snake_case_         = cosine_distance(a__							,				self.special_care_embeds				)
										snake_case_         = 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_         = 0.0

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

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

										return images, has_nsfw_concepts











"""simple docstring"""





def      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		int				,		SCREAMING_SNAKE_CASE	:		int     ):







      '''simple docstring'''




      if a < 0 or b < 0:
            raise ValueError("""the value of both inputs must be positive"""     )
      lowerCAmelCase    		=			str(bin(SCREAMING_SNAKE_CASE     )     )[2:]  # remove the leading "0b"
      lowerCAmelCase    		=			str(bin(SCREAMING_SNAKE_CASE     )     )[2:]
      lowerCAmelCase    		=			max(len(SCREAMING_SNAKE_CASE     )				,		len(SCREAMING_SNAKE_CASE     )     )
      return "0b" + "".join(
          str(int("""1""" in (char_a, char_b)     )     )
          for char_a, char_b in zip(a_binary.zfill(SCREAMING_SNAKE_CASE     )				,		b_binary.zfill(SCREAMING_SNAKE_CASE     )     )     )


if __name__ == "__main__":
       import doctest

       doctest.testmod()






