


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    (      _a       ,		_a       ,		_a=1E-12		)					->							List[Any]:
     '''simple docstring'''





     lowercase_       :Any     =       jnp.divide(emb_a.T       ,		jnp.clip(jnp.linalg.norm(_a       ,		axis=1		)       ,		a_min=_a		)		).T
     lowercase_       :Tuple     =       jnp.divide(emb_a.T       ,		jnp.clip(jnp.linalg.norm(_a       ,		axis=1		)       ,		a_min=_a		)		).T
     return jnp.matmul(_a       ,		norm_emb_a.T		)






class      UpperCamelCase					(		nn.Module     ):


   '''simple docstring'''







   lowercase					:	CLIPConfig
   lowercase					:	jnp.dtype									=jnp.floataa
   def 		UpperCamelCase  (     self		):
        lowercase_       :List[Any]     =       FlaxCLIPVisionModule(self.config.vision_config		)
        lowercase_       :Any     =       nn.Dense(self.config.projection_dim   ,							use_bias=UpperCamelCase_   ,							dtype=self.dtype		)

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

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




   def __call__(     self   ,							UpperCamelCase_		):
        lowercase_       :Optional[Any]     =       self.vision_model(UpperCamelCase_		)[1]
        lowercase_       :int     =       self.visual_projection(UpperCamelCase_		)

        lowercase_       :List[Any]     =       jax_cosine_distance(UpperCamelCase_   ,							self.special_care_embeds		)
        lowercase_       :List[Any]     =       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_       :Optional[Any]     =       0.0

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

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

        return has_nsfw_concepts






class      UpperCamelCase					(		lowercase__     ):


   '''simple docstring'''







   lowercase					:	Optional[int]									=CLIPConfig
   lowercase					:	str									="""clip_input"""
   lowercase					:	Optional[int]									=FlaxStableDiffusionSafetyCheckerModule
   def __init__(     self   ,							UpperCamelCase_   ,							UpperCamelCase_ = None   ,							UpperCamelCase_ = 0   ,							UpperCamelCase_ = jnp.floataa   ,							UpperCamelCase_ = True   ,							**UpperCamelCase_   ,							):
        if input_shape is None:
             lowercase_       :str     =       (1, 224, 224, 3)
        lowercase_       :Optional[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 		UpperCamelCase  (     self   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_ = None		):
        # init input tensor
        lowercase_       :Any     =       jax.random.normal(UpperCamelCase_   ,							UpperCamelCase_		)

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

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

        return random_params




   def __call__(     self   ,							UpperCamelCase_   ,							UpperCamelCase_ = None   ,							):
        lowercase_       :str     =       jnp.transpose(UpperCamelCase_   ,							(0, 2, 3, 1)		)

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





import argparse

import torch

from transformers import OpenAIGPTConfig, OpenAIGPTModel, load_tf_weights_in_openai_gpt
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()




def     UpperCamelCase    (      _a       ,		_a       ,		_a		)					->							List[str]:
     '''simple docstring'''





     if openai_config_file == "":
          lowercase_       :str     =       OpenAIGPTConfig()
     else:
          lowercase_       :int     =       OpenAIGPTConfig.from_json_file(_a		)
     lowercase_       :int     =       OpenAIGPTModel(_a		)

     # Load weights from numpy
     load_tf_weights_in_openai_gpt(_a       ,		_a       ,		_a		)

     # Save pytorch-model
     lowercase_       :Optional[int]     =       pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
     lowercase_       :List[str]     =       pytorch_dump_folder_path + '''/''' + CONFIG_NAME
     print(f"Save PyTorch model to {pytorch_weights_dump_path}"		)
     torch.save(model.state_dict()       ,		_a		)
     print(f"Save configuration file to {pytorch_config_dump_path}"		)
     with open(_a       ,		'''w'''       ,		encoding='''utf-8'''		) as f:
          f.write(config.to_json_string()		)


if __name__ == "__main__":
   SCREAMING_SNAKE_CASE       :		List[str]         =				argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       "--openai_checkpoint_folder_path",
       default=None,
       type=str,
       required=True,
       help="Path to the TensorFlow checkpoint path.",
   )
   parser.add_argument(
       "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
   )
   parser.add_argument(
       "--openai_config_file",
       default="",
       type=str,
       help=(
           "An optional config json file corresponding to the pre-trained OpenAI model. \n"
           "This specifies the model architecture."
       ),
   )
   SCREAMING_SNAKE_CASE       :		Optional[Any]         =				parser.parse_args()
   convert_openai_checkpoint_to_pytorch(
       args.openai_checkpoint_folder_path, args.openai_config_file, args.pytorch_dump_folder_path
   )

