import torch
import torch.nn as nn
from transformers import CLIPConfig, CLIPVisionModel, PreTrainedModel

from ...utils import logging


snake_case__				= logging.get_logger(__name__)
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase		)						->    Any:



      '''simple docstring'''

      _lowerCamelCase			    =				nn.functional.normalize(__UpperCAmelCase		)
      _lowerCamelCase			    =				nn.functional.normalize(__UpperCAmelCase		)
      return torch.mm(__UpperCAmelCase ,						normalized_text_embeds.t()		)




class  UpperCamelCase				( __lowercase ):

     '''simple docstring'''

     A_										=  CLIPConfig

     A_										=  ['CLIPEncoderLayer']



     def __init__(					self				,				A_						) ->				Any:






           """simple docstring"""
           super().__init__(A_						)

           _lowerCamelCase			    =				CLIPVisionModel(config.vision_config						)
           _lowerCamelCase			    =				nn.Linear(config.vision_config.hidden_size				,				config.projection_dim				,				bias=A_						)

           _lowerCamelCase			    =				nn.Parameter(torch.ones(17				,				config.projection_dim						)				,				requires_grad=A_						)
           _lowerCamelCase			    =				nn.Parameter(torch.ones(3				,				config.projection_dim						)				,				requires_grad=A_						)

           _lowerCamelCase			    =				nn.Parameter(torch.ones(17						)				,				requires_grad=A_						)
           _lowerCamelCase			    =				nn.Parameter(torch.ones(3						)				,				requires_grad=A_						)



     @torch.no_grad()
     def 		UpperCamelCase_						(					self				,				A_				,				A_						) ->				int:






           """simple docstring"""
           _lowerCamelCase			    =				self.vision_model(A_						)[1]  # pooled_output
           _lowerCamelCase			    =				self.visual_projection(A_						)

           # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
           _lowerCamelCase			    =				cosine_distance(A_				,				self.special_care_embeds						).cpu().float().numpy()
           _lowerCamelCase			    =				cosine_distance(A_				,				self.concept_embeds						).cpu().float().numpy()

           _lowerCamelCase			    =				[]
           _lowerCamelCase			    =				image_embeds.shape[0]
           for i in range(A_						):
                 _lowerCamelCase			    =				{'''special_scores''': {}, '''special_care''': [], '''concept_scores''': {}, '''bad_concepts''': []}

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

                 for concept_idx in range(len(special_cos_dist[0]						)						):
                       _lowerCamelCase			    =				special_cos_dist[i][concept_idx]
                       _lowerCamelCase			    =				self.special_care_embeds_weights[concept_idx].item()
                       _lowerCamelCase			    =				round(concept_cos - concept_threshold + adjustment				,				3						)
                       if result_img["special_scores"][concept_idx] > 0:
                             result_img["special_care"].append({concept_idx, result_img['''special_scores'''][concept_idx]}						)
                             _lowerCamelCase			    =				0.01

                 for concept_idx in range(len(cos_dist[0]						)						):
                       _lowerCamelCase			    =				cos_dist[i][concept_idx]
                       _lowerCamelCase			    =				self.concept_embeds_weights[concept_idx].item()
                       _lowerCamelCase			    =				round(concept_cos - concept_threshold + adjustment				,				3						)
                       if result_img["concept_scores"][concept_idx] > 0:
                             result_img["bad_concepts"].append(A_						)

                 result.append(A_						)

           _lowerCamelCase			    =				[len(res['''bad_concepts''']						) > 0 for res in result]

           return images, has_nsfw_concepts





     @torch.no_grad()
     def 		UpperCamelCase_						(					self				,				A_				,				A_						) ->				Optional[Any]:






           """simple docstring"""
           _lowerCamelCase			    =				self.vision_model(A_						)[1]  # pooled_output
           _lowerCamelCase			    =				self.visual_projection(A_						)

           _lowerCamelCase			    =				cosine_distance(A_				,				self.special_care_embeds						)
           _lowerCamelCase			    =				cosine_distance(A_				,				self.concept_embeds						)

           # increase this value to create a stronger `nsfw` filter
           # at the cost of increasing the possibility of filtering benign images
           _lowerCamelCase			    =				0.0

           _lowerCamelCase			    =				special_cos_dist - self.special_care_embeds_weights + adjustment
           # special_scores = special_scores.round(decimals=3)
           _lowerCamelCase			    =				torch.any(special_scores > 0				,				dim=1						)
           _lowerCamelCase			    =				special_care * 0.01
           _lowerCamelCase			    =				special_adjustment.unsqueeze(1						).expand(-1				,				cos_dist.shape[1]						)

           _lowerCamelCase			    =				(cos_dist - self.concept_embeds_weights) + special_adjustment
           # concept_scores = concept_scores.round(decimals=3)
           _lowerCamelCase			    =				torch.any(concept_scores > 0				,				dim=1						)

           return images, has_nsfw_concepts
