






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

from ...utils import logging


__magic_name__   :  Union[str, Any]        =							logging.get_logger(__name__)
def 		lowercase__  (				_UpperCamelCase       ,							_UpperCamelCase) ->			List[Any]:
      """simple docstring"""

      UpperCamelCase					       =  nn.functional.normalize(_UpperCamelCase)
      UpperCamelCase					       =  nn.functional.normalize(_UpperCamelCase)
      return torch.mm(_UpperCamelCase       ,							normalized_text_embeds.t())





class 					A__     (     __snake_case ):

     '''simple docstring'''






     snake_case__					=					CLIPConfig

     snake_case__					=					["""CLIPEncoderLayer"""]







     def __init__( self      :    str   ,						_SCREAMING_SNAKE_CASE      :    CLIPConfig ):






           """simple docstring"""




           super().__init__(_SCREAMING_SNAKE_CASE )

           UpperCamelCase					       =  CLIPVisionModel(config.vision_config )
           UpperCamelCase					       =  nn.Linear(config.vision_config.hidden_size   ,						config.projection_dim   ,						bias=_SCREAMING_SNAKE_CASE )

           UpperCamelCase					       =  nn.Parameter(torch.ones(17   ,						config.projection_dim )   ,						requires_grad=_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  nn.Parameter(torch.ones(3   ,						config.projection_dim )   ,						requires_grad=_SCREAMING_SNAKE_CASE )

           UpperCamelCase					       =  nn.Parameter(torch.ones(17 )   ,						requires_grad=_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  nn.Parameter(torch.ones(3 )   ,						requires_grad=_SCREAMING_SNAKE_CASE )







     @torch.no_grad()
     def   _SCREAMING_SNAKE_CASE     ( self      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    int   ,						_SCREAMING_SNAKE_CASE      :    Any ):






           """simple docstring"""




           UpperCamelCase					       =  self.vision_model(_SCREAMING_SNAKE_CASE )[1]  # pooled_output
           UpperCamelCase					       =  self.visual_projection(_SCREAMING_SNAKE_CASE )

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

           UpperCamelCase					       =  []
           UpperCamelCase					       =  image_embeds.shape[0]
           for i in range(_SCREAMING_SNAKE_CASE ):
                 UpperCamelCase					       =  {'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
                 UpperCamelCase					       =  0.0

                 for concept_idx in range(len(special_cos_dist[0] ) ):
                       UpperCamelCase					       =  special_cos_dist[i][concept_idx]
                       UpperCamelCase					       =  self.special_care_embeds_weights[concept_idx].item()
                       UpperCamelCase					       =  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]} )
                             UpperCamelCase					       =  0.0_1

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

                 result.append(_SCREAMING_SNAKE_CASE )

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

           return images, has_nsfw_concepts







     @torch.no_grad()
     def   _SCREAMING_SNAKE_CASE     ( self      :    int   ,						_SCREAMING_SNAKE_CASE      :    torch.FloatTensor   ,						_SCREAMING_SNAKE_CASE      :    torch.FloatTensor ):






           """simple docstring"""




           UpperCamelCase					       =  self.vision_model(_SCREAMING_SNAKE_CASE )[1]  # pooled_output
           UpperCamelCase					       =  self.visual_projection(_SCREAMING_SNAKE_CASE )

           UpperCamelCase					       =  cosine_distance(_SCREAMING_SNAKE_CASE   ,						self.special_care_embeds )
           UpperCamelCase					       =  cosine_distance(_SCREAMING_SNAKE_CASE   ,						self.concept_embeds )

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

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

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

           return images, has_nsfw_concepts











