




"""simple docstring"""


import sys
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import AutoFeatureExtractor, WavaVecaFeatureExtractor
from transformers.testing_utils import TOKEN, USER, get_tests_dir, is_staging_test


sys.path.append(str(Path(__file__).parent.parent / """utils"""))

from test_module.custom_feature_extraction import CustomFeatureExtractor  # noqa E402


_A													=      get_tests_dir("""fixtures""")


class _lowerCamelCase    (					unittest.TestCase						):






   def 	_lowerCAmelCase   (		self       :		str   )			->						Tuple:

         """simple docstring"""
         lowerCAmelCase__      :			Any   =	mock.Mock()
         lowerCAmelCase__      :			List[str]   =	5_00
         lowerCAmelCase__      :			Tuple   =	{}
         lowerCAmelCase__      :			Union[str, Any]   =	HTTPError
         lowerCAmelCase__      :			Union[str, Any]   =	{}

         # Download this model to make sure it's in the cache.
         lowerCAmelCase__      :			int   =	WavaVecaFeatureExtractor.from_pretrained("""hf-internal-testing/tiny-random-wav2vec2"""   )
         # Under the mock environment we get a 500 error when trying to reach the model.
         with mock.patch("""requests.Session.request"""   ,     return_value=UpperCamelCase   ) as mock_head:
               lowerCAmelCase__      :			str   =	WavaVecaFeatureExtractor.from_pretrained("""hf-internal-testing/tiny-random-wav2vec2"""   )
               # This check we did call the fake head request
               mock_head.assert_called()






   def 	_lowerCAmelCase   (		self       :		List[Any]   )			->						Union[str, Any]:

         """simple docstring"""
         lowerCAmelCase__      :			str   =	WavaVecaFeatureExtractor.from_pretrained(
             """https://huggingface.co/hf-internal-testing/tiny-random-wav2vec2/resolve/main/preprocessor_config.json"""   )






