

import os
import tempfile
import unittest

from transformers import NezhaConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import (
         MODEL_FOR_PRETRAINING_MAPPING,
         NezhaForMaskedLM,
         NezhaForMultipleChoice,
         NezhaForNextSentencePrediction,
         NezhaForPreTraining,
         NezhaForQuestionAnswering,
         NezhaForSequenceClassification,
         NezhaForTokenClassification,
         NezhaModel,
     )
     from transformers.models.nezha.modeling_nezha import NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST







class 			lowercase    :




    def __init__( self				,       snake_case				,       snake_case=13				,       snake_case=7				,       snake_case=True				,       snake_case=True				,       snake_case=True				,       snake_case=True				,       snake_case=99				,       snake_case=32				,       snake_case=5				,       snake_case=4				,       snake_case=37				,       snake_case="gelu"				,       snake_case=0.1				,       snake_case=0.1				,       snake_case=128				,       snake_case=32				,       snake_case=16				,       snake_case=2				,       snake_case=0.02				,       snake_case=3				,       snake_case=4				,       snake_case=None				,       ):
           snake_case_          =   parent
           snake_case_          =   batch_size
           snake_case_          =   seq_length
           snake_case_          =   is_training
           snake_case_          =   use_input_mask
           snake_case_          =   use_token_type_ids
           snake_case_          =   use_labels
           snake_case_          =   vocab_size
           snake_case_          =   hidden_size
           snake_case_          =   num_hidden_layers
           snake_case_          =   num_attention_heads
           snake_case_          =   intermediate_size
           snake_case_          =   hidden_act
           snake_case_          =   hidden_dropout_prob
           snake_case_          =   attention_probs_dropout_prob
           snake_case_          =   max_position_embeddings
           snake_case_          =   type_vocab_size
           snake_case_          =   type_sequence_label_size
           snake_case_          =   initializer_range
           snake_case_          =   num_labels
           snake_case_          =   num_choices
           snake_case_          =   scope




    def    a							( self	):
           snake_case_          =   ids_tensor([self.batch_size, self.seq_length]				,       self.vocab_size	)

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

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

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

           snake_case_          =   self.get_config()

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




    def    a							( self	):
           return NezhaConfig(
               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				,       is_decoder=snake_case				,       initializer_range=self.initializer_range				,       )




    def    a							( self	):
           (
               (
               snake_case_   
           )    ,   (
               snake_case_   
           )    ,   (
               snake_case_   
           )    ,   (
               snake_case_   
           )    ,   (
               snake_case_   
           )    ,   (
               snake_case_   
           )    ,   (
               snake_case_   
           )    ,      
           )          =   self.prepare_config_and_inputs()

           snake_case_          =   True
           snake_case_          =   floats_tensor([self.batch_size, self.seq_length, self.hidden_size]	)
           snake_case_          =   ids_tensor([self.batch_size, self.seq_length]				,       vocab_size=2	)

           return (
               config,
               input_ids,
               token_type_ids,
               input_mask,
               sequence_labels,
               token_labels,
               choice_labels,
               encoder_hidden_states,
               encoder_attention_mask,
           )




    def    a							( self				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case	):
           snake_case_          =   NezhaModel(config=snake_case	)
           model.to(snake_case	)
           model.eval()
           snake_case_          =   model(snake_case				,       attention_mask=snake_case				,       token_type_ids=snake_case	)
           snake_case_          =   model(snake_case				,       token_type_ids=snake_case	)
           snake_case_          =   model(snake_case	)
           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				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       ):
           snake_case_          =   True
           snake_case_          =   NezhaModel(snake_case	)
           model.to(snake_case	)
           model.eval()
           snake_case_          =   model(
               snake_case				,       attention_mask=snake_case				,       token_type_ids=snake_case				,       encoder_hidden_states=snake_case				,       encoder_attention_mask=snake_case				,       )
           snake_case_          =   model(
               snake_case				,       attention_mask=snake_case				,       token_type_ids=snake_case				,       encoder_hidden_states=snake_case				,       )
           snake_case_          =   model(snake_case				,       attention_mask=snake_case				,       token_type_ids=snake_case	)
           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				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case	):
           snake_case_          =   NezhaForMaskedLM(config=snake_case	)
           model.to(snake_case	)
           model.eval()
           snake_case_          =   model(snake_case				,       attention_mask=snake_case				,       token_type_ids=snake_case				,       labels=snake_case	)
           self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.seq_length, self.vocab_size)	)




    def    a							( self				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case	):
           snake_case_          =   NezhaForNextSentencePrediction(config=snake_case	)
           model.to(snake_case	)
           model.eval()
           snake_case_          =   model(
               snake_case				,       attention_mask=snake_case				,       token_type_ids=snake_case				,       labels=snake_case				,       )
           self.parent.assertEqual(result.logits.shape				,       (self.batch_size, 2)	)




    def    a							( self				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case	):
           snake_case_          =   NezhaForPreTraining(config=snake_case	)
           model.to(snake_case	)
           model.eval()
           snake_case_          =   model(
               snake_case				,       attention_mask=snake_case				,       token_type_ids=snake_case				,       labels=snake_case				,       next_sentence_label=snake_case				,       )
           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				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case	):
           snake_case_          =   NezhaForQuestionAnswering(config=snake_case	)
           model.to(snake_case	)
           model.eval()
           snake_case_          =   model(
               snake_case				,       attention_mask=snake_case				,       token_type_ids=snake_case				,       start_positions=snake_case				,       end_positions=snake_case				,       )
           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				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case	):
           snake_case_          =   self.num_labels
           snake_case_          =   NezhaForSequenceClassification(snake_case	)
           model.to(snake_case	)
           model.eval()
           snake_case_          =   model(snake_case				,       attention_mask=snake_case				,       token_type_ids=snake_case				,       labels=snake_case	)
           self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.num_labels)	)




    def    a							( self				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case	):
           snake_case_          =   self.num_labels
           snake_case_          =   NezhaForTokenClassification(config=snake_case	)
           model.to(snake_case	)
           model.eval()
           snake_case_          =   model(snake_case				,       attention_mask=snake_case				,       token_type_ids=snake_case				,       labels=snake_case	)
           self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.seq_length, self.num_labels)	)




    def    a							( self				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case	):
           snake_case_          =   self.num_choices
           snake_case_          =   NezhaForMultipleChoice(config=snake_case	)
           model.to(snake_case	)
           model.eval()
           snake_case_          =   input_ids.unsqueeze(1	).expand(-1				,       self.num_choices				,       -1	).contiguous()
           snake_case_          =   token_type_ids.unsqueeze(1	).expand(-1				,       self.num_choices				,       -1	).contiguous()
           snake_case_          =   input_mask.unsqueeze(1	).expand(-1				,       self.num_choices				,       -1	).contiguous()
           snake_case_          =   model(
               snake_case				,       attention_mask=snake_case				,       token_type_ids=snake_case				,       labels=snake_case				,       )
           self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.num_choices)	)




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







