



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

from ...utils import logging


__UpperCAmelCase			  =			logging.get_logger(__name__)
def 						UpperCamelCase      (  snake_case__			:       Optional[int]			,	snake_case__			:       int      )      ->	int:
   UpperCamelCase       :    Any		      =      nn.functional.normalize(snake_case__      )
   UpperCamelCase       :    List[Any]		      =      nn.functional.normalize(snake_case__      )
   return torch.mm(snake_case__			,	normalized_text_embeds.t()      )



class 							lowerCAmelCase_						(    a__				):
       UpperCAmelCase__							:   Optional[int]  	=    CLIPConfig

       UpperCAmelCase__							:   Union[str, Any]  	=    ["CLIPEncoderLayer"]
       def __init__(     self,     SCREAMING_SNAKE_CASE_			)	->						List[Any]:
          super().__init__(SCREAMING_SNAKE_CASE_			)

          UpperCamelCase       :    List[Any]		      =      CLIPVisionModel(config.vision_config			)
          UpperCamelCase       :    str		      =      nn.Linear(config.vision_config.hidden_size,     config.projection_dim,     bias=SCREAMING_SNAKE_CASE_			)

          UpperCamelCase       :    Any		      =      nn.Parameter(torch.ones(17,     config.projection_dim			),     requires_grad=SCREAMING_SNAKE_CASE_			)
          UpperCamelCase       :    Tuple		      =      nn.Parameter(torch.ones(3,     config.projection_dim			),     requires_grad=SCREAMING_SNAKE_CASE_			)

          UpperCamelCase       :    Dict		      =      nn.Parameter(torch.ones(17			),     requires_grad=SCREAMING_SNAKE_CASE_			)
          UpperCamelCase       :    Any		      =      nn.Parameter(torch.ones(3			),     requires_grad=SCREAMING_SNAKE_CASE_			)
       @torch.no_grad()
       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			)	->						str:
          UpperCamelCase       :    List[str]		      =      self.vision_model(SCREAMING_SNAKE_CASE_			)[1]  # pooled_output
          UpperCamelCase       :    str		      =      self.visual_projection(SCREAMING_SNAKE_CASE_			)

          # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
          UpperCamelCase       :    Tuple		      =      cosine_distance(SCREAMING_SNAKE_CASE_,     self.special_care_embeds			).cpu().float().numpy()
          UpperCamelCase       :    Union[str, Any]		      =      cosine_distance(SCREAMING_SNAKE_CASE_,     self.concept_embeds			).cpu().float().numpy()

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

             for concept_idx in range(len(special_cos_dist[0]			)			):
                UpperCamelCase       :    Any		      =      special_cos_dist[i][concept_idx]
                UpperCamelCase       :    Optional[int]		      =      self.special_care_embeds_weights[concept_idx].item()
                UpperCamelCase       :    Optional[Any]		      =      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       :    List[Any]		      =      0.01

             for concept_idx in range(len(cos_dist[0]			)			):
                UpperCamelCase       :    List[str]		      =      cos_dist[i][concept_idx]
                UpperCamelCase       :    List[Any]		      =      self.concept_embeds_weights[concept_idx].item()
                UpperCamelCase       :    Optional[int]		      =      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       :    List[str]		      =      [len(res['bad_concepts']			) > 0 for res in result]

          return images, has_nsfw_concepts






       @torch.no_grad()
       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			)	->						Tuple:
          UpperCamelCase       :    int		      =      self.vision_model(SCREAMING_SNAKE_CASE_			)[1]  # pooled_output
          UpperCamelCase       :    Dict		      =      self.visual_projection(SCREAMING_SNAKE_CASE_			)

          UpperCamelCase       :    Union[str, Any]		      =      cosine_distance(SCREAMING_SNAKE_CASE_,     self.special_care_embeds			)
          UpperCamelCase       :    Optional[int]		      =      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       :    Optional[Any]		      =      0.0

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

          UpperCamelCase       :    Dict		      =      (cos_dist - self.concept_embeds_weights) + special_adjustment
          # concept_scores = concept_scores.round(decimals=3)
          UpperCamelCase       :    Optional[Any]		      =      torch.any(concept_scores > 0,     dim=1			)

          return images, has_nsfw_concepts









from collections import OrderedDict
from typing import TYPE_CHECKING, Any, List, Mapping, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import TensorType, logging


if TYPE_CHECKING:
						from ...onnx.config import PatchingSpec
						from ...tokenization_utils_base import PreTrainedTokenizerBase


SCREAMING_SNAKE_CASE:     str									= logging.get_logger(__name__)

SCREAMING_SNAKE_CASE:     Optional[int]									= {
    '''allenai/longformer-base-4096''': '''https://huggingface.co/allenai/longformer-base-4096/resolve/main/config.json''',
    '''allenai/longformer-large-4096''': '''https://huggingface.co/allenai/longformer-large-4096/resolve/main/config.json''',
    '''allenai/longformer-large-4096-finetuned-triviaqa''': (
        '''https://huggingface.co/allenai/longformer-large-4096-finetuned-triviaqa/resolve/main/config.json'''
    ),
    '''allenai/longformer-base-4096-extra.pos.embd.only''': (
        '''https://huggingface.co/allenai/longformer-base-4096-extra.pos.embd.only/resolve/main/config.json'''
    ),
    '''allenai/longformer-large-4096-extra.pos.embd.only''': (
        '''https://huggingface.co/allenai/longformer-large-4096-extra.pos.embd.only/resolve/main/config.json'''
    ),
}