@is_staging_test
class _lowerCamelCase    (					unittest.TestCase						):






   @classmethod
   def 	_lowerCAmelCase   (		cls       :		str   )			->						str:

         """simple docstring"""
         lowerCAmelCase__      :			Optional[Any]   =	TOKEN
         HfFolder.save_token(UpperCamelCase   )






   @classmethod
   def 	_lowerCAmelCase   (		cls       :		Any   )			->						List[Any]:

         """simple docstring"""
         try:
               delete_repo(token=cls._token   ,     repo_id="""test-feature-extractor"""   )
         except HTTPError:
               pass

         try:
               delete_repo(token=cls._token   ,     repo_id="""valid_org/test-feature-extractor-org"""   )
         except HTTPError:
               pass

         try:
               delete_repo(token=cls._token   ,     repo_id="""test-dynamic-feature-extractor"""   )
         except HTTPError:
               pass






   def 	_lowerCAmelCase   (		self       :		Optional[Any]   )			->						str:

         """simple docstring"""
         lowerCAmelCase__      :			Tuple   =	WavaVecaFeatureExtractor.from_pretrained(UpperCamelCase   )
         feature_extractor.push_to_hub("""test-feature-extractor"""   ,     use_auth_token=self._token   )

         lowerCAmelCase__      :			List[Any]   =	WavaVecaFeatureExtractor.from_pretrained(f"""{USER}/test-feature-extractor"""   )
         for k, v in feature_extractor.__dict__.items():
               self.assertEqual(UpperCamelCase   ,     getattr(UpperCamelCase   ,     UpperCamelCase   )   )

         # Reset repo
         delete_repo(token=self._token   ,     repo_id="""test-feature-extractor"""   )

         # Push to hub via save_pretrained
         with tempfile.TemporaryDirectory() as tmp_dir:
               feature_extractor.save_pretrained(
                   UpperCamelCase   ,     repo_id="""test-feature-extractor"""   ,     push_to_hub=UpperCamelCase   ,     use_auth_token=self._token   )

         lowerCAmelCase__      :			Optional[Any]   =	WavaVecaFeatureExtractor.from_pretrained(f"""{USER}/test-feature-extractor"""   )
         for k, v in feature_extractor.__dict__.items():
               self.assertEqual(UpperCamelCase   ,     getattr(UpperCamelCase   ,     UpperCamelCase   )   )






   def 	_lowerCAmelCase   (		self       :		Optional[Any]   )			->						Optional[Any]:

         """simple docstring"""
         lowerCAmelCase__      :			int   =	WavaVecaFeatureExtractor.from_pretrained(UpperCamelCase   )
         feature_extractor.push_to_hub("""valid_org/test-feature-extractor"""   ,     use_auth_token=self._token   )

         lowerCAmelCase__      :			Optional[Any]   =	WavaVecaFeatureExtractor.from_pretrained("""valid_org/test-feature-extractor"""   )
         for k, v in feature_extractor.__dict__.items():
               self.assertEqual(UpperCamelCase   ,     getattr(UpperCamelCase   ,     UpperCamelCase   )   )

         # Reset repo
         delete_repo(token=self._token   ,     repo_id="""valid_org/test-feature-extractor"""   )

         # Push to hub via save_pretrained
         with tempfile.TemporaryDirectory() as tmp_dir:
               feature_extractor.save_pretrained(
                   UpperCamelCase   ,     repo_id="""valid_org/test-feature-extractor-org"""   ,     push_to_hub=UpperCamelCase   ,     use_auth_token=self._token   )

         lowerCAmelCase__      :			str   =	WavaVecaFeatureExtractor.from_pretrained("""valid_org/test-feature-extractor-org"""   )
         for k, v in feature_extractor.__dict__.items():
               self.assertEqual(UpperCamelCase   ,     getattr(UpperCamelCase   ,     UpperCamelCase   )   )






   def 	_lowerCAmelCase   (		self       :		Union[str, Any]   )			->						Union[str, Any]:

         """simple docstring"""
         CustomFeatureExtractor.register_for_auto_class()
         lowerCAmelCase__      :			Tuple   =	CustomFeatureExtractor.from_pretrained(UpperCamelCase   )

         feature_extractor.push_to_hub("""test-dynamic-feature-extractor"""   ,     use_auth_token=self._token   )

         # This has added the proper auto_map field to the config
         self.assertDictEqual(
             feature_extractor.auto_map   ,     {"""AutoFeatureExtractor""": """custom_feature_extraction.CustomFeatureExtractor"""}   ,     )

         lowerCAmelCase__      :			Tuple   =	AutoFeatureExtractor.from_pretrained(
             f"""{USER}/test-dynamic-feature-extractor"""   ,     trust_remote_code=UpperCamelCase   )
         # Can't make an isinstance check because the new_feature_extractor is from the CustomFeatureExtractor class of a dynamic module
         self.assertEqual(new_feature_extractor.__class__.__name__   ,     """CustomFeatureExtractor"""   )











