

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







class  UpperCamelCase	(    nn.Module						):
   lowercase							=						42
   lowercase							=						jnp.floataa
   def     _UpperCAmelCase     (    self    )			->   Union[str, Any]:


        '''simple docstring'''


        lowercase_			:       Optional[Any]									=			FlaxCLIPVisionModule(self.config.vision_config    )
        lowercase_			:       Union[str, Any]									=			nn.Dense(self.config.projection_dim  ,use_bias=__UpperCamelCase  ,dtype=self.dtype    )

        lowercase_			:       List[Any]									=			self.param('concept_embeds'  ,jax.nn.initializers.ones  ,(17, self.config.projection_dim)    )
        lowercase_			:       Dict									=			self.param(
            'special_care_embeds'  ,jax.nn.initializers.ones  ,(3, self.config.projection_dim)    )

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







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


        '''simple docstring'''


        lowercase_			:       List[str]									=			self.vision_model(__UpperCamelCase    )[1]
        lowercase_			:       Optional[int]									=			self.visual_projection(__UpperCamelCase    )

        lowercase_			:       Tuple									=			jax_cosine_distance(__UpperCamelCase  ,self.special_care_embeds    )
        lowercase_			:       str									=			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
        lowercase_			:       List[Any]									=			0.0

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

        lowercase_			:       Optional[Any]									=			cos_dist - self.concept_embeds_weights[None, :] + special_adjustment
        lowercase_			:       Optional[int]									=			jnp.round(__UpperCamelCase  ,3    )
        lowercase_			:       int									=			jnp.any(concept_scores > 0  ,axis=1    )

        return has_nsfw_concepts







class  UpperCamelCase	(    lowercase__						):
   lowercase							=						CLIPConfig
   lowercase							=						'clip_input'
   lowercase							=						FlaxStableDiffusionSafetyCheckerModule
   def __init__(    self  ,__UpperCamelCase  ,__UpperCamelCase = None  ,__UpperCamelCase = 0  ,__UpperCamelCase = jnp.floataa  ,__UpperCamelCase = True  ,**__UpperCamelCase  ,)			->   Tuple:


        '''simple docstring'''


        if input_shape is None:
             lowercase_			:       Optional[Any]									=			(1, 224, 224, 3)
        lowercase_			:       Optional[int]									=			self.module_class(config=__UpperCamelCase  ,dtype=__UpperCamelCase  ,**__UpperCamelCase    )
        super().__init__(__UpperCamelCase  ,__UpperCamelCase  ,input_shape=__UpperCamelCase  ,seed=__UpperCamelCase  ,dtype=__UpperCamelCase  ,_do_init=_do_init    )
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase = None    )			->   Any:


        '''simple docstring'''


        lowercase_			:       Union[str, Any]									=			jax.random.normal(__UpperCamelCase  ,__UpperCamelCase    )

        lowercase_							,						lowercase_			:       Dict									=			jax.random.split(__UpperCamelCase    )
        lowercase_			:       Union[str, Any]									=			{'params': params_rng, 'dropout': dropout_rng}

        lowercase_			:       List[str]									=			self.module.init(__UpperCamelCase  ,__UpperCamelCase    )['params']

        return random_params







   def __call__(    self  ,__UpperCamelCase  ,__UpperCamelCase = None  ,)			->   Dict:


        '''simple docstring'''


        lowercase_			:       Optional[Any]									=			jnp.transpose(__UpperCamelCase  ,(0, 2, 3, 1)    )

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


"""simple docstring"""







import functools
import operator

from ...configuration_utils import PretrainedConfig
from ...utils import logging


__SCREAMING_SNAKE_CASE								=logging.get_logger(__name__)

__SCREAMING_SNAKE_CASE								={
    "facebook/wav2vec2-base-960h": "https://huggingface.co/facebook/wav2vec2-base-960h/resolve/main/config.json",
    # See all Wav2Vec2 models at https://huggingface.co/models?filter=wav2vec2
}







