





"""simple docstring"""




import tempfile
import unittest

import numpy as np

from diffusers import (
    DDIMScheduler,
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionPipeline,
    PNDMScheduler,
)
from diffusers.utils.testing_utils import is_onnx_available, nightly, require_onnxruntime, require_torch_gpu

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
  import onnxruntime as ort


class 			__magic_name__       (       SCREAMING_SNAKE_CASE__   ,					unittest.TestCase ):
       UpperCamelCase_   =		'''hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'''
       def 							lowercase_							(     self ,				A_=0       )     ->	Union[str, Any]:





         """simple docstring"""




         _lowercase:			Union[str, Any]       =     np.random.RandomState(A_       )
         _lowercase:			Any       =     {
             '''prompt''': '''A painting of a squirrel eating a burger''',
             '''generator''': generator,
             '''num_inference_steps''': 2,
             '''guidance_scale''': 7.5,
             '''output_type''': '''numpy''',
         }
         return inputs
       def 							lowercase_							(     self       )     ->	Optional[int]:





         """simple docstring"""




         _lowercase:			Any       =     OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint ,				provider='''CPUExecutionProvider'''       )
         pipe.set_progress_bar_config(disable=A_       )

         _lowercase:			List[str]       =     self.get_dummy_inputs()
         _lowercase:			Union[str, Any]       =     pipe(**A_       ).images
         _lowercase:			int       =     image[0, -3:, -3:, -1]

         assert image.shape == (1, 128, 128, 3)
         _lowercase:			Dict       =     np.array([0.6_50_72, 0.5_84_92, 0.4_82_19, 0.5_55_21, 0.5_31_80, 0.5_59_39, 0.5_06_97, 0.3_98_00, 0.4_64_55]       )

         assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-2
       def 							lowercase_							(     self       )     ->	str:





         """simple docstring"""




         _lowercase:			Optional[int]       =     OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint ,				provider='''CPUExecutionProvider'''       )
         _lowercase:			Any       =     PNDMScheduler.from_config(pipe.scheduler.config ,				skip_prk_steps=A_       )
         pipe.set_progress_bar_config(disable=A_       )

         _lowercase:			Tuple       =     self.get_dummy_inputs()
         _lowercase:			Optional[int]       =     pipe(**A_       ).images
         _lowercase:			Optional[Any]       =     image[0, -3:, -3:, -1]

         assert image.shape == (1, 128, 128, 3)
         _lowercase:			Union[str, Any]       =     np.array([0.6_58_63, 0.5_94_25, 0.4_93_26, 0.5_63_13, 0.5_38_75, 0.5_66_27, 0.5_10_65, 0.3_97_77, 0.4_63_30]       )

         assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-2
       def 							lowercase_							(     self       )     ->	Any:





         """simple docstring"""




         _lowercase:			str       =     OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint ,				provider='''CPUExecutionProvider'''       )
         _lowercase:			int       =     LMSDiscreteScheduler.from_config(pipe.scheduler.config       )
         pipe.set_progress_bar_config(disable=A_       )

         _lowercase:			List[Any]       =     self.get_dummy_inputs()
         _lowercase:			Optional[int]       =     pipe(**A_       ).images
         _lowercase:			str       =     image[0, -3:, -3:, -1]

         assert image.shape == (1, 128, 128, 3)
         _lowercase:			str       =     np.array([0.5_37_55, 0.6_07_86, 0.4_74_02, 0.4_94_88, 0.5_18_69, 0.4_98_19, 0.4_79_85, 0.3_89_57, 0.4_42_79]       )

         assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-2
       def 							lowercase_							(     self       )     ->	Dict:





         """simple docstring"""




         _lowercase:			Any       =     OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint ,				provider='''CPUExecutionProvider'''       )
         _lowercase:			Tuple       =     EulerDiscreteScheduler.from_config(pipe.scheduler.config       )
         pipe.set_progress_bar_config(disable=A_       )

         _lowercase:			int       =     self.get_dummy_inputs()
         _lowercase:			Tuple       =     pipe(**A_       ).images
         _lowercase:			Tuple       =     image[0, -3:, -3:, -1]

         assert image.shape == (1, 128, 128, 3)
         _lowercase:			List[Any]       =     np.array([0.5_37_55, 0.6_07_86, 0.4_74_02, 0.4_94_88, 0.5_18_69, 0.4_98_19, 0.4_79_85, 0.3_89_57, 0.4_42_79]       )

         assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-2
       def 							lowercase_							(     self       )     ->	Union[str, Any]:





         """simple docstring"""




         _lowercase:			Union[str, Any]       =     OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint ,				provider='''CPUExecutionProvider'''       )
         _lowercase:			Optional[int]       =     EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config       )
         pipe.set_progress_bar_config(disable=A_       )

         _lowercase:			int       =     self.get_dummy_inputs()
         _lowercase:			List[Any]       =     pipe(**A_       ).images
         _lowercase:			Optional[Any]       =     image[0, -3:, -3:, -1]

         assert image.shape == (1, 128, 128, 3)
         _lowercase:			Optional[int]       =     np.array([0.5_38_17, 0.6_08_12, 0.4_73_84, 0.4_95_30, 0.5_18_94, 0.4_98_14, 0.4_79_84, 0.3_89_58, 0.4_42_71]       )

         assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-2
       def 							lowercase_							(     self       )     ->	Any:





         """simple docstring"""




         _lowercase:			int       =     OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint ,				provider='''CPUExecutionProvider'''       )
         _lowercase:			List[Any]       =     DPMSolverMultistepScheduler.from_config(pipe.scheduler.config       )
         pipe.set_progress_bar_config(disable=A_       )

         _lowercase:			Dict       =     self.get_dummy_inputs()
         _lowercase:			List[Any]       =     pipe(**A_       ).images
         _lowercase:			List[Any]       =     image[0, -3:, -3:, -1]

         assert image.shape == (1, 128, 128, 3)
         _lowercase:			List[Any]       =     np.array([0.5_38_95, 0.6_08_08, 0.4_79_33, 0.4_96_08, 0.5_18_86, 0.4_99_50, 0.4_80_53, 0.3_89_57, 0.4_42_00]       )

         assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-2
       def 							lowercase_							(     self       )     ->	List[Any]:





         """simple docstring"""




         _lowercase:			Tuple       =     OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint ,				provider='''CPUExecutionProvider'''       )
         pipe.set_progress_bar_config(disable=A_       )

         _lowercase:			List[str]       =     self.get_dummy_inputs()
         _lowercase:			Optional[Any]       =     3 * [inputs['''prompt''']]

         # forward
         _lowercase:			int       =     pipe(**A_       )
         _lowercase:			List[str]       =     output.images[0, -3:, -3:, -1]

         _lowercase:			Union[str, Any]       =     self.get_dummy_inputs()
         _lowercase:			Any       =     3 * [inputs.pop('''prompt'''       )]

         _lowercase:			Optional[int]       =     pipe.tokenizer(
             A_ ,				padding='''max_length''' ,				max_length=pipe.tokenizer.model_max_length ,				truncation=A_ ,				return_tensors='''np''' ,				)
         _lowercase:			List[str]       =     text_inputs['''input_ids''']

         _lowercase:			Union[str, Any]       =     pipe.text_encoder(input_ids=text_inputs.astype(np.intaa       )       )[0]

         _lowercase:			Optional[int]       =     prompt_embeds

         # forward
         _lowercase:			Dict       =     pipe(**A_       )
         _lowercase:			List[str]       =     output.images[0, -3:, -3:, -1]

         assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()       ).max() < 1E-4



       def 							lowercase_							(     self       )     ->	Optional[int]:





         """simple docstring"""




         _lowercase:			Union[str, Any]       =     OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint ,				provider='''CPUExecutionProvider'''       )
         pipe.set_progress_bar_config(disable=A_       )

         _lowercase:			int       =     self.get_dummy_inputs()
         _lowercase:			Any       =     3 * ['''this is a negative prompt''']
         _lowercase:			Any       =     negative_prompt
         _lowercase:			List[Any]       =     3 * [inputs['''prompt''']]

         # forward
         _lowercase:			Dict       =     pipe(**A_       )
         _lowercase:			Any       =     output.images[0, -3:, -3:, -1]

         _lowercase:			Optional[Any]       =     self.get_dummy_inputs()
         _lowercase:			str       =     3 * [inputs.pop('''prompt'''       )]

         _lowercase:			List[Any]       =     []
         for p in [prompt, negative_prompt]:
           _lowercase:			Any       =     pipe.tokenizer(
               A_ ,				padding='''max_length''' ,				max_length=pipe.tokenizer.model_max_length ,				truncation=A_ ,				return_tensors='''np''' ,				)
           _lowercase:			int       =     text_inputs['''input_ids''']

           embeds.append(pipe.text_encoder(input_ids=text_inputs.astype(np.intaa       )       )[0]       )

         _lowercase:			int       =     embeds

         # forward
         _lowercase:			Any       =     pipe(**A_       )
         _lowercase:			Any       =     output.images[0, -3:, -3:, -1]

         assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()       ).max() < 1E-4


