






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

from ...utils import logging


__snake_case     :     Dict       =  logging.get_logger(__name__)







def 						_UpperCamelCase (     UpperCamelCase_							:       Dict     ,							UpperCamelCase_							:       List[Any]		)	-> Optional[Any]:

						"""simple docstring"""



						lowerCAmelCase__         =    nn.functional.normalize(UpperCamelCase_		)
						lowerCAmelCase__         =    nn.functional.normalize(UpperCamelCase_		)
						return torch.mm(UpperCamelCase_     ,							normalized_text_embeds.t()		)


class      __SCREAMING_SNAKE_CASE       ( __lowercase):
						_SCREAMING_SNAKE_CASE						:							Any				=       CLIPConfig

						_SCREAMING_SNAKE_CASE						:							str				=       ['''CLIPEncoderLayer''']






						def __init__(						self  ,    _UpperCamelCase ):






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

												lowerCAmelCase__         =    CLIPVisionModel(config.vision_config )
												lowerCAmelCase__         =    nn.Linear(config.vision_config.hidden_size  ,    config.projection_dim  ,    bias=_UpperCamelCase )

												lowerCAmelCase__         =    nn.Parameter(torch.ones(17  ,    config.projection_dim )  ,    requires_grad=_UpperCamelCase )
												lowerCAmelCase__         =    nn.Parameter(torch.ones(3  ,    config.projection_dim )  ,    requires_grad=_UpperCamelCase )

												lowerCAmelCase__         =    nn.Parameter(torch.ones(17 )  ,    requires_grad=_UpperCamelCase )
												lowerCAmelCase__         =    nn.Parameter(torch.ones(3 )  ,    requires_grad=_UpperCamelCase )






						@torch.no_grad()
						def    UpperCamelCase__   (						self  ,    _UpperCamelCase  ,    _UpperCamelCase ):






												"""simple docstring"""
												lowerCAmelCase__         =    self.vision_model(_UpperCamelCase )[1]  # pooled_output
												lowerCAmelCase__         =    self.visual_projection(_UpperCamelCase )

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

												lowerCAmelCase__         =    []
												lowerCAmelCase__         =    image_embeds.shape[0]
												for i in range(_UpperCamelCase ):
																		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(_UpperCamelCase )

																		result.append(_UpperCamelCase )

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

												for idx, has_nsfw_concept in enumerate(_UpperCamelCase ):
																		if has_nsfw_concept:
																								if torch.is_tensor(_UpperCamelCase ) or torch.is_tensor(images[0] ):
																														lowerCAmelCase__         =    torch.zeros_like(images[idx] )  # black image
																								else:
																														lowerCAmelCase__         =    np.zeros(images[idx].shape )  # black image

												if any(_UpperCamelCase ):
																		logger.warning(
																		    'Potential NSFW content was detected in one or more images. A black image will be returned instead.'
																		    ' Try again with a different prompt and/or seed.' )

												return images, has_nsfw_concepts







						@torch.no_grad()
						def    UpperCamelCase__   (						self  ,    _UpperCamelCase  ,    _UpperCamelCase ):






												"""simple docstring"""
												lowerCAmelCase__         =    self.vision_model(_UpperCamelCase )[1]  # pooled_output
												lowerCAmelCase__         =    self.visual_projection(_UpperCamelCase )

												lowerCAmelCase__         =    cosine_distance(_UpperCamelCase  ,    self.special_care_embeds )
												lowerCAmelCase__         =    cosine_distance(_UpperCamelCase  ,    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 )

												lowerCAmelCase__         =    0.0  # black image

												return images, has_nsfw_concepts













from transformers import (
    DPRContextEncoderTokenizer,
    DPRContextEncoderTokenizerFast,
    DPRQuestionEncoderTokenizer,
    DPRQuestionEncoderTokenizerFast,
    DPRReaderOutput,
    DPRReaderTokenizer,
    DPRReaderTokenizerFast,
)
from transformers.testing_utils import require_tokenizers, slow
from transformers.tokenization_utils_base import BatchEncoding

