
from __future__ import annotations

import unittest

from transformers import MobileBertConfig, is_tf_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
							import tensorflow as tf

							from transformers import (
							    TF_MODEL_FOR_PRETRAINING_MAPPING,
							    TFMobileBertForMaskedLM,
							    TFMobileBertForMultipleChoice,
							    TFMobileBertForNextSentencePrediction,
							    TFMobileBertForPreTraining,
							    TFMobileBertForQuestionAnswering,
							    TFMobileBertForSequenceClassification,
							    TFMobileBertForTokenClassification,
							    TFMobileBertModel,
							)







@require_tf
class        lowercase	(      SCREAMING_SNAKE_CASE__ ,						SCREAMING_SNAKE_CASE__ ,						unittest.TestCase	):
							lowercase_ : Dict       =(
							    (
							        TFMobileBertModel,
							        TFMobileBertForMaskedLM,
							        TFMobileBertForNextSentencePrediction,
							        TFMobileBertForPreTraining,
							        TFMobileBertForQuestionAnswering,
							        TFMobileBertForSequenceClassification,
							        TFMobileBertForTokenClassification,
							        TFMobileBertForMultipleChoice,
							    )
							    if is_tf_available()
							    else ()
							)
							lowercase_ : Dict       =(
							    {
							        '''feature-extraction''': TFMobileBertModel,
							        '''fill-mask''': TFMobileBertForMaskedLM,
							        '''question-answering''': TFMobileBertForQuestionAnswering,
							        '''text-classification''': TFMobileBertForSequenceClassification,
							        '''token-classification''': TFMobileBertForTokenClassification,
							        '''zero-shot''': TFMobileBertForSequenceClassification,
							    }
							    if is_tf_available()
							    else {}
							)
							lowercase_ : Tuple       =False
							lowercase_ : Any       =False




							def 				A__   (       self     ,A__     ,A__     ,A__=False):
												lowercase					=		super()._prepare_for_class(A__     ,A__     ,return_labels=A__)

												if return_labels:
																	if model_class in get_values(A__):
																						lowercase					=		tf.zeros(self.model_tester.batch_size     ,dtype=tf.intaa)

												return inputs_dict







							class        lowercase	(      SCREAMING_SNAKE_CASE__	):




														def __init__(       self     ,A__     ,A__=1_3     ,A__=7     ,A__=True     ,A__=True     ,A__=True     ,A__=True     ,A__=9_9     ,A__=3_2     ,A__=3_2     ,A__=2     ,A__=4     ,A__=3_7     ,A__="gelu"     ,A__=0.1     ,A__=0.1     ,A__=5_1_2     ,A__=1_6     ,A__=2     ,A__=0.02     ,A__=3     ,A__=4     ,A__=None     ,):
																			lowercase					=		parent
																			lowercase					=		batch_size
																			lowercase					=		seq_length
																			lowercase					=		is_training
																			lowercase					=		use_input_mask
																			lowercase					=		use_token_type_ids
																			lowercase					=		use_labels
																			lowercase					=		vocab_size
																			lowercase					=		hidden_size
																			lowercase					=		num_hidden_layers
																			lowercase					=		num_attention_heads
																			lowercase					=		intermediate_size
																			lowercase					=		hidden_act
																			lowercase					=		hidden_dropout_prob
																			lowercase					=		attention_probs_dropout_prob
																			lowercase					=		max_position_embeddings
																			lowercase					=		type_vocab_size
																			lowercase					=		type_sequence_label_size
																			lowercase					=		initializer_range
																			lowercase					=		num_labels
																			lowercase					=		num_choices
																			lowercase					=		scope
																			lowercase					=		embedding_size




														def 				A__   (       self):
																			lowercase					=		ids_tensor([self.batch_size, self.seq_length]     ,self.vocab_size)

																			lowercase					=		None
																			if self.use_input_mask:
																								lowercase					=		random_attention_mask([self.batch_size, self.seq_length])

																			lowercase					=		None
																			if self.use_token_type_ids:
																								lowercase					=		ids_tensor([self.batch_size, self.seq_length]     ,self.type_vocab_size)

																			lowercase					=		None
																			lowercase					=		None
																			lowercase					=		None
																			if self.use_labels:
																								lowercase					=		ids_tensor([self.batch_size]     ,self.type_sequence_label_size)
																								lowercase					=		ids_tensor([self.batch_size, self.seq_length]     ,self.num_labels)
																								lowercase					=		ids_tensor([self.batch_size]     ,self.num_choices)

																			lowercase					=		MobileBertConfig(
																			    vocab_size=self.vocab_size     ,hidden_size=self.hidden_size     ,num_hidden_layers=self.num_hidden_layers     ,num_attention_heads=self.num_attention_heads     ,intermediate_size=self.intermediate_size     ,hidden_act=self.hidden_act     ,hidden_dropout_prob=self.hidden_dropout_prob     ,attention_probs_dropout_prob=self.attention_probs_dropout_prob     ,max_position_embeddings=self.max_position_embeddings     ,type_vocab_size=self.type_vocab_size     ,initializer_range=self.initializer_range     ,embedding_size=self.embedding_size     ,)

																			return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels




														def 				A__   (       self     ,A__     ,A__     ,A__     ,A__     ,A__     ,A__     ,A__):
																			lowercase					=		TFMobileBertModel(config=A__)
																			lowercase					=		{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
																			lowercase					=		model(A__)

																			lowercase					=		[input_ids, input_mask]
																			lowercase					=		model(A__)

																			lowercase					=		model(A__)

																			self.parent.assertEqual(
																			    result.last_hidden_state.shape     ,(self.batch_size, self.seq_length, self.hidden_size))
																			self.parent.assertEqual(result.pooler_output.shape     ,(self.batch_size, self.hidden_size))




														def 				A__   (       self     ,A__     ,A__     ,A__     ,A__     ,A__     ,A__     ,A__):
																			lowercase					=		TFMobileBertForMaskedLM(config=A__)
																			lowercase					=		{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
																			lowercase					=		model(A__)
																			self.parent.assertEqual(result.logits.shape     ,(self.batch_size, self.seq_length, self.vocab_size))




														def 				A__   (       self     ,A__     ,A__     ,A__     ,A__     ,A__     ,A__     ,A__):
																			lowercase					=		TFMobileBertForNextSentencePrediction(config=A__)
																			lowercase					=		{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
																			lowercase					=		model(A__)
																			self.parent.assertEqual(result.logits.shape     ,(self.batch_size, 2))




														def 				A__   (       self     ,A__     ,A__     ,A__     ,A__     ,A__     ,A__     ,A__):
																			lowercase					=		TFMobileBertForPreTraining(config=A__)
																			lowercase					=		{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
																			lowercase					=		model(A__)
																			self.parent.assertEqual(
																			    result.prediction_logits.shape     ,(self.batch_size, self.seq_length, self.vocab_size))
																			self.parent.assertEqual(result.seq_relationship_logits.shape     ,(self.batch_size, 2))




														def 				A__   (       self     ,A__     ,A__     ,A__     ,A__     ,A__     ,A__     ,A__):
																			lowercase					=		self.num_labels
																			lowercase					=		TFMobileBertForSequenceClassification(config=A__)
																			lowercase					=		{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
																			lowercase					=		model(A__)
																			self.parent.assertEqual(result.logits.shape     ,(self.batch_size, self.num_labels))




														def 				A__   (       self     ,A__     ,A__     ,A__     ,A__     ,A__     ,A__     ,A__):
																			lowercase					=		self.num_choices
																			lowercase					=		TFMobileBertForMultipleChoice(config=A__)
																			lowercase					=		tf.tile(tf.expand_dims(A__     ,1)     ,(1, self.num_choices, 1))
																			lowercase					=		tf.tile(tf.expand_dims(A__     ,1)     ,(1, self.num_choices, 1))
																			lowercase					=		tf.tile(tf.expand_dims(A__     ,1)     ,(1, self.num_choices, 1))
																			lowercase					=		{
																			    '''input_ids''': multiple_choice_inputs_ids,
																			    '''attention_mask''': multiple_choice_input_mask,
																			    '''token_type_ids''': multiple_choice_token_type_ids,
																			}
																			lowercase					=		model(A__)
																			self.parent.assertEqual(result.logits.shape     ,(self.batch_size, self.num_choices))




														def 				A__   (       self     ,A__     ,A__     ,A__     ,A__     ,A__     ,A__     ,A__):
																			lowercase					=		self.num_labels
																			lowercase					=		TFMobileBertForTokenClassification(config=A__)
																			lowercase					=		{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
																			lowercase					=		model(A__)
																			self.parent.assertEqual(result.logits.shape     ,(self.batch_size, self.seq_length, self.num_labels))




														def 				A__   (       self     ,A__     ,A__     ,A__     ,A__     ,A__     ,A__     ,A__):
																			lowercase					=		TFMobileBertForQuestionAnswering(config=A__)
																			lowercase					=		{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
																			lowercase					=		model(A__)
																			self.parent.assertEqual(result.start_logits.shape     ,(self.batch_size, self.seq_length))
																			self.parent.assertEqual(result.end_logits.shape     ,(self.batch_size, self.seq_length))





														def 				A__   (       self):
																			lowercase					=		self.prepare_config_and_inputs()
																			(
																			    (
																			    lowercase
																			)			,		(
																			    lowercase
																			)			,		(
																			    lowercase
																			)			,		(
																			    lowercase
																			)			,		(
																			    lowercase
																			)			,		(
																			    lowercase
																			)			,		(
																			    lowercase
																			)			,		
																			)					=		config_and_inputs
																			lowercase					=		{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
																			return config, inputs_dict




							def 				A__   (       self):
												lowercase					=		TFMobileBertModelTest.TFMobileBertModelTester(self)
												lowercase					=		ConfigTester(self     ,config_class=A__     ,hidden_size=3_7)




							def 				A__   (       self):
												self.config_tester.run_common_tests()




							def 				A__   (       self):
												lowercase					=		self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_mobilebert_model(*A__)




							def 				A__   (       self):
												lowercase					=		self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_mobilebert_for_masked_lm(*A__)




							def 				A__   (       self):
												lowercase					=		self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_mobilebert_for_multiple_choice(*A__)




							def 				A__   (       self):
												lowercase					=		self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*A__)




							def 				A__   (       self):
												lowercase					=		self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_mobilebert_for_pretraining(*A__)




							def 				A__   (       self):
												lowercase					=		self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_mobilebert_for_question_answering(*A__)




							def 				A__   (       self):
												lowercase					=		self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_mobilebert_for_sequence_classification(*A__)




							def 				A__   (       self):
												lowercase					=		self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_mobilebert_for_token_classification(*A__)





							@slow
							def 				A__   (       self):
												# for model_name in TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												for model_name in ["google/mobilebert-uncased"]:
																	lowercase					=		TFMobileBertModel.from_pretrained(A__)
																	self.assertIsNotNone(A__)







@require_tf
class        lowercase	(      unittest.TestCase	):




							@slow
							def 				A__   (       self):
												lowercase					=		TFMobileBertForPreTraining.from_pretrained('''google/mobilebert-uncased''')
												lowercase					=		tf.constant([[0, 1, 2, 3, 4, 5]])
												lowercase					=		model(A__)[0]

												lowercase					=		[1, 6, 3_0_5_2_2]
												self.assertEqual(output.shape     ,A__)

												lowercase					=		tf.constant(
												    [
												        [
												            [-4.5919547, -9.248295, -9.645256],
												            [-6.7306175, -6.440284, -6.6052837],
												            [-7.2743506, -6.7847915, -6.024673],
												        ]
												    ])
												tf.debugging.assert_near(output[:, :3, :3]     ,A__     ,atol=1E-4)






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

from ...utils import logging


lowercase__	:Union[str, Any]										=					logging.get_logger(__name__)







class        lowercase	(      SCREAMING_SNAKE_CASE__	):
							lowercase_ : Union[str, Any]       =CLIPConfig

							lowercase_ : str       =['''CLIPEncoderLayer''']




							def __init__(       self     ,A__):
												super().__init__(A__)

												lowercase					=		CLIPVisionModelWithProjection(config.vision_config)

												lowercase					=		nn.Linear(config.vision_config.projection_dim     ,1)
												lowercase					=		nn.Linear(config.vision_config.projection_dim     ,1)





							@torch.no_grad()
							def 				A__   (       self     ,A__     ,A__     ,A__=0.5     ,A__=0.5):
												lowercase					=		self.vision_model(A__)[0]

												lowercase					=		self.p_head(A__)
												lowercase					=		nsfw_detected.flatten()
												lowercase					=		nsfw_detected > p_threshold
												lowercase					=		nsfw_detected.tolist()

												if any(A__):
																	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.''')

												for idx, nsfw_detected_ in enumerate(A__):
																	if nsfw_detected_:
																						lowercase					=		np.zeros(images[idx].shape)

												lowercase					=		self.w_head(A__)
												lowercase					=		watermark_detected.flatten()
												lowercase					=		watermark_detected > w_threshold
												lowercase					=		watermark_detected.tolist()

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

												for idx, watermark_detected_ in enumerate(A__):
																	if watermark_detected_:
																						lowercase					=		np.zeros(images[idx].shape)

												return images, nsfw_detected, watermark_detected