@nightly
@require_onnxruntime
@require_torch_gpu
class 			__magic_name__       (       unittest.TestCase ):
       @property
       def 							lowercase_							(     self       )     ->	str:





         """simple docstring"""




         return (
             "CUDAExecutionProvider",
             {
                 "gpu_mem_limit": "15000000000",  # 15GB
                 "arena_extend_strategy": "kSameAsRequested",
             },
         )
       @property
       def 							lowercase_							(     self       )     ->	Optional[int]:





         """simple docstring"""




         _lowercase:			str       =     ort.SessionOptions()
         _lowercase:			Tuple       =     False
         return options
       def 							lowercase_							(     self       )     ->	Optional[Any]:





         """simple docstring"""




         _lowercase:			Dict       =     OnnxStableDiffusionPipeline.from_pretrained(
             '''CompVis/stable-diffusion-v1-4''' ,				revision='''onnx''' ,				safety_checker=A_ ,				feature_extractor=A_ ,				provider=self.gpu_provider ,				sess_options=self.gpu_options ,				)
         sd_pipe.set_progress_bar_config(disable=A_       )

         _lowercase:			Optional[Any]       =     '''A painting of a squirrel eating a burger'''
         np.random.seed(0       )
         _lowercase:			Optional[int]       =     sd_pipe([prompt] ,				guidance_scale=6.0 ,				num_inference_steps=10 ,				output_type='''np'''       )
         _lowercase:			Tuple       =     output.images

         _lowercase:			Dict       =     image[0, -3:, -3:, -1]

         assert image.shape == (1, 512, 512, 3)
         _lowercase:			Union[str, Any]       =     np.array([0.04_52, 0.03_90, 0.00_87, 0.03_50, 0.06_17, 0.03_64, 0.05_44, 0.05_23, 0.07_20]       )

         assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-3
       def 							lowercase_							(     self       )     ->	Any:





         """simple docstring"""




         _lowercase:			Tuple       =     DDIMScheduler.from_pretrained(
             '''runwayml/stable-diffusion-v1-5''' ,				subfolder='''scheduler''' ,				revision='''onnx'''       )
         _lowercase:			int       =     OnnxStableDiffusionPipeline.from_pretrained(
             '''runwayml/stable-diffusion-v1-5''' ,				revision='''onnx''' ,				scheduler=A_ ,				safety_checker=A_ ,				feature_extractor=A_ ,				provider=self.gpu_provider ,				sess_options=self.gpu_options ,				)
         sd_pipe.set_progress_bar_config(disable=A_       )

         _lowercase:			int       =     '''open neural network exchange'''
         _lowercase:			List[Any]       =     np.random.RandomState(0       )
         _lowercase:			Union[str, Any]       =     sd_pipe([prompt] ,				guidance_scale=7.5 ,				num_inference_steps=10 ,				generator=A_ ,				output_type='''np'''       )
         _lowercase:			List[Any]       =     output.images
         _lowercase:			Tuple       =     image[0, -3:, -3:, -1]

         assert image.shape == (1, 512, 512, 3)
         _lowercase:			Any       =     np.array([0.28_67, 0.19_74, 0.14_81, 0.72_94, 0.72_51, 0.66_67, 0.41_94, 0.56_42, 0.64_86]       )

         assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-3
       def 							lowercase_							(     self       )     ->	int:





         """simple docstring"""




         _lowercase:			int       =     LMSDiscreteScheduler.from_pretrained(
             '''runwayml/stable-diffusion-v1-5''' ,				subfolder='''scheduler''' ,				revision='''onnx'''       )
         _lowercase:			str       =     OnnxStableDiffusionPipeline.from_pretrained(
             '''runwayml/stable-diffusion-v1-5''' ,				revision='''onnx''' ,				scheduler=A_ ,				safety_checker=A_ ,				feature_extractor=A_ ,				provider=self.gpu_provider ,				sess_options=self.gpu_options ,				)
         sd_pipe.set_progress_bar_config(disable=A_       )

         _lowercase:			Optional[Any]       =     '''open neural network exchange'''
         _lowercase:			Tuple       =     np.random.RandomState(0       )
         _lowercase:			Optional[int]       =     sd_pipe([prompt] ,				guidance_scale=7.5 ,				num_inference_steps=10 ,				generator=A_ ,				output_type='''np'''       )
         _lowercase:			Dict       =     output.images
         _lowercase:			Optional[Any]       =     image[0, -3:, -3:, -1]

         assert image.shape == (1, 512, 512, 3)
         _lowercase:			List[Any]       =     np.array([0.23_06, 0.19_59, 0.15_93, 0.65_49, 0.63_94, 0.54_08, 0.50_65, 0.60_10, 0.61_61]       )

         assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-3
       def 							lowercase_							(     self       )     ->	Tuple:





         """simple docstring"""




         _lowercase:			Any       =     0

         def test_callback_fn(A_ ,				A_ ,				A_       ) -> None:
           _lowercase:			Tuple       =     True
           nonlocal number_of_steps
           number_of_steps += 1
           if step == 0:
             assert latents.shape == (1, 4, 64, 64)
             _lowercase:			Any       =     latents[0, -3:, -3:, -1]
             _lowercase:			str       =     np.array(
                 [-0.67_72, -0.38_35, -1.24_56, 0.19_05, -1.09_74, 0.69_67, -1.93_53, 0.01_78, 1.01_67]       )

             assert np.abs(latents_slice.flatten() - expected_slice       ).max() < 1E-3
           elif step == 5:
             assert latents.shape == (1, 4, 64, 64)
             _lowercase:			str       =     latents[0, -3:, -3:, -1]
             _lowercase:			str       =     np.array(
                 [-0.33_51, 0.22_41, -0.18_37, -0.23_25, -0.65_77, 0.33_93, -0.02_41, 0.58_99, 1.38_75]       )

             assert np.abs(latents_slice.flatten() - expected_slice       ).max() < 1E-3

         _lowercase:			Tuple       =     False

         _lowercase:			Union[str, Any]       =     OnnxStableDiffusionPipeline.from_pretrained(
             '''runwayml/stable-diffusion-v1-5''' ,				revision='''onnx''' ,				safety_checker=A_ ,				feature_extractor=A_ ,				provider=self.gpu_provider ,				sess_options=self.gpu_options ,				)
         pipe.set_progress_bar_config(disable=A_       )

         _lowercase:			List[Any]       =     '''Andromeda galaxy in a bottle'''

         _lowercase:			Any       =     np.random.RandomState(0       )
         pipe(
             prompt=A_ ,				num_inference_steps=5 ,				guidance_scale=7.5 ,				generator=A_ ,				callback=A_ ,				callback_steps=1 ,				)
         assert test_callback_fn.has_been_called
         assert number_of_steps == 6



       def 							lowercase_							(     self       )     ->	int:





         """simple docstring"""




         _lowercase:			Optional[int]       =     OnnxStableDiffusionPipeline.from_pretrained(
             '''runwayml/stable-diffusion-v1-5''' ,				revision='''onnx''' ,				safety_checker=A_ ,				feature_extractor=A_ ,				provider=self.gpu_provider ,				sess_options=self.gpu_options ,				)
         assert isinstance(A_ ,				A_       )
         assert pipe.safety_checker is None

         _lowercase:			Optional[int]       =     pipe('''example prompt''' ,				num_inference_steps=2       ).images[0]
         assert image is not None

         # check that there's no error when saving a pipeline with one of the models being None
         with tempfile.TemporaryDirectory() as tmpdirname:
           pipe.save_pretrained(A_       )
           _lowercase:			str       =     OnnxStableDiffusionPipeline.from_pretrained(A_       )

         # sanity check that the pipeline still works
         assert pipe.safety_checker is None
         _lowercase:			Optional[Any]       =     pipe('''example prompt''' ,				num_inference_steps=2       ).images[0]
         assert image is not None