class    lowercase_		(SCREAMING_SNAKE_CASE__ ):
						lowerCAmelCase__			      ="longformer"



						def __init__(			self	:   Optional[Any]					,							snake_case__	:   Union[List[int], int] = 5_12					,							snake_case__	:   int = 2					,							snake_case__	:   int = 1					,							snake_case__	:   int = 0					,							snake_case__	:   int = 2					,							snake_case__	:   int = 3_05_22					,							snake_case__	:   int = 7_68					,							snake_case__	:   int = 12					,							snake_case__	:   int = 12					,							snake_case__	:   int = 30_72					,							snake_case__	:   str = "gelu"					,							snake_case__	:   float = 0.1					,							snake_case__	:   float = 0.1					,							snake_case__	:   int = 5_12					,							snake_case__	:   int = 2					,							snake_case__	:   float = 0.02					,							snake_case__	:   float = 1e-12					,							snake_case__	:   bool = False					,							**snake_case__	:   Tuple					,							):

								"""simple docstring"""

								super().__init__(pad_token_id=snake_case__					,							**snake_case__  )

								SCREAMING_SNAKE_CASE_					      =							attention_window
								SCREAMING_SNAKE_CASE_					      =							sep_token_id
								SCREAMING_SNAKE_CASE_					      =							bos_token_id
								SCREAMING_SNAKE_CASE_					      =							eos_token_id
								SCREAMING_SNAKE_CASE_					      =							vocab_size
								SCREAMING_SNAKE_CASE_					      =							hidden_size
								SCREAMING_SNAKE_CASE_					      =							num_hidden_layers
								SCREAMING_SNAKE_CASE_					      =							num_attention_heads
								SCREAMING_SNAKE_CASE_					      =							hidden_act
								SCREAMING_SNAKE_CASE_					      =							intermediate_size
								SCREAMING_SNAKE_CASE_					      =							hidden_dropout_prob
								SCREAMING_SNAKE_CASE_					      =							attention_probs_dropout_prob
								SCREAMING_SNAKE_CASE_					      =							max_position_embeddings
								SCREAMING_SNAKE_CASE_					      =							type_vocab_size
								SCREAMING_SNAKE_CASE_					      =							initializer_range
								SCREAMING_SNAKE_CASE_					      =							layer_norm_eps
								SCREAMING_SNAKE_CASE_					      =							onnx_export



class    lowercase_		(SCREAMING_SNAKE_CASE__ ):



						def __init__(			self	:   Dict					,							snake_case__	:   "PretrainedConfig"					,							snake_case__	:   str = "default"					,							snake_case__	:   "List[PatchingSpec]" = None  ):

								"""simple docstring"""

								super().__init__(snake_case__					,							snake_case__					,							snake_case__  )
								SCREAMING_SNAKE_CASE_					      =							True



						@property
						def 	__a					(			self	:   Union[str, Any]  ):

								"""simple docstring"""

								if self.task == "multiple-choice":
										SCREAMING_SNAKE_CASE_					      =							{0: 'batch', 1: 'choice', 2: 'sequence'}
								else:
										SCREAMING_SNAKE_CASE_					      =							{0: 'batch', 1: 'sequence'}
								return OrderedDict(
								    [
								        ('input_ids', dynamic_axis),
								        ('attention_mask', dynamic_axis),
								        ('global_attention_mask', dynamic_axis),
								    ]  )



						@property
						def 	__a					(			self	:   Optional[int]  ):

								"""simple docstring"""

								SCREAMING_SNAKE_CASE_					      =							super().outputs
								if self.task == "default":
										SCREAMING_SNAKE_CASE_					      =							{0: 'batch'}
								return outputs



						@property
						def 	__a					(			self	:   Optional[int]  ):

								"""simple docstring"""

								return 1e-4



						@property
						def 	__a					(			self	:   List[str]  ):

								"""simple docstring"""

								return max(super().default_onnx_opset					,							14  )






						def 	__a					(			self	:   List[str]					,							snake_case__	:   "PreTrainedTokenizerBase"					,							snake_case__	:   int = -1					,							snake_case__	:   int = -1					,							snake_case__	:   bool = False					,							snake_case__	:   Optional[TensorType] = None					,							):

								"""simple docstring"""

								SCREAMING_SNAKE_CASE_					      =							super().generate_dummy_inputs(
								    preprocessor=snake_case__					,							batch_size=snake_case__					,							seq_length=snake_case__					,							is_pair=snake_case__					,							framework=snake_case__  )
								import torch

								# for some reason, replacing this code by inputs["global_attention_mask"] = torch.randint(2, inputs["input_ids"].shape, dtype=torch.int64)
								# makes the export fail randomly
								SCREAMING_SNAKE_CASE_					      =							torch.zeros_like(inputs['input_ids']  )
								# make every second token global
								SCREAMING_SNAKE_CASE_					      =							1

								return inputs