






from typing import Optional, Tuple

import jax
import jax.numpy as jnp
from flax import linen as nn
from flax.core.frozen_dict import FrozenDict
from transformers import CLIPConfig, FlaxPreTrainedModel
from transformers.models.clip.modeling_flax_clip import FlaxCLIPVisionModule






def       _lowerCAmelCase       (      lowerCAmelCase	,      lowerCAmelCase	,      lowerCAmelCase=1e-12		):







       '''simple docstring'''
       UpperCAmelCase	    =			jnp.divide(emb_a.T	,      jnp.clip(jnp.linalg.norm(lowerCAmelCase	,      axis=1		)	,      a_min=lowerCAmelCase		)		).T
       UpperCAmelCase	    =			jnp.divide(emb_a.T	,      jnp.clip(jnp.linalg.norm(lowerCAmelCase	,      axis=1		)	,      a_min=lowerCAmelCase		)		).T
       return jnp.matmul(lowerCAmelCase	,      norm_emb_a.T		)





class   UpperCamelCase_ (       nn.Module  ):
  _A       :					CLIPConfig
  _A       :					jnp.dtype				=	jnp.floataa



  def 		UpperCamelCase_				(      self				)			->			Tuple:



         """simple docstring"""





         UpperCAmelCase	    =			FlaxCLIPVisionModule(self.config.vision_config				)
         UpperCAmelCase	    =			nn.Dense(self.config.projection_dim ,		use_bias=snake_case__ ,		dtype=self.dtype				)

         UpperCAmelCase	    =			self.param("""concept_embeds""" ,		jax.nn.initializers.ones ,		(17, self.config.projection_dim)				)
         UpperCAmelCase	    =			self.param(
             """special_care_embeds""" ,		jax.nn.initializers.ones ,		(3, self.config.projection_dim)				)

         UpperCAmelCase	    =			self.param("""concept_embeds_weights""" ,		jax.nn.initializers.ones ,		(17,)				)
         UpperCAmelCase	    =			self.param("""special_care_embeds_weights""" ,		jax.nn.initializers.ones ,		(3,)				)



  def __call__(      self ,		snake_case__				)			->			Any:



         """simple docstring"""





         UpperCAmelCase	    =			self.vision_model(snake_case__				)[1]
         UpperCAmelCase	    =			self.visual_projection(snake_case__				)

         UpperCAmelCase	    =			jax_cosine_distance(snake_case__ ,		self.special_care_embeds				)
         UpperCAmelCase	    =			jax_cosine_distance(snake_case__ ,		self.concept_embeds				)

         # increase this value to create a stronger `nfsw` filter
         # at the cost of increasing the possibility of filtering benign image inputs
         UpperCAmelCase	    =			0.0

         UpperCAmelCase	    =			special_cos_dist - self.special_care_embeds_weights[None, :] + adjustment
         UpperCAmelCase	    =			jnp.round(snake_case__ ,		3				)
         UpperCAmelCase	    =			jnp.any(special_scores > 0 ,		axis=1 ,		keepdims=snake_case__				)
         # Use a lower threshold if an image has any special care concept
         UpperCAmelCase	    =			is_special_care * 0.01

         UpperCAmelCase	    =			cos_dist - self.concept_embeds_weights[None, :] + special_adjustment
         UpperCAmelCase	    =			jnp.round(snake_case__ ,		3				)
         UpperCAmelCase	    =			jnp.any(concept_scores > 0 ,		axis=1				)

         return has_nsfw_concepts





class   UpperCamelCase_ (       a_  ):
  _A       :					str				=	CLIPConfig
  _A       :					List[Any]				=	'clip_input'
  _A       :					Dict				=	FlaxStableDiffusionSafetyCheckerModule



  def __init__(      self ,		snake_case__ ,		snake_case__ = None ,		snake_case__ = 0 ,		snake_case__ = jnp.floataa ,		snake_case__ = True ,		**snake_case__ ,		)			->			Optional[int]:



         """simple docstring"""





         if input_shape is None:
                UpperCAmelCase	    =			(1, 2_24, 2_24, 3)
         UpperCAmelCase	    =			self.module_class(config=snake_case__ ,		dtype=snake_case__ ,		**snake_case__				)
         super().__init__(snake_case__ ,		snake_case__ ,		input_shape=snake_case__ ,		seed=snake_case__ ,		dtype=snake_case__ ,		_do_init=_do_init				)



  def 		UpperCamelCase_				(      self ,		snake_case__ ,		snake_case__ ,		snake_case__ = None				)			->			FrozenDict:



         """simple docstring"""





         UpperCAmelCase	    =			jax.random.normal(snake_case__ ,		snake_case__				)

         UpperCAmelCase    ,       UpperCAmelCase	    =			jax.random.split(snake_case__				)
         UpperCAmelCase	    =			{"""params""": params_rng, """dropout""": dropout_rng}

         UpperCAmelCase	    =			self.module.init(snake_case__ ,		snake_case__				)["""params"""]

         return random_params



  def __call__(      self ,		snake_case__ ,		snake_case__ = None ,		)			->			Optional[int]:



         """simple docstring"""





         UpperCAmelCase	    =			jnp.transpose(snake_case__ ,		(0, 2, 3, 1)				)

         return self.module.apply(
             {"""params""": params or self.params} ,		jnp.array(snake_case__ ,		dtype=jnp.floataa				) ,		rngs={} ,		)








"""simple docstring"""







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


lowerCAmelCase_	:    Optional[int]            =      {
    '''configuration_graphormer''': ['''GRAPHORMER_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''GraphormerConfig'''],
}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowerCAmelCase_	:    str            =      [
     '''GRAPHORMER_PRETRAINED_MODEL_ARCHIVE_LIST''',
     '''GraphormerForGraphClassification''',
     '''GraphormerModel''',
     '''GraphormerPreTrainedModel''',
 ]


if TYPE_CHECKING:
 from .configuration_graphormer import GRAPHORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, GraphormerConfig

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_graphormer import (
      GRAPHORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
      GraphormerForGraphClassification,
      GraphormerModel,
      GraphormerPreTrainedModel,
  )


else:
 import sys

 lowerCAmelCase_	:    Optional[Any]            =      _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)