"""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 	lowercase_				(			__UpperCAmelCase    ,		__UpperCAmelCase    ,		__UpperCAmelCase=1E-12				)						->    Dict:
      lowerCAmelCase__      :			Optional[int]   =	jnp.divide(emb_a.T    ,		jnp.clip(jnp.linalg.norm(__UpperCAmelCase    ,		axis=1				)    ,		a_min=__UpperCAmelCase				)				).T
      lowerCAmelCase__      :			Any   =	jnp.divide(emb_a.T    ,		jnp.clip(jnp.linalg.norm(__UpperCAmelCase    ,		axis=1				)    ,		a_min=__UpperCAmelCase				)				).T
      return jnp.matmul(__UpperCAmelCase    ,		norm_emb_a.T				)


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






   def 	_lowerCAmelCase   (		self       :		List[str]   )			->						Tuple:

         """simple docstring"""
         lowerCAmelCase__      :			Optional[int]   =	FlaxCLIPVisionModule(self.config.vision_config   )
         lowerCAmelCase__      :			Dict   =	nn.Dense(self.config.projection_dim   ,     use_bias=UpperCamelCase   ,     dtype=self.dtype   )

         lowerCAmelCase__      :			Optional[int]   =	self.param("""concept_embeds"""   ,     jax.nn.initializers.ones   ,     (17, self.config.projection_dim)   )
         lowerCAmelCase__      :			Optional[int]   =	self.param(
             """special_care_embeds"""   ,     jax.nn.initializers.ones   ,     (3, self.config.projection_dim)   )

         lowerCAmelCase__      :			List[Any]   =	self.param("""concept_embeds_weights"""   ,     jax.nn.initializers.ones   ,     (17,)   )
         lowerCAmelCase__      :			List[Any]   =	self.param("""special_care_embeds_weights"""   ,     jax.nn.initializers.ones   ,     (3,)   )






   def __call__(		self       :		str   ,     UpperCamelCase       :		List[Any]   )			->						Union[str, Any]:

         """simple docstring"""
         lowerCAmelCase__      :			Union[str, Any]   =	self.vision_model(UpperCamelCase   )[1]
         lowerCAmelCase__      :			List[Any]   =	self.visual_projection(UpperCamelCase   )

         lowerCAmelCase__      :			str   =	jax_cosine_distance(UpperCamelCase   ,     self.special_care_embeds   )
         lowerCAmelCase__      :			Optional[int]   =	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__      :			Any   =	0.0

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

         lowerCAmelCase__      :			List[Any]   =	cos_dist - self.concept_embeds_weights[None, :] + special_adjustment
         lowerCAmelCase__      :			Dict   =	jnp.round(UpperCamelCase   ,     3   )
         lowerCAmelCase__      :			List[str]   =	jnp.any(concept_scores > 0   ,     axis=1   )

         return has_nsfw_concepts






class _lowerCamelCase    (					a_						):
   _lowerCamelCase    :Optional[Any]     				=     CLIPConfig
   _lowerCamelCase    :Dict     				=     "clip_input"
   _lowerCamelCase    :int     				=     FlaxStableDiffusionSafetyCheckerModule






   def __init__(		self       :		int   ,     UpperCamelCase       :		CLIPConfig   ,     UpperCamelCase       :		Optional[Tuple] = None   ,     UpperCamelCase       :		int = 0   ,     UpperCamelCase       :		jnp.dtype = jnp.floataa   ,     UpperCamelCase       :		bool = True   ,     **UpperCamelCase       :		int   ,     )			->						str:

         """simple docstring"""
         if input_shape is None:
               lowerCAmelCase__      :			str   =	(1, 2_24, 2_24, 3)
         lowerCAmelCase__      :			Union[str, Any]   =	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       :		Dict   ,     UpperCamelCase       :		jax.random.KeyArray   ,     UpperCamelCase       :		Tuple   ,     UpperCamelCase       :		FrozenDict = None   )			->						FrozenDict:

         """simple docstring"""
         # init input tensor
         lowerCAmelCase__      :			Union[str, Any]   =	jax.random.normal(UpperCamelCase   ,     UpperCamelCase   )

         lowerCAmelCase__ ,   lowerCAmelCase__      :			Union[str, Any]   =	jax.random.split(UpperCamelCase   )
         lowerCAmelCase__      :			List[Any]   =	{"""params""": params_rng, """dropout""": dropout_rng}

         lowerCAmelCase__      :			str   =	self.module.init(UpperCamelCase   ,     UpperCamelCase   )["""params"""]

         return random_params






   def __call__(		self       :		Dict   ,     UpperCamelCase       :		Tuple   ,     UpperCamelCase       :		dict = None   ,     )			->						List[Any]:

         """simple docstring"""
         lowerCAmelCase__      :			Union[str, Any]   =	jnp.transpose(UpperCamelCase   ,     (0, 2, 3, 1)   )

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





