



"""simple docstring"""







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

from ...utils import logging


A		=						logging.get_logger(__name__)

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							Any     ,		lowerCamelCase_:							Optional[Any] ):
						"""simple docstring"""






						snake_case   :			int						    = nn.functional.normalize(lowerCamelCase_ )
						snake_case   :			List[str]						    = nn.functional.normalize(lowerCamelCase_ )
						return torch.mm(lowerCamelCase_     ,		normalized_text_embeds.t() )

class       _a  ( SCREAMING_SNAKE_CASE__):
			__magic_name__     					=     CLIPConfig

			__magic_name__     					=     ["""CLIPEncoderLayer"""]



			def __init__(			self							:				int   ,	_lowercase							:				CLIPConfig		)							->					Tuple:
									super().__init__(_lowercase		)

									snake_case   :			Optional[int]						    = CLIPVisionModel(config.vision_config		)
									snake_case   :			Dict						    = nn.Linear(config.vision_config.hidden_size   ,	config.projection_dim   ,	bias=_lowercase		)

									snake_case   :			int						    = nn.Parameter(torch.ones(17   ,	config.projection_dim		)   ,	requires_grad=_lowercase		)
									snake_case   :			Optional[int]						    = nn.Parameter(torch.ones(3   ,	config.projection_dim		)   ,	requires_grad=_lowercase		)

									snake_case   :			int						    = nn.Parameter(torch.ones(17		)   ,	requires_grad=_lowercase		)
									snake_case   :			Tuple						    = nn.Parameter(torch.ones(3		)   ,	requires_grad=_lowercase		)



			@torch.no_grad()
			def 				__lowercase					(			self							:				int   ,	_lowercase							:				List[Any]   ,	_lowercase							:				List[Any]		)							->					int:
									snake_case   :			Optional[int]						    = self.vision_model(_lowercase		)[1]  # pooled_output
									snake_case   :			Optional[Any]						    = self.visual_projection(_lowercase		)

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

									snake_case   :			Tuple						    = []
									snake_case   :			Union[str, Any]						    = image_embeds.shape[0]
									for i in range(_lowercase		):
															snake_case   :			Tuple						    = {"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   :			int						    = 0.0

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

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

															result.append(_lowercase		)

									snake_case   :			str						    = [len(res["bad_concepts"]		) > 0 for res in result]

									return images, has_nsfw_concepts



			@torch.no_grad()
			def 				__lowercase					(			self							:				Optional[int]   ,	_lowercase							:				torch.FloatTensor   ,	_lowercase							:				torch.FloatTensor		)							->					Optional[int]:
									snake_case   :			str						    = self.vision_model(_lowercase		)[1]  # pooled_output
									snake_case   :			List[str]						    = self.visual_projection(_lowercase		)

									snake_case   :			Dict						    = cosine_distance(_lowercase   ,	self.special_care_embeds		)
									snake_case   :			Tuple						    = cosine_distance(_lowercase   ,	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   :			Optional[Any]						    = 0.0

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

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

									return images, has_nsfw_concepts








"""simple docstring"""







# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import re

from ..utils import cached_file


# docstyle-ignore
A		=						'\nHuman: <<task>>\n\nAssistant: '


A		=						'huggingface-tools/default-prompts'
A		=						{'chat': 'chat_prompt_template.txt', 'run': 'run_prompt_template.txt'}

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							Optional[int]     ,		lowerCamelCase_:							str     ,		lowerCamelCase_:							Tuple="run" ):
						"""simple docstring"""






						if prompt_or_repo_id is None:
												snake_case   :			Any						    = DEFAULT_PROMPTS_REPO

						# prompt is considered a repo ID when it does not contain any kind of space
						if re.search("\\s"     ,		lowerCamelCase_ ) is not None:
												return prompt_or_repo_id

						snake_case   :			Optional[int]						    = cached_file(
						    lowerCamelCase_     ,		PROMPT_FILES[mode]     ,		repo_type="dataset"     ,		user_agent={"agent": agent_name} )
						with open(lowerCamelCase_     ,		"r"     ,		encoding="utf-8" ) as f:
												return f.read()