@require_torch
class 			lowercase    (	lowercase_     ,  lowercase_     ,  lowercase_     ,  unittest.TestCase      ):
    __SCREAMING_SNAKE_CASE :			Optional[int]	      =      (
        (
            NezhaModel,
            NezhaForMaskedLM,
            NezhaForMultipleChoice,
            NezhaForNextSentencePrediction,
            NezhaForPreTraining,
            NezhaForQuestionAnswering,
            NezhaForSequenceClassification,
            NezhaForTokenClassification,
        )
        if is_torch_available()
        else ()
    )
    __SCREAMING_SNAKE_CASE :			Union[str, Any]	      =      (
        {
            '''feature-extraction''': NezhaModel,
            '''fill-mask''': NezhaForMaskedLM,
            '''question-answering''': NezhaForQuestionAnswering,
            '''text-classification''': NezhaForSequenceClassification,
            '''token-classification''': NezhaForTokenClassification,
            '''zero-shot''': NezhaForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    __SCREAMING_SNAKE_CASE :			List[Any]	      =      True




    def    a							( self				,       snake_case				,       snake_case				,       snake_case=False	):
           snake_case_          =   super()._prepare_for_class(snake_case				,       snake_case				,       return_labels=snake_case	)

           if return_labels:
                  if model_class in get_values(snake_case	):
                         snake_case_          =   torch.zeros(
                             (self.model_tester.batch_size, self.model_tester.seq_length)				,       dtype=torch.long				,       device=snake_case	)
                         snake_case_          =   torch.zeros(
                             self.model_tester.batch_size				,       dtype=torch.long				,       device=snake_case	)
           return inputs_dict




    def    a							( self	):
           snake_case_          =   NezhaModelTester(self	)
           snake_case_          =   ConfigTester(self				,       config_class=snake_case				,       hidden_size=37	)




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




    def    a							( self	):
           snake_case_          =   self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*snake_case	)




    def    a							( self	):
           snake_case_          =   self.model_tester.prepare_config_and_inputs_for_decoder()
           self.model_tester.create_and_check_model_as_decoder(*snake_case	)




    def    a							( self	):
           # This regression test was failing with PyTorch < 1.3
           (
               (
               snake_case_   
           )    ,   (
               snake_case_   
           )    ,   (
               snake_case_   
           )    ,   (
               snake_case_   
           )    ,   (
               snake_case_   
           )    ,   (
               snake_case_   
           )    ,   (
               snake_case_   
           )    ,   (
               snake_case_   
           )    ,   (
               snake_case_   
           )    ,      
           )          =   self.model_tester.prepare_config_and_inputs_for_decoder()

           snake_case_          =   None

           self.model_tester.create_and_check_model_as_decoder(
               snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       )




    def    a							( self	):
           snake_case_          =   self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_lm(*snake_case	)




    def    a							( self	):
           snake_case_          =   self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_multiple_choice(*snake_case	)




    def    a							( self	):
           snake_case_          =   self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_next_sequence_prediction(*snake_case	)




    def    a							( self	):
           snake_case_          =   self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_pretraining(*snake_case	)




    def    a							( self	):
           snake_case_          =   self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_question_answering(*snake_case	)




    def    a							( self	):
           snake_case_          =   self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_sequence_classification(*snake_case	)




    def    a							( self	):
           snake_case_          =   self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_token_classification(*snake_case	)




    @slow
    def    a							( self	):
           for model_name in NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  snake_case_          =   NezhaModel.from_pretrained(snake_case	)
                  self.assertIsNotNone(snake_case	)




    @slow
    @require_torch_gpu
    def    a							( self	):
           snake_case_    ,   snake_case_          =   self.model_tester.prepare_config_and_inputs_for_common()
           for model_class in self.all_model_classes:
                  # NezhaForMultipleChoice behaves incorrectly in JIT environments.
                  if model_class == NezhaForMultipleChoice:
                         return

                  snake_case_          =   True
                  snake_case_          =   model_class(config=snake_case	)

                  snake_case_          =   self._prepare_for_class(snake_case				,       snake_case	)
                  snake_case_          =   torch.jit.trace(
                      snake_case				,       (inputs_dict['input_ids'].to('cpu'	), inputs_dict['attention_mask'].to('cpu'	))	)

                  with tempfile.TemporaryDirectory() as tmp:
                         torch.jit.save(snake_case				,       os.path.join(snake_case				,       'bert.pt'	)	)
                         snake_case_          =   torch.jit.load(os.path.join(snake_case				,       'bert.pt'	)				,       map_location=snake_case	)
                         loaded(inputs_dict['input_ids'].to(snake_case	)				,       inputs_dict['attention_mask'].to(snake_case	)	)







@require_torch
class 			lowercase    (	unittest.TestCase      ):




    @slow
    def    a							( self	):
           snake_case_          =   NezhaModel.from_pretrained('sijunhe/nezha-cn-base'	)
           snake_case_          =   torch.tensor([[0, 1, 2, 3, 4, 5]]	)
           snake_case_          =   torch.tensor([[0, 1, 1, 1, 1, 1]]	)
           with torch.no_grad():
                  snake_case_          =   model(snake_case				,       attention_mask=snake_case	)[0]
           snake_case_          =   torch.Size((1, 6, 768)	)
           self.assertEqual(output.shape				,       snake_case	)
           snake_case_          =   torch.tensor([[[0.06_85, 0.24_41, 0.11_02], [0.06_00, 0.19_06, 0.13_49], [0.02_21, 0.08_19, 0.05_86]]]	)

           self.assertTrue(torch.allclose(output[:, 1:4, 1:4]				,       snake_case				,       atol=1e-4	)	)




    @slow
    def    a							( self	):
           snake_case_          =   NezhaForMaskedLM.from_pretrained('sijunhe/nezha-cn-base'	)
           snake_case_          =   torch.tensor([[0, 1, 2, 3, 4, 5]]	)
           snake_case_          =   torch.tensor([[1, 1, 1, 1, 1, 1]]	)
           with torch.no_grad():
                  snake_case_          =   model(snake_case				,       attention_mask=snake_case	)[0]
           snake_case_          =   torch.Size((1, 6, 2_1128)	)
           self.assertEqual(output.shape				,       snake_case	)
           snake_case_          =   torch.tensor(
               [[-2.79_39, -1.79_02, -2.21_89], [-2.85_85, -1.89_08, -2.37_23], [-2.64_99, -1.77_50, -2.25_58]]	)

           self.assertTrue(torch.allclose(output[:, 1:4, 1:4]				,       snake_case				,       atol=1e-4	)	)







import gc
import unittest

from diffusers import FlaxControlNetModel, FlaxStableDiffusionControlNetPipeline
from diffusers.utils import is_flax_available, load_image, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
					import jax
					import jax.numpy as jnp
					from flax.jax_utils import replicate
					from flax.training.common_utils import shard







@slow
@require_flax
class 			lowercase    (	unittest.TestCase      ):




				def    a							( self	):
											# clean up the VRAM after each test
											super().tearDown()
											gc.collect()




				def    a							( self	):
											snake_case_    ,   snake_case_          =   FlaxControlNetModel.from_pretrained(
											    'lllyasviel/sd-controlnet-canny'				,       from_pt=snake_case				,       dtype=jnp.bfloataa	)
											snake_case_    ,   snake_case_          =   FlaxStableDiffusionControlNetPipeline.from_pretrained(
											    'runwayml/stable-diffusion-v1-5'				,       controlnet=snake_case				,       from_pt=snake_case				,       dtype=jnp.bfloataa	)
											snake_case_          =   controlnet_params

											snake_case_          =   'bird'
											snake_case_          =   jax.device_count()
											snake_case_          =   pipe.prepare_text_inputs([prompts] * num_samples	)

											snake_case_          =   load_image(
											    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png'	)
											snake_case_          =   pipe.prepare_image_inputs([canny_image] * num_samples	)

											snake_case_          =   jax.random.PRNGKey(0	)
											snake_case_          =   jax.random.split(snake_case				,       jax.device_count()	)

											snake_case_          =   replicate(snake_case	)
											snake_case_          =   shard(snake_case	)
											snake_case_          =   shard(snake_case	)

											snake_case_          =   pipe(
											    prompt_ids=snake_case				,       image=snake_case				,       params=snake_case				,       prng_seed=snake_case				,       num_inference_steps=50				,       jit=snake_case				,       ).images
											assert images.shape == (jax.device_count(), 1, 768, 512, 3)

											snake_case_          =   images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]	)
											snake_case_          =   images[0, 253:256, 253:256, -1]

											snake_case_          =   jnp.asarray(jax.device_get(image_slice.flatten()	)	)
											snake_case_          =   jnp.array(
											    [0.16_79_69, 0.11_66_99, 0.08_15_43, 0.15_42_97, 0.13_28_12, 0.10_88_87, 0.16_99_22, 0.16_99_22, 0.20_50_78]	)
											print(F'''output_slice: {output_slice}'''	)
											assert jnp.abs(output_slice - expected_slice	).max() < 1e-2




				def    a							( self	):
											snake_case_    ,   snake_case_          =   FlaxControlNetModel.from_pretrained(
											    'lllyasviel/sd-controlnet-openpose'				,       from_pt=snake_case				,       dtype=jnp.bfloataa	)
											snake_case_    ,   snake_case_          =   FlaxStableDiffusionControlNetPipeline.from_pretrained(
											    'runwayml/stable-diffusion-v1-5'				,       controlnet=snake_case				,       from_pt=snake_case				,       dtype=jnp.bfloataa	)
											snake_case_          =   controlnet_params

											snake_case_          =   'Chef in the kitchen'
											snake_case_          =   jax.device_count()
											snake_case_          =   pipe.prepare_text_inputs([prompts] * num_samples	)

											snake_case_          =   load_image(
											    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/pose.png'	)
											snake_case_          =   pipe.prepare_image_inputs([pose_image] * num_samples	)

											snake_case_          =   jax.random.PRNGKey(0	)
											snake_case_          =   jax.random.split(snake_case				,       jax.device_count()	)

											snake_case_          =   replicate(snake_case	)
											snake_case_          =   shard(snake_case	)
											snake_case_          =   shard(snake_case	)

											snake_case_          =   pipe(
											    prompt_ids=snake_case				,       image=snake_case				,       params=snake_case				,       prng_seed=snake_case				,       num_inference_steps=50				,       jit=snake_case				,       ).images
											assert images.shape == (jax.device_count(), 1, 768, 512, 3)

											snake_case_          =   images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]	)
											snake_case_          =   images[0, 253:256, 253:256, -1]

											snake_case_          =   jnp.asarray(jax.device_get(image_slice.flatten()	)	)
											snake_case_          =   jnp.array(
											    [[0.27_14_84, 0.26_17_19, 0.27_53_91, 0.27_73_44, 0.27_92_97, 0.29_10_16, 0.29_49_22, 0.30_27_34, 0.30_27_34]]	)
											print(F'''output_slice: {output_slice}'''	)
											assert jnp.abs(output_slice - expected_slice	).max() < 1e-2




