

'''simple docstring'''



import contextlib
import importlib
import io
import unittest

import transformers

# Try to import everything from transformers to ensure every object can be loaded.
from transformers import *  # noqa F406
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, require_flax, require_tf, require_torch
from transformers.utils import ContextManagers, find_labels, is_flax_available, is_tf_available, is_torch_available


if is_torch_available():
   from transformers import BertForPreTraining, BertForQuestionAnswering, BertForSequenceClassification

if is_tf_available():
   from transformers import TFBertForPreTraining, TFBertForQuestionAnswering, TFBertForSequenceClassification

if is_flax_available():
   from transformers import FlaxBertForPreTraining, FlaxBertForQuestionAnswering, FlaxBertForSequenceClassification


SCREAMING_SNAKE_CASE_     = DUMMY_UNKNOWN_IDENTIFIER
# An actual model hosted on huggingface.co

SCREAMING_SNAKE_CASE_     = "main"
# Default branch name
SCREAMING_SNAKE_CASE_     = "f2c752cfc5c0ab6f4bdec59acea69eefbee381c2"
# One particular commit (not the top of `main`)
SCREAMING_SNAKE_CASE_     = "aaaaaaa"
# This commit does not exist, so we should 404.

SCREAMING_SNAKE_CASE_     = "d9e9f15bc825e4b2c9249e9578f884bbcb5e3684"
# Sha-1 of config.json on the top of `main`, for checking purposes
SCREAMING_SNAKE_CASE_     = "4b243c475af8d0a7754e87d7d096c92e5199ec2fe168a2ee7998e3b8e9bcb1d3"

@contextlib.contextmanager
def 	lowerCamelCase__							(   )		->     Any:



    """simple docstring"""







    print('Welcome!')
    yield
    print('Bye!')

@contextlib.contextmanager
def 	lowerCamelCase__							(   )		->     List[str]:



    """simple docstring"""







    print('Bonjour!')
    yield
    print('Au revoir!')

class       SCREAMING_SNAKE_CASE       (				unittest.TestCase       ):






      '''simple docstring'''






      def 						__UpperCAmelCase      (       self			:		List[str] ):



          """simple docstring"""





          assert transformers.__spec__ is not None
          assert importlib.util.find_spec('transformers' ) is not None



class       SCREAMING_SNAKE_CASE       (				unittest.TestCase       ):






      '''simple docstring'''






      @unittest.mock.patch('sys.stdout'  ,       new_callable=io.StringIO )
      def 						__UpperCAmelCase      (       self			:		Optional[int]  ,       snake_case			:		int ):



          """simple docstring"""





          with ContextManagers([] ):
              print('Transformers are awesome!' )
          # The print statement adds a new line at the end of the output
          self.assertEqual(mock_stdout.getvalue()  ,       'Transformers are awesome!\n' )




      @unittest.mock.patch('sys.stdout'  ,       new_callable=io.StringIO )
      def 						__UpperCAmelCase      (       self			:		List[Any]  ,       snake_case			:		str ):



          """simple docstring"""





          with ContextManagers([context_en()] ):
              print('Transformers are awesome!' )
          # The output should be wrapped with an English welcome and goodbye
          self.assertEqual(mock_stdout.getvalue()  ,       'Welcome!\nTransformers are awesome!\nBye!\n' )




      @unittest.mock.patch('sys.stdout'  ,       new_callable=io.StringIO )
      def 						__UpperCAmelCase      (       self			:		List[str]  ,       snake_case			:		Any ):



          """simple docstring"""





          with ContextManagers([context_fr(), context_en()] ):
              print('Transformers are awesome!' )
          # The output should be wrapped with an English and French welcome and goodbye
          self.assertEqual(mock_stdout.getvalue()  ,       'Bonjour!\nWelcome!\nTransformers are awesome!\nBye!\nAu revoir!\n' )




      @require_torch
      def 						__UpperCAmelCase      (       self			:		List[str] ):



          """simple docstring"""





          self.assertEqual(find_labels(snake_case )  ,       ['labels'] )
          self.assertEqual(find_labels(snake_case )  ,       ['labels', 'next_sentence_label'] )
          self.assertEqual(find_labels(snake_case )  ,       ['start_positions', 'end_positions'] )

          class       SCREAMING_SNAKE_CASE       (				__snake_case       ):






                '''simple docstring'''






                pass



          self.assertEqual(find_labels(snake_case )  ,       ['labels'] )




      @require_tf
      def 						__UpperCAmelCase      (       self			:		str ):



          """simple docstring"""





          self.assertEqual(find_labels(snake_case )  ,       ['labels'] )
          self.assertEqual(find_labels(snake_case )  ,       ['labels', 'next_sentence_label'] )
          self.assertEqual(find_labels(snake_case )  ,       ['start_positions', 'end_positions'] )

          class       SCREAMING_SNAKE_CASE       (				__snake_case       ):






                '''simple docstring'''






                pass



          self.assertEqual(find_labels(snake_case )  ,       ['labels'] )




      @require_flax
      def 						__UpperCAmelCase      (       self			:		str ):



          """simple docstring"""





          self.assertEqual(find_labels(snake_case )  ,       [] )
          self.assertEqual(find_labels(snake_case )  ,       [] )
          self.assertEqual(find_labels(snake_case )  ,       [] )

          class       SCREAMING_SNAKE_CASE       (				__snake_case       ):






                '''simple docstring'''






                pass



          self.assertEqual(find_labels(snake_case )  ,       [] )