"""simple docstring"""




import unittest

from transformers import MobileBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import (
      MODEL_FOR_PRETRAINING_MAPPING,
      MobileBertForMaskedLM,
      MobileBertForMultipleChoice,
      MobileBertForNextSentencePrediction,
      MobileBertForPreTraining,
      MobileBertForQuestionAnswering,
      MobileBertForSequenceClassification,
      MobileBertForTokenClassification,
      MobileBertModel,
  )


class 			__magic_name__       :
       def __init__(     self ,				A_ ,				A_=13 ,				A_=7 ,				A_=True ,				A_=True ,				A_=True ,				A_=True ,				A_=99 ,				A_=64 ,				A_=32 ,				A_=5 ,				A_=4 ,				A_=37 ,				A_="gelu" ,				A_=0.1 ,				A_=0.1 ,				A_=512 ,				A_=16 ,				A_=2 ,				A_=0.02 ,				A_=3 ,				A_=4 ,				A_=None ,				)     ->	Dict:





         """simple docstring"""




         _lowercase:			List[Any]       =     parent
         _lowercase:			str       =     batch_size
         _lowercase:			List[str]       =     seq_length
         _lowercase:			Any       =     is_training
         _lowercase:			Union[str, Any]       =     use_input_mask
         _lowercase:			Any       =     use_token_type_ids
         _lowercase:			int       =     use_labels
         _lowercase:			Any       =     vocab_size
         _lowercase:			Optional[int]       =     hidden_size
         _lowercase:			Optional[int]       =     embedding_size
         _lowercase:			Union[str, Any]       =     num_hidden_layers
         _lowercase:			Optional[int]       =     num_attention_heads
         _lowercase:			Optional[int]       =     intermediate_size
         _lowercase:			List[str]       =     hidden_act
         _lowercase:			Dict       =     hidden_dropout_prob
         _lowercase:			Any       =     attention_probs_dropout_prob
         _lowercase:			Tuple       =     max_position_embeddings
         _lowercase:			List[str]       =     type_vocab_size
         _lowercase:			Any       =     type_sequence_label_size
         _lowercase:			int       =     initializer_range
         _lowercase:			Any       =     num_labels
         _lowercase:			int       =     num_choices
         _lowercase:			List[str]       =     scope
       def 							lowercase_							(     self       )     ->	Optional[Any]:





         """simple docstring"""




         _lowercase:			Tuple       =     ids_tensor([self.batch_size, self.seq_length] ,				self.vocab_size       )

         _lowercase:			Optional[Any]       =     None
         if self.use_input_mask:
           _lowercase:			Any       =     random_attention_mask([self.batch_size, self.seq_length]       )

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

         _lowercase:			Dict       =     None
         _lowercase:			Tuple       =     None
         _lowercase:			List[str]       =     None
         if self.use_labels:
           _lowercase:			str       =     ids_tensor([self.batch_size] ,				self.type_sequence_label_size       )
           _lowercase:			int       =     ids_tensor([self.batch_size, self.seq_length] ,				self.num_labels       )
           _lowercase:			Dict       =     ids_tensor([self.batch_size] ,				self.num_choices       )

         _lowercase:			str       =     self.get_config()

         return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
       def 							lowercase_							(     self       )     ->	List[Any]:





         """simple docstring"""




         return 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 ,				embedding_size=self.embedding_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 ,				is_decoder=A_ ,				initializer_range=self.initializer_range ,				)
       def 							lowercase_							(     self ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_       )     ->	Any:





         """simple docstring"""




         _lowercase:			Tuple       =     MobileBertModel(config=A_       )
         model.to(A_       )
         model.eval()
         _lowercase:			Tuple       =     model(A_ ,				attention_mask=A_ ,				token_type_ids=A_       )
         _lowercase:			Dict       =     model(A_ ,				token_type_ids=A_       )
         _lowercase:			Optional[int]       =     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 							lowercase_							(     self ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_       )     ->	int:





         """simple docstring"""




         _lowercase:			int       =     MobileBertForMaskedLM(config=A_       )
         model.to(A_       )
         model.eval()
         _lowercase:			Optional[int]       =     model(A_ ,				attention_mask=A_ ,				token_type_ids=A_ ,				labels=A_       )
         self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.seq_length, self.vocab_size)       )
       def 							lowercase_							(     self ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_       )     ->	int:





         """simple docstring"""




         _lowercase:			Any       =     MobileBertForNextSentencePrediction(config=A_       )
         model.to(A_       )
         model.eval()
         _lowercase:			int       =     model(
             A_ ,				attention_mask=A_ ,				token_type_ids=A_ ,				labels=A_ ,				)
         self.parent.assertEqual(result.logits.shape ,				(self.batch_size, 2)       )
       def 							lowercase_							(     self ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_       )     ->	Optional[int]:





         """simple docstring"""




         _lowercase:			str       =     MobileBertForPreTraining(config=A_       )
         model.to(A_       )
         model.eval()
         _lowercase:			List[str]       =     model(
             A_ ,				attention_mask=A_ ,				token_type_ids=A_ ,				labels=A_ ,				next_sentence_label=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 							lowercase_							(     self ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_       )     ->	List[str]:





         """simple docstring"""




         _lowercase:			Optional[Any]       =     MobileBertForQuestionAnswering(config=A_       )
         model.to(A_       )
         model.eval()
         _lowercase:			Dict       =     model(
             A_ ,				attention_mask=A_ ,				token_type_ids=A_ ,				start_positions=A_ ,				end_positions=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 							lowercase_							(     self ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_       )     ->	List[Any]:





         """simple docstring"""




         _lowercase:			Optional[Any]       =     self.num_labels
         _lowercase:			int       =     MobileBertForSequenceClassification(A_       )
         model.to(A_       )
         model.eval()
         _lowercase:			List[Any]       =     model(A_ ,				attention_mask=A_ ,				token_type_ids=A_ ,				labels=A_       )
         self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.num_labels)       )
       def 							lowercase_							(     self ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_       )     ->	Tuple:





         """simple docstring"""




         _lowercase:			Dict       =     self.num_labels
         _lowercase:			Optional[Any]       =     MobileBertForTokenClassification(config=A_       )
         model.to(A_       )
         model.eval()
         _lowercase:			Optional[int]       =     model(A_ ,				attention_mask=A_ ,				token_type_ids=A_ ,				labels=A_       )
         self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.seq_length, self.num_labels)       )
       def 							lowercase_							(     self ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_ ,				A_       )     ->	List[str]:





         """simple docstring"""




         _lowercase:			Union[str, Any]       =     self.num_choices
         _lowercase:			List[Any]       =     MobileBertForMultipleChoice(config=A_       )
         model.to(A_       )
         model.eval()
         _lowercase:			Any       =     input_ids.unsqueeze(1       ).expand(-1 ,				self.num_choices ,				-1       ).contiguous()
         _lowercase:			Tuple       =     token_type_ids.unsqueeze(1       ).expand(-1 ,				self.num_choices ,				-1       ).contiguous()
         _lowercase:			List[str]       =     input_mask.unsqueeze(1       ).expand(-1 ,				self.num_choices ,				-1       ).contiguous()
         _lowercase:			Any       =     model(
             A_ ,				attention_mask=A_ ,				token_type_ids=A_ ,				labels=A_ ,				)
         self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.num_choices)       )



       def 							lowercase_							(     self       )     ->	Optional[int]:





         """simple docstring"""




         _lowercase:			Any       =     self.prepare_config_and_inputs()
         (
             (
             _lowercase
         ) ,				(
             _lowercase
         ) ,				(
             _lowercase
         ) ,				(
             _lowercase
         ) ,				(
             _lowercase
         ) ,				(
             _lowercase
         ) ,				(
             _lowercase
         ) ,				
         ):			Dict       =     config_and_inputs
         _lowercase:			int       =     {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
         return config, inputs_dict


@require_torch
class 			__magic_name__       (       SCREAMING_SNAKE_CASE__   ,					SCREAMING_SNAKE_CASE__   ,					unittest.TestCase ):
       UpperCamelCase_   =		(
           (
               MobileBertModel,
               MobileBertForMaskedLM,
               MobileBertForMultipleChoice,
               MobileBertForNextSentencePrediction,
               MobileBertForPreTraining,
               MobileBertForQuestionAnswering,
               MobileBertForSequenceClassification,
               MobileBertForTokenClassification,
           )
           if is_torch_available()
           else ()
       )
       UpperCamelCase_   =		(
           {
               '''feature-extraction''': MobileBertModel,
               '''fill-mask''': MobileBertForMaskedLM,
               '''question-answering''': MobileBertForQuestionAnswering,
               '''text-classification''': MobileBertForSequenceClassification,
               '''token-classification''': MobileBertForTokenClassification,
               '''zero-shot''': MobileBertForSequenceClassification,
           }
           if is_torch_available()
           else {}
       )
       UpperCamelCase_   =		True
       def 							lowercase_							(     self ,				A_ ,				A_ ,				A_=False       )     ->	str:





         """simple docstring"""




         _lowercase:			Optional[Any]       =     super()._prepare_for_class(A_ ,				A_ ,				return_labels=A_       )

         if return_labels:
           if model_class in get_values(A_       ):
             _lowercase:			Tuple       =     torch.zeros(
                 (self.model_tester.batch_size, self.model_tester.seq_length) ,				dtype=torch.long ,				device=A_       )
             _lowercase:			Tuple       =     torch.zeros(
                 self.model_tester.batch_size ,				dtype=torch.long ,				device=A_       )
         return inputs_dict
       def 							lowercase_							(     self       )     ->	Tuple:





         """simple docstring"""




         _lowercase:			Optional[Any]       =     MobileBertModelTester(self       )
         _lowercase:			str       =     ConfigTester(self ,				config_class=A_ ,				hidden_size=37       )
       def 							lowercase_							(     self       )     ->	Any:





         """simple docstring"""




         self.config_tester.run_common_tests()
       def 							lowercase_							(     self       )     ->	Union[str, Any]:





         """simple docstring"""




         _lowercase:			str       =     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_model(*A_       )
       def 							lowercase_							(     self       )     ->	Optional[Any]:





         """simple docstring"""




         _lowercase:			Union[str, Any]       =     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_masked_lm(*A_       )
       def 							lowercase_							(     self       )     ->	Any:





         """simple docstring"""




         _lowercase:			List[str]       =     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_multiple_choice(*A_       )
       def 							lowercase_							(     self       )     ->	Union[str, Any]:





         """simple docstring"""




         _lowercase:			Any       =     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*A_       )
       def 							lowercase_							(     self       )     ->	Optional[Any]:





         """simple docstring"""




         _lowercase:			Optional[Any]       =     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_pretraining(*A_       )
       def 							lowercase_							(     self       )     ->	Tuple:





         """simple docstring"""




         _lowercase:			List[Any]       =     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_question_answering(*A_       )
       def 							lowercase_							(     self       )     ->	List[str]:





         """simple docstring"""




         _lowercase:			str       =     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_sequence_classification(*A_       )



       def 							lowercase_							(     self       )     ->	List[Any]:





         """simple docstring"""




         _lowercase:			Dict       =     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_token_classification(*A_       )