from ..bert.test_tokenization_bert import BertTokenizationTest


@require_tokenizers
class      __SCREAMING_SNAKE_CASE       ( __lowercase):
						_SCREAMING_SNAKE_CASE						:							Tuple				=       DPRContextEncoderTokenizer
						_SCREAMING_SNAKE_CASE						:							Any				=       DPRContextEncoderTokenizerFast
						_SCREAMING_SNAKE_CASE						:							str				=       True




@require_tokenizers
class      __SCREAMING_SNAKE_CASE       ( __lowercase):
						_SCREAMING_SNAKE_CASE						:							Optional[Any]				=       DPRQuestionEncoderTokenizer
						_SCREAMING_SNAKE_CASE						:							List[Any]				=       DPRQuestionEncoderTokenizerFast
						_SCREAMING_SNAKE_CASE						:							Optional[Any]				=       True




@require_tokenizers
class      __SCREAMING_SNAKE_CASE       ( __lowercase):
						_SCREAMING_SNAKE_CASE						:							Any				=       DPRReaderTokenizer
						_SCREAMING_SNAKE_CASE						:							int				=       DPRReaderTokenizerFast
						_SCREAMING_SNAKE_CASE						:							List[str]				=       True






						@slow
						def    UpperCamelCase__   (						self ):






												"""simple docstring"""
												lowerCAmelCase__         =    self.tokenizer_class.from_pretrained('bert-base-uncased' )

												lowerCAmelCase__         =    tokenizer.encode('question sequence'  ,    add_special_tokens=_UpperCamelCase )
												lowerCAmelCase__         =    tokenizer.encode('title sequence'  ,    add_special_tokens=_UpperCamelCase )
												lowerCAmelCase__         =    tokenizer.encode('text sequence ' * 4  ,    add_special_tokens=_UpperCamelCase )
												lowerCAmelCase__         =    [[1_01] + text_a + [1_02] + text_a + [1_02] + text_a]
												lowerCAmelCase__         =    BatchEncoding({'input_ids': input_ids} )

												lowerCAmelCase__         =    [[0] * len(input_ids[0] )]
												lowerCAmelCase__         =    [[0] * len(input_ids[0] )]
												lowerCAmelCase__         =    [0]
												lowerCAmelCase__         =    DPRReaderOutput(_UpperCamelCase  ,    _UpperCamelCase  ,    _UpperCamelCase )

												lowerCAmelCase__    ,							lowerCAmelCase__         =    8, 9
												lowerCAmelCase__         =    10
												lowerCAmelCase__         =    10
												lowerCAmelCase__         =    tokenizer.decode_best_spans(_UpperCamelCase  ,    _UpperCamelCase )
												self.assertEqual(predicted_spans[0].start_index  ,    _UpperCamelCase )
												self.assertEqual(predicted_spans[0].end_index  ,    _UpperCamelCase )
												self.assertEqual(predicted_spans[0].doc_id  ,    0 )







						@slow
						def    UpperCamelCase__   (						self ):






												"""simple docstring"""
												lowerCAmelCase__         =    self.tokenizer_class.from_pretrained('bert-base-uncased' )

												lowerCAmelCase__         =    tokenizer.encode('question sequence'  ,    add_special_tokens=_UpperCamelCase )
												lowerCAmelCase__         =    tokenizer.encode('title sequence'  ,    add_special_tokens=_UpperCamelCase )
												lowerCAmelCase__         =    tokenizer.encode('text sequence'  ,    add_special_tokens=_UpperCamelCase )
												lowerCAmelCase__         =    [1_01] + text_a + [1_02] + text_a + [1_02] + text_a
												lowerCAmelCase__         =    tokenizer(questions=['question sequence']  ,    titles=['title sequence']  ,    texts=['text sequence'] )
												self.assertIn('input_ids'  ,    _UpperCamelCase )
												self.assertIn('attention_mask'  ,    _UpperCamelCase )
												self.assertListEqual(encoded_input['input_ids'][0]  ,    _UpperCamelCase )