from typing import List

import numpy as np
def       __magic_name__(  __UpperCAmelCase		)						->    int:



      '''simple docstring'''

      _lowerCamelCase			    =				{key: len(__UpperCAmelCase		) for key, value in gen_kwargs.items() if isinstance(__UpperCAmelCase ,						__UpperCAmelCase		)}
      if len(set(lists_lengths.values()		)		) > 1:
            raise RuntimeError(
                (
                    '''Sharding is ambiguous for this dataset: '''
                    + '''we found several data sources lists of different lengths, and we don\'t know over which list we should parallelize:\n'''
                    + '''\n'''.join(F'\t- key {key} has length {length}' for key, length in lists_lengths.items()		)
                    + '''\nTo fix this, check the \'gen_kwargs\' and make sure to use lists only for data sources, '''
                    + '''and use tuples otherwise. In the end there should only be one single list, or several lists with the same length.'''
                )		)
      _lowerCamelCase			    =				max(lists_lengths.values() ,						default=0		)
      return max(1 ,						__UpperCAmelCase		)
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase		)						->    List[range]:



      '''simple docstring'''

      _lowerCamelCase			    =				[]
      for group_idx in range(__UpperCAmelCase		):
            _lowerCamelCase			    =				num_shards // max_num_jobs + (group_idx < (num_shards % max_num_jobs))
            if num_shards_to_add == 0:
                  break
            _lowerCamelCase			    =				shards_indices_per_group[-1].stop if shards_indices_per_group else 0
            _lowerCamelCase			    =				range(__UpperCAmelCase ,						start + num_shards_to_add		)
            shards_indices_per_group.append(__UpperCAmelCase		)
      return shards_indices_per_group
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase		)						->    List[dict]:



      '''simple docstring'''

      _lowerCamelCase			    =				_number_of_shards_in_gen_kwargs(__UpperCAmelCase		)
      if num_shards == 1:
            return [dict(__UpperCAmelCase		)]
      else:
            _lowerCamelCase			    =				_distribute_shards(num_shards=__UpperCAmelCase ,						max_num_jobs=__UpperCAmelCase		)
            return [
                {
                    key: [value[shard_idx] for shard_idx in shard_indices_per_group[group_idx]]
                    if isinstance(__UpperCAmelCase ,						__UpperCAmelCase		)
                    else value
                    for key, value in gen_kwargs.items()
                }
                for group_idx in range(len(__UpperCAmelCase		)		)
            ]
def       __magic_name__(  __UpperCAmelCase		)						->    dict:



      '''simple docstring'''

      return {
          key: [value for gen_kwargs in gen_kwargs_list for value in gen_kwargs[key]]
          if isinstance(gen_kwargs_list[0][key] ,						__UpperCAmelCase		)
          else gen_kwargs_list[0][key]
          for key in gen_kwargs_list[0]
      }




def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase		)						->    dict:



      '''simple docstring'''

      _lowerCamelCase			    =				{len(__UpperCAmelCase		) for value in gen_kwargs.values() if isinstance(__UpperCAmelCase ,						__UpperCAmelCase		)}
      _lowerCamelCase			    =				{}
      for size in list_sizes:
            _lowerCamelCase			    =				list(range(__UpperCAmelCase		)		)
            rng.shuffle(indices_per_size[size]		)
      # Now let's copy the gen_kwargs and shuffle the lists based on their sizes
      _lowerCamelCase			    =				dict(__UpperCAmelCase		)
      for key, value in shuffled_kwargs.items():
            if isinstance(__UpperCAmelCase ,						__UpperCAmelCase		):
                  _lowerCamelCase			    =				[value[i] for i in indices_per_size[len(__UpperCAmelCase		)]]
      return shuffled_kwargs