class  UpperCamelCase	(    lowercase_						):
   lowercase							=						'wav2vec2'
   def __init__(    self  ,__UpperCamelCase=32  ,__UpperCamelCase=768  ,__UpperCamelCase=12  ,__UpperCamelCase=12  ,__UpperCamelCase=3072  ,__UpperCamelCase="gelu"  ,__UpperCamelCase=0.1  ,__UpperCamelCase=0.1  ,__UpperCamelCase=0.1  ,__UpperCamelCase=0.0  ,__UpperCamelCase=0.0  ,__UpperCamelCase=0.1  ,__UpperCamelCase=0.1  ,__UpperCamelCase=0.02  ,__UpperCamelCase=1e-5  ,__UpperCamelCase="group"  ,__UpperCamelCase="gelu"  ,__UpperCamelCase=(512, 512, 512, 512, 512, 512, 512)  ,__UpperCamelCase=(5, 2, 2, 2, 2, 2, 2)  ,__UpperCamelCase=(10, 3, 3, 3, 3, 2, 2)  ,__UpperCamelCase=False  ,__UpperCamelCase=128  ,__UpperCamelCase=16  ,__UpperCamelCase=False  ,__UpperCamelCase=True  ,__UpperCamelCase=0.05  ,__UpperCamelCase=10  ,__UpperCamelCase=2  ,__UpperCamelCase=0.0  ,__UpperCamelCase=10  ,__UpperCamelCase=0  ,__UpperCamelCase=320  ,__UpperCamelCase=2  ,__UpperCamelCase=0.1  ,__UpperCamelCase=100  ,__UpperCamelCase=256  ,__UpperCamelCase=256  ,__UpperCamelCase=0.1  ,__UpperCamelCase="sum"  ,__UpperCamelCase=False  ,__UpperCamelCase=False  ,__UpperCamelCase=256  ,__UpperCamelCase=(512, 512, 512, 512, 1500)  ,__UpperCamelCase=(5, 3, 3, 1, 1)  ,__UpperCamelCase=(1, 2, 3, 1, 1)  ,__UpperCamelCase=512  ,__UpperCamelCase=0  ,__UpperCamelCase=1  ,__UpperCamelCase=2  ,__UpperCamelCase=False  ,__UpperCamelCase=3  ,__UpperCamelCase=2  ,__UpperCamelCase=3  ,__UpperCamelCase=None  ,__UpperCamelCase=None  ,**__UpperCamelCase  ,)			->   Optional[Any]:


        '''simple docstring'''


        super().__init__(**__UpperCamelCase  ,pad_token_id=__UpperCamelCase  ,bos_token_id=__UpperCamelCase  ,eos_token_id=__UpperCamelCase    )
        lowercase_			:       Optional[Any]									=			hidden_size
        lowercase_			:       Tuple									=			feat_extract_norm
        lowercase_			:       Dict									=			feat_extract_activation
        lowercase_			:       List[str]									=			list(__UpperCamelCase    )
        lowercase_			:       str									=			list(__UpperCamelCase    )
        lowercase_			:       Dict									=			list(__UpperCamelCase    )
        lowercase_			:       Optional[Any]									=			conv_bias
        lowercase_			:       Dict									=			num_conv_pos_embeddings
        lowercase_			:       List[str]									=			num_conv_pos_embedding_groups
        lowercase_			:       Optional[Any]									=			len(self.conv_dim    )
        lowercase_			:       Any									=			num_hidden_layers
        lowercase_			:       List[Any]									=			intermediate_size
        lowercase_			:       List[Any]									=			hidden_act
        lowercase_			:       Optional[int]									=			num_attention_heads
        lowercase_			:       int									=			hidden_dropout
        lowercase_			:       Dict									=			attention_dropout
        lowercase_			:       Union[str, Any]									=			activation_dropout
        lowercase_			:       Tuple									=			feat_proj_dropout
        lowercase_			:       List[str]									=			final_dropout
        lowercase_			:       Union[str, Any]									=			layerdrop
        lowercase_			:       List[str]									=			layer_norm_eps
        lowercase_			:       Optional[int]									=			initializer_range
        lowercase_			:       List[Any]									=			vocab_size
        lowercase_			:       Optional[int]									=			do_stable_layer_norm
        lowercase_			:       Union[str, Any]									=			use_weighted_layer_sum

        if (
            (len(self.conv_stride    ) != self.num_feat_extract_layers)
            or (len(self.conv_kernel    ) != self.num_feat_extract_layers)
            or (len(self.conv_dim    ) != self.num_feat_extract_layers)
        ):
             raise ValueError(
                 'Configuration for convolutional layers is incorrect. It is required that `len(config.conv_dim)` =='
                 ' `len(config.conv_stride)` == `len(config.conv_kernel)`, but is `len(config.conv_dim) ='
                 f''' {len(self.conv_dim    )}`, `len(config.conv_stride) = {len(self.conv_stride    )}`,'''
                 f''' `len(config.conv_kernel) = {len(self.conv_kernel    )}`.'''    )

        # fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
        lowercase_			:       Dict									=			apply_spec_augment
        lowercase_			:       Optional[int]									=			mask_time_prob
        lowercase_			:       Union[str, Any]									=			mask_time_length
        lowercase_			:       List[str]									=			mask_time_min_masks
        lowercase_			:       List[str]									=			mask_feature_prob
        lowercase_			:       Any									=			mask_feature_length
        lowercase_			:       List[Any]									=			mask_feature_min_masks

        # parameters for pretraining with codevector quantized representations
        lowercase_			:       List[Any]									=			num_codevectors_per_group
        lowercase_			:       Optional[int]									=			num_codevector_groups
        lowercase_			:       Dict									=			contrastive_logits_temperature
        lowercase_			:       int									=			feat_quantizer_dropout
        lowercase_			:       Optional[int]									=			num_negatives
        lowercase_			:       str									=			codevector_dim
        lowercase_			:       str									=			proj_codevector_dim
        lowercase_			:       Optional[Any]									=			diversity_loss_weight

        # ctc loss
        lowercase_			:       Tuple									=			ctc_loss_reduction
        lowercase_			:       int									=			ctc_zero_infinity

        # adapter
        lowercase_			:       int									=			add_adapter
        lowercase_			:       Dict									=			adapter_kernel_size
        lowercase_			:       List[str]									=			adapter_stride
        lowercase_			:       Dict									=			num_adapter_layers
        lowercase_			:       Dict									=			output_hidden_size or hidden_size
        lowercase_			:       Optional[Any]									=			adapter_attn_dim

        # SequenceClassification-specific parameter. Feel free to ignore for other classes.
        lowercase_			:       Dict									=			classifier_proj_size

        # XVector-specific parameters. Feel free to ignore for other classes.
        lowercase_			:       Any									=			list(__UpperCamelCase    )
        lowercase_			:       str									=			list(__UpperCamelCase    )
        lowercase_			:       Any									=			list(__UpperCamelCase    )
        lowercase_			:       Tuple									=			xvector_output_dim







   @property
   def     _UpperCAmelCase     (    self    )			->   str:


        '''simple docstring'''


        return functools.reduce(operator.mul  ,self.conv_stride  ,1    )