class 					A__     :

     '''simple docstring'''






     def __init__( self      :    Optional[int]   ,						_SCREAMING_SNAKE_CASE      :    str = ""   ,						_SCREAMING_SNAKE_CASE      :    bool = False ):






           """simple docstring"""




           UpperCamelCase					       =  {}

           # A node will be a leaf if the tree contains its word
           UpperCamelCase					       =  is_leaf

           UpperCamelCase					       =  prefix







     def   _SCREAMING_SNAKE_CASE     ( self      :    Optional[int]   ,						_SCREAMING_SNAKE_CASE      :    str ):






           """simple docstring"""




           UpperCamelCase					       =  0
           for q, w in zip(self.prefix   ,						_SCREAMING_SNAKE_CASE ):
                 if q != w:
                       break

                 x += 1

           return self.prefix[:x], self.prefix[x:], word[x:]







     def   _SCREAMING_SNAKE_CASE     ( self      :    Any   ,						_SCREAMING_SNAKE_CASE      :    list[str] ):






           """simple docstring"""




           for word in words:
                 self.insert(_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    Union[str, Any]   ,						_SCREAMING_SNAKE_CASE      :    str ):






           """simple docstring"""




           if self.prefix == word:
                 UpperCamelCase					       =  True

           # Case 2: The node has no edges that have a prefix to the word
           # Solution: We create an edge from the current node to a new one
           # containing the word
           elif word[0] not in self.nodes:
                 UpperCamelCase					       =  RadixNode(prefix=_SCREAMING_SNAKE_CASE   ,						is_leaf=_SCREAMING_SNAKE_CASE )

           else:
                 UpperCamelCase					       =  self.nodes[word[0]]
                 UpperCamelCase       ,  UpperCamelCase       ,  UpperCamelCase					       =  incoming_node.match(
                     _SCREAMING_SNAKE_CASE )

                 # Case 3: The node prefix is equal to the matching
                 # Solution: We insert remaining word on the next node
                 if remaining_prefix == "":
                       self.nodes[matching_string[0]].insert(_SCREAMING_SNAKE_CASE )

                 # Case 4: The word is greater equal to the matching
                 # Solution: Create a node in between both nodes, change
                 # prefixes and add the new node for the remaining word
                 else:
                       UpperCamelCase					       =  remaining_prefix

                       UpperCamelCase					       =  self.nodes[matching_string[0]]
                       UpperCamelCase					       =  RadixNode(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )
                       UpperCamelCase					       =  aux_node

                       if remaining_word == "":
                             UpperCamelCase					       =  True
                       else:
                             self.nodes[matching_string[0]].insert(_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    Any   ,						_SCREAMING_SNAKE_CASE      :    str ):






           """simple docstring"""




           UpperCamelCase					       =  self.nodes.get(word[0]   ,						_SCREAMING_SNAKE_CASE )
           if not incoming_node:
                 return False
           else:
                 UpperCamelCase       ,  UpperCamelCase       ,  UpperCamelCase					       =  incoming_node.match(
                     _SCREAMING_SNAKE_CASE )
                 # If there is remaining prefix, the word can't be on the tree
                 if remaining_prefix != "":
                       return False
                 # This applies when the word and the prefix are equal
                 elif remaining_word == "":
                       return incoming_node.is_leaf
                 # We have word remaining so we check the next node
                 else:
                       return incoming_node.find(_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    str ):






           """simple docstring"""




           UpperCamelCase					       =  self.nodes.get(word[0]   ,						_SCREAMING_SNAKE_CASE )
           if not incoming_node:
                 return False
           else:
                 UpperCamelCase       ,  UpperCamelCase       ,  UpperCamelCase					       =  incoming_node.match(
                     _SCREAMING_SNAKE_CASE )
                 # If there is remaining prefix, the word can't be on the tree
                 if remaining_prefix != "":
                       return False
                 # We have word remaining so we check the next node
                 elif remaining_word != "":
                       return incoming_node.delete(_SCREAMING_SNAKE_CASE )
                 else:
                       # If it is not a leaf, we don't have to delete
                       if not incoming_node.is_leaf:
                             return False
                       else:
                             # We delete the nodes if no edges go from it
                             if len(incoming_node.nodes ) == 0:
                                   del self.nodes[word[0]]
                                   # We merge the current node with its only child
                                   if len(self.nodes ) == 1 and not self.is_leaf:
                                         UpperCamelCase					       =  list(self.nodes.values() )[0]
                                         UpperCamelCase					       =  merging_node.is_leaf
                                         self.prefix += merging_node.prefix
                                         UpperCamelCase					       =  merging_node.nodes
                    # If there is more than 1 edge, we just mark it as non-leaf
                             elif len(incoming_node.nodes ) > 1:
                                   UpperCamelCase					       =  False
                             # If there is 1 edge, we merge it with its child
                             else:
                                   UpperCamelCase					       =  list(incoming_node.nodes.values() )[0]
                                   UpperCamelCase					       =  merging_node.is_leaf
                                   incoming_node.prefix += merging_node.prefix
                                   UpperCamelCase					       =  merging_node.nodes

                             return True







     def   _SCREAMING_SNAKE_CASE     ( self      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    int = 0 ):






           """simple docstring"""




           if self.prefix != "":
                 print('-' * height   ,						self.prefix   ,						'  (leaf)' if self.is_leaf else '' )

           for value in self.nodes.values():
                 value.print_tree(height + 1 )

def 		lowercase__  (				) ->			bool:
      """simple docstring"""

      UpperCamelCase					       =  'banana bananas bandana band apple all beast'.split()
      UpperCamelCase					       =  RadixNode()
      root.insert_many(_UpperCamelCase)

      assert all(root.find(_UpperCamelCase) for word in words)
      assert not root.find('bandanas')
      assert not root.find('apps')
      root.delete('all')
      assert not root.find('all')
      root.delete('banana')
      assert not root.find('banana')
      assert root.find('bananas')

      return True
def 		lowercase__  (				) ->			None:
      """simple docstring"""

      assert test_trie()
def 		lowercase__  (				) ->			None:
      """simple docstring"""

      UpperCamelCase					       =  RadixNode()
      UpperCamelCase					       =  'banana bananas bandanas bandana band apple all beast'.split()
      root.insert_many(_UpperCamelCase)

      print('Words:'       ,							_UpperCamelCase)
      print('Tree:')
      root.print_tree()


if __name__ == "__main__":
       main()



