



'''simple docstring'''


from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_speech_available,
    is_torch_available,
)


_lowerCamelCase	:       Tuple      	=	{
    "configuration_trocr": ["TROCR_PRETRAINED_CONFIG_ARCHIVE_MAP", "TrOCRConfig"],
    "processing_trocr": ["TrOCRProcessor"],
}


try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _lowerCamelCase	:       Any      	=	[
       "TROCR_PRETRAINED_MODEL_ARCHIVE_LIST",
       "TrOCRForCausalLM",
       "TrOCRPreTrainedModel",
   ]


if TYPE_CHECKING:
   from .configuration_trocr import TROCR_PRETRAINED_CONFIG_ARCHIVE_MAP, TrOCRConfig
   from .processing_trocr import TrOCRProcessor

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_trocr import TROCR_PRETRAINED_MODEL_ARCHIVE_LIST, TrOCRForCausalLM, TrOCRPreTrainedModel

else:
   import sys

   _lowerCamelCase	:       int      	=	_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)





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

from ...utils import logging


__lowerCAmelCase	: Dict  				=logging.get_logger(__name__)


def        _UpperCamelCase							(      lowercase__    ,       lowercase__      ):
						__SCREAMING_SNAKE_CASE						:     List[Any]					=		nn.functional.normalize(lowercase__      )
						__SCREAMING_SNAKE_CASE						:     Tuple					=		nn.functional.normalize(lowercase__      )
						return torch.mm(lowercase__    ,       normalized_text_embeds.t()      )
class        _lowercase      (			A__       ):






				'''simple docstring'''



				SCREAMING_SNAKE_CASE__				:			Optional[Any]						=		CLIPConfig

				SCREAMING_SNAKE_CASE__				:			List[str]						=		['''CLIPEncoderLayer''']


				def __init__(       self			:str			,			lowerCAmelCase__			:CLIPConfig	)		->			Tuple:
										super().__init__(lowerCAmelCase__	)

										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		CLIPVisionModel(config.vision_config	)
										__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		nn.Linear(config.vision_config.hidden_size			,			config.projection_dim			,			bias=lowerCAmelCase__	)

										__SCREAMING_SNAKE_CASE						:     Any					=		nn.Parameter(torch.ones(17			,			config.projection_dim	)			,			requires_grad=lowerCAmelCase__	)
										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		nn.Parameter(torch.ones(3			,			config.projection_dim	)			,			requires_grad=lowerCAmelCase__	)

										__SCREAMING_SNAKE_CASE						:     Dict					=		nn.Parameter(torch.ones(17	)			,			requires_grad=lowerCAmelCase__	)
										__SCREAMING_SNAKE_CASE						:     str					=		nn.Parameter(torch.ones(3	)			,			requires_grad=lowerCAmelCase__	)


				@torch.no_grad()
				def 		__magic_name__(       self			:Optional[int]			,			lowerCAmelCase__			:int			,			lowerCAmelCase__			:Optional[int]	)		->			Dict:
										__SCREAMING_SNAKE_CASE						:     int					=		self.vision_model(lowerCAmelCase__	)[1]  # pooled_output
										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		self.visual_projection(lowerCAmelCase__	)

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

										__SCREAMING_SNAKE_CASE						:     List[Any]					=		[]
										__SCREAMING_SNAKE_CASE						:     List[Any]					=		image_embeds.shape[0]
										for i in range(lowerCAmelCase__	):
																__SCREAMING_SNAKE_CASE						:     Optional[int]					=		{'''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
																__SCREAMING_SNAKE_CASE						:     List[Any]					=		0.0

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

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

																result.append(lowerCAmelCase__	)

										__SCREAMING_SNAKE_CASE						:     List[str]					=		[len(res['''bad_concepts''']	) > 0 for res in result]

										return images, has_nsfw_concepts



				@torch.no_grad()
				def 		__magic_name__(       self			:Optional[int]			,			lowerCAmelCase__			:torch.FloatTensor			,			lowerCAmelCase__			:torch.FloatTensor	)		->			Any:
										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		self.vision_model(lowerCAmelCase__	)[1]  # pooled_output
										__SCREAMING_SNAKE_CASE						:     List[str]					=		self.visual_projection(lowerCAmelCase__	)

										__SCREAMING_SNAKE_CASE						:     Dict					=		cosine_distance(lowerCAmelCase__			,			self.special_care_embeds	)
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		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
										__SCREAMING_SNAKE_CASE						:     List[Any]					=		0.0

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

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

										return images, has_nsfw_concepts
