"""simple docstring"""






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

from ...utils import logging


lowerCAmelCase__						:						Dict		  =			logging.get_logger(__name__)



def 						a_		(      lowerCamelCase      ,    lowerCamelCase ):
				UpperCAmelCase__    =							nn.functional.normalize(lowerCamelCase )
				UpperCAmelCase__    =							nn.functional.normalize(lowerCamelCase )
				return torch.mm(lowerCamelCase      ,    normalized_text_embeds.t() )







class        snake_case (				__UpperCAmelCase       ):




					"""simple docstring"""

					snake_case__       					=				CLIPConfig

					snake_case__       					=				["CLIPEncoderLayer"]







					def __init__(       self     :			str					,lowerCamelCase__     :			CLIPConfig    ):
									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 	__lowerCAmelCase		(       self     :			List[str]					,lowerCamelCase__     :			Dict					,lowerCamelCase__     :			Tuple    ):
									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.0_1

													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 	__lowerCAmelCase		(       self     :			str					,lowerCamelCase__     :			torch.FloatTensor					,lowerCamelCase__     :			torch.FloatTensor    ):
									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.0_1
									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"""






def 						a_		(      lowerCamelCase      ,    lowerCamelCase ):
				return x if y == 0 else greatest_common_divisor(lowerCamelCase      ,    x % y )



def 						a_		(      lowerCamelCase      ,    lowerCamelCase ):
				return (x * y) // greatest_common_divisor(lowerCamelCase      ,    lowerCamelCase )



def 						a_		(      lowerCamelCase = 2_0 ):
				UpperCAmelCase__    =							1
				for i in range(1      ,    n + 1 ):
								UpperCAmelCase__    =							lcm(lowerCamelCase      ,    lowerCamelCase )
				return g


if __name__ == "__main__":
						print(F"""{solution() = }""")