"""simple docstring"""



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      UpperCAmelCase     (       _lowercase :   str			,   _lowercase :   int			,   _lowercase :   Optional[int]=1E-12	)						-> List[str]:


	"""simple docstring"""




	lowerCAmelCase_      							=			jnp.divide(emb_a.T			,   jnp.clip(jnp.linalg.norm(_lowercase			,   axis=1	)			,   a_min=_lowercase	)	).T
	lowerCAmelCase_      							=			jnp.divide(emb_a.T			,   jnp.clip(jnp.linalg.norm(_lowercase			,   axis=1	)			,   a_min=_lowercase	)	).T
	return jnp.matmul(_lowercase			,   norm_emb_a.T	)







class    __a						( nn.Module							):
				lowerCamelCase     :       CLIPConfig
				lowerCamelCase     :       jnp.dtype			       =jnp.floataa







				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






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

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

					lowerCAmelCase_      							=			self.param('''concept_embeds_weights'''      ,	jax.nn.initializers.ones      ,	(17,)   )
					lowerCAmelCase_      							=			self.param('''special_care_embeds_weights'''      ,	jax.nn.initializers.ones      ,	(3,)   )







				def __call__( self      ,	UpperCAmelCase   ):

					'''simple docstring'''






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

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

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

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

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

					return has_nsfw_concepts







class    __a						( __snake_case							):
				lowerCamelCase     :       Optional[Any]			       =CLIPConfig
				lowerCamelCase     :       List[str]			       ='clip_input'
				lowerCamelCase     :       int			       =FlaxStableDiffusionSafetyCheckerModule







				def __init__( self      ,	UpperCAmelCase      ,	UpperCAmelCase = None      ,	UpperCAmelCase = 0      ,	UpperCAmelCase = jnp.floataa      ,	UpperCAmelCase = True      ,	**UpperCAmelCase      ,	):

					'''simple docstring'''






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







				def      lowerCamelCase_     ( self      ,	UpperCAmelCase      ,	UpperCAmelCase      ,	UpperCAmelCase = None   ):

					'''simple docstring'''






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

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

					lowerCAmelCase_      							=			self.module.init(UpperCAmelCase      ,	UpperCAmelCase   )['''params''']

					return random_params







				def __call__( self      ,	UpperCAmelCase      ,	UpperCAmelCase = None      ,	):

					'''simple docstring'''






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

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