def  _lowerCAmelCase		(		_UpperCamelCase						):






  """simple docstring"""



  return torch.tensor(
      _UpperCamelCase ,				dtype=torch.long ,				device=_UpperCamelCase ,				)


A__   :					Optional[Any]							=   1e-3







@require_torch
@require_sentencepiece
@require_tokenizers
class 			__magic_name__       (       unittest.TestCase ):
       @slow
       def 							lowercase_							(     self       )     ->	int:





         """simple docstring"""




         _lowercase:			Optional[int]       =     MobileBertModel.from_pretrained('''google/mobilebert-uncased'''       ).to(A_       )
         _lowercase:			str       =     _long_tensor([[101, 7110, 1005, 1056, 2023, 1_1333, 1_7413, 1029, 102]]       )
         with torch.no_grad():
           _lowercase:			Dict       =     model(A_       )[0]
         _lowercase:			List[Any]       =     torch.Size((1, 9, 512)       )
         self.assertEqual(output.shape ,				A_       )
         _lowercase:			Tuple       =     torch.tensor(
             [
                 [
                     [-2.4_73_65_26E07, 8.2_69_16_56E04, 1.6_52_18_38E05],
                     [-5.7_54_17_04E-01, 3.9_05_60_22E00, 4.4_01_15_07E00],
                     [2.6_04_73_59E00, 1.5_67_76_52E00, -1.7_32_41_88E-01],
                 ]
             ] ,				device=A_ ,				)

         # MobileBERT results range from 10e0 to 10e8. Even a 0.0000001% difference with a value of 10e8 results in a
         # ~1 difference, it's therefore not a good idea to measure using addition.
         # Here, we instead divide the expected result with the result in order to obtain ~1. We then check that the
         # result is held between bounds: 1 - TOLERANCE < expected_result / result < 1 + TOLERANCE
         _lowercase:			Tuple       =     torch.all((expected_slice / output[..., :3, :3]) >= 1 - TOLERANCE       )
         _lowercase:			Dict       =     torch.all((expected_slice / output[..., :3, :3]) <= 1 + TOLERANCE       )

         self.assertTrue(lower_bound and upper_bound       )
