





'''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    _lowerCamelCase  (__lowerCamelCase  :     Optional[int]					,  __lowerCamelCase  :     Tuple					,  __lowerCamelCase  :     Any=1e-12     )   ->    Union[str, Any]:
   a__     = jnp.divide(emb_a.T					,  jnp.clip(jnp.linalg.norm(__lowerCamelCase					,  axis=1     )					,  a_min=__lowerCamelCase     )     ).T
   a__     = 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     ):
 lowerCAmelCase__					:			CLIPConfig
 lowerCAmelCase__					:			jnp.dtype								=						jnp.floataa







 def 						__a   (				self			:   str		):




    '''simple docstring'''
    a__     = FlaxCLIPVisionModule(self.config.vision_config		)
    a__     = nn.Dense(self.config.projection_dim , use_bias=lowerCamelCase , dtype=self.dtype		)

    a__     = self.param("concept_embeds" , jax.nn.initializers.ones , (1_7, self.config.projection_dim)		)
    a__     = self.param(
        "special_care_embeds" , jax.nn.initializers.ones , (3, self.config.projection_dim)		)

    a__     = self.param("concept_embeds_weights" , jax.nn.initializers.ones , (1_7,)		)
    a__     = self.param("special_care_embeds_weights" , jax.nn.initializers.ones , (3,)		)







 def __call__(				self			:   Union[str, Any] , lowerCamelCase			:   Union[str, Any]		):




    '''simple docstring'''
    a__     = self.vision_model(lowerCamelCase		)[1]
    a__     = self.visual_projection(lowerCamelCase		)

    a__     = jax_cosine_distance(lowerCamelCase , self.special_care_embeds		)
    a__     = jax_cosine_distance(lowerCamelCase , self.concept_embeds		)

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

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

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

    return has_nsfw_concepts






class       UpperCamelCase__  (		__lowerCAmelCase     ):
 lowerCAmelCase__					:			List[Any]								=						CLIPConfig
 lowerCAmelCase__					:			Optional[int]								=						"clip_input"
 lowerCAmelCase__					:			Optional[Any]								=						FlaxStableDiffusionSafetyCheckerModule







 def __init__(				self			:   Optional[int] , lowerCamelCase			:   CLIPConfig , lowerCamelCase			:   Optional[Tuple] = None , lowerCamelCase			:   int = 0 , lowerCamelCase			:   jnp.dtype = jnp.floataa , lowerCamelCase			:   bool = True , **lowerCamelCase			:   Dict , ):




    '''simple docstring'''
    if input_shape is None:
       a__     = (1, 2_2_4, 2_2_4, 3)
    a__     = self.module_class(config=lowerCamelCase , dtype=lowerCamelCase , **lowerCamelCase		)
    super().__init__(lowerCamelCase , lowerCamelCase , input_shape=lowerCamelCase , seed=lowerCamelCase , dtype=lowerCamelCase , _do_init=_do_init		)







 def 						__a   (				self			:   Tuple , lowerCamelCase			:   jax.random.KeyArray , lowerCamelCase			:   Tuple , lowerCamelCase			:   FrozenDict = None		):




    '''simple docstring'''
    a__     = jax.random.normal(lowerCamelCase , lowerCamelCase		)

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

    a__     = self.module.init(lowerCamelCase , lowerCamelCase		)["params"]

    return random_params







 def __call__(				self			:   int , lowerCamelCase			:   Dict , lowerCamelCase			:   dict = None , ):




    '''simple docstring'''
    a__     = jnp.transpose(lowerCamelCase , (0, 2, 3, 1)		)

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










'''simple docstring'''


import argparse
import json

from tqdm import tqdm







def    _lowerCamelCase  ()   ->    Any:
			a__     = argparse.ArgumentParser()

			# Required parameters
			parser.add_argument(
			    "--src_path"					,  type=__lowerCamelCase					,  default="biencoder-nq-dev.json"					,  help="Path to raw DPR training data"					,  )
			parser.add_argument(
			    "--evaluation_set"					,  type=__lowerCamelCase					,  help="where to store parsed evaluation_set file"					,  )
			parser.add_argument(
			    "--gold_data_path"					,  type=__lowerCamelCase					,  help="where to store parsed gold_data_path file"					,  )
			a__     = parser.parse_args()

			with open(args.src_path					,  "r"     ) as src_file, open(args.evaluation_set					,  "w"     ) as eval_file, open(
			    args.gold_data_path					,  "w"     ) as gold_file:
						a__     = json.load(__lowerCamelCase     )
						for dpr_record in tqdm(__lowerCamelCase     ):
									a__     = dpr_record["question"]
									a__     = [context["title"] for context in dpr_record["positive_ctxs"]]
									eval_file.write(question + "\n"     )
									gold_file.write("\t".join(__lowerCamelCase     ) + "\n"     )


if __name__ == "__main__":
		main()



