






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 					A__     :

     '''simple docstring'''






     def __init__( self      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    Dict=13   ,						_SCREAMING_SNAKE_CASE      :    int=7   ,						_SCREAMING_SNAKE_CASE      :    Optional[int]=True   ,						_SCREAMING_SNAKE_CASE      :    int=True   ,						_SCREAMING_SNAKE_CASE      :    Any=True   ,						_SCREAMING_SNAKE_CASE      :    str=True   ,						_SCREAMING_SNAKE_CASE      :    Tuple=99   ,						_SCREAMING_SNAKE_CASE      :    Optional[int]=32   ,						_SCREAMING_SNAKE_CASE      :    str=5   ,						_SCREAMING_SNAKE_CASE      :    Optional[int]=4   ,						_SCREAMING_SNAKE_CASE      :    List[Any]=37   ,						_SCREAMING_SNAKE_CASE      :    Any="gelu"   ,						_SCREAMING_SNAKE_CASE      :    Tuple=0.1   ,						_SCREAMING_SNAKE_CASE      :    Union[str, Any]=0.1   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]=128   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]=32   ,						_SCREAMING_SNAKE_CASE      :    int=16   ,						_SCREAMING_SNAKE_CASE      :    Dict=2   ,						_SCREAMING_SNAKE_CASE      :    Dict=0.0_2   ,						_SCREAMING_SNAKE_CASE      :    Optional[int]=3   ,						_SCREAMING_SNAKE_CASE      :    int=4   ,						_SCREAMING_SNAKE_CASE      :    Optional[int]=None   ,						):






           """simple docstring"""




           UpperCamelCase					       =  parent
           UpperCamelCase					       =  batch_size
           UpperCamelCase					       =  seq_length
           UpperCamelCase					       =  is_training
           UpperCamelCase					       =  use_input_mask
           UpperCamelCase					       =  use_token_type_ids
           UpperCamelCase					       =  use_labels
           UpperCamelCase					       =  vocab_size
           UpperCamelCase					       =  hidden_size
           UpperCamelCase					       =  num_hidden_layers
           UpperCamelCase					       =  num_attention_heads
           UpperCamelCase					       =  intermediate_size
           UpperCamelCase					       =  hidden_act
           UpperCamelCase					       =  hidden_dropout_prob
           UpperCamelCase					       =  attention_probs_dropout_prob
           UpperCamelCase					       =  max_position_embeddings
           UpperCamelCase					       =  type_vocab_size
           UpperCamelCase					       =  type_sequence_label_size
           UpperCamelCase					       =  initializer_range
           UpperCamelCase					       =  num_labels
           UpperCamelCase					       =  num_choices
           UpperCamelCase					       =  scope







     def   _SCREAMING_SNAKE_CASE     ( self      :    Any ):






           """simple docstring"""




           UpperCamelCase					       =  ids_tensor([self.batch_size, self.seq_length]   ,						self.vocab_size )

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

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

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

           UpperCamelCase					       =  self.get_config()

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







     def   _SCREAMING_SNAKE_CASE     ( self      :    Optional[int] ):






           """simple docstring"""




           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=_SCREAMING_SNAKE_CASE   ,						initializer_range=self.initializer_range   ,						)







     def   _SCREAMING_SNAKE_CASE     ( self      :    List[Any] ):






           """simple docstring"""




           (
               (
               UpperCamelCase
           )       ,  (
               UpperCamelCase
           )       ,  (
               UpperCamelCase
           )       ,  (
               UpperCamelCase
           )       ,  (
               UpperCamelCase
           )       ,  (
               UpperCamelCase
           )       ,  (
               UpperCamelCase
           )       ,  
           )					       =  self.prepare_config_and_inputs()

           UpperCamelCase					       =  True
           UpperCamelCase					       =  floats_tensor([self.batch_size, self.seq_length, self.hidden_size] )
           UpperCamelCase					       =  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   _SCREAMING_SNAKE_CASE     ( self      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    Any   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    int   ,						_SCREAMING_SNAKE_CASE      :    Any   ,						_SCREAMING_SNAKE_CASE      :    Dict   ,						_SCREAMING_SNAKE_CASE      :    str ):






           """simple docstring"""




           UpperCamelCase					       =  NezhaModel(config=_SCREAMING_SNAKE_CASE )
           model.to(_SCREAMING_SNAKE_CASE )
           model.eval()
           UpperCamelCase					       =  model(_SCREAMING_SNAKE_CASE   ,						attention_mask=_SCREAMING_SNAKE_CASE   ,						token_type_ids=_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  model(_SCREAMING_SNAKE_CASE   ,						token_type_ids=_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  model(_SCREAMING_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   _SCREAMING_SNAKE_CASE     ( self      :    Any   ,						_SCREAMING_SNAKE_CASE      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    Dict   ,						_SCREAMING_SNAKE_CASE      :    str   ,						_SCREAMING_SNAKE_CASE      :    Any   ,						_SCREAMING_SNAKE_CASE      :    Dict   ,						_SCREAMING_SNAKE_CASE      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    List[str]   ,						_SCREAMING_SNAKE_CASE      :    List[str]   ,						_SCREAMING_SNAKE_CASE      :    Tuple   ,						):






           """simple docstring"""




           UpperCamelCase					       =  True
           UpperCamelCase					       =  NezhaModel(_SCREAMING_SNAKE_CASE )
           model.to(_SCREAMING_SNAKE_CASE )
           model.eval()
           UpperCamelCase					       =  model(
               _SCREAMING_SNAKE_CASE   ,						attention_mask=_SCREAMING_SNAKE_CASE   ,						token_type_ids=_SCREAMING_SNAKE_CASE   ,						encoder_hidden_states=_SCREAMING_SNAKE_CASE   ,						encoder_attention_mask=_SCREAMING_SNAKE_CASE   ,						)
           UpperCamelCase					       =  model(
               _SCREAMING_SNAKE_CASE   ,						attention_mask=_SCREAMING_SNAKE_CASE   ,						token_type_ids=_SCREAMING_SNAKE_CASE   ,						encoder_hidden_states=_SCREAMING_SNAKE_CASE   ,						)
           UpperCamelCase					       =  model(_SCREAMING_SNAKE_CASE   ,						attention_mask=_SCREAMING_SNAKE_CASE   ,						token_type_ids=_SCREAMING_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   _SCREAMING_SNAKE_CASE     ( self      :    List[str]   ,						_SCREAMING_SNAKE_CASE      :    Optional[int]   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    int   ,						_SCREAMING_SNAKE_CASE      :    List[str]   ,						_SCREAMING_SNAKE_CASE      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    List[str]   ,						_SCREAMING_SNAKE_CASE      :    Any ):






           """simple docstring"""




           UpperCamelCase					       =  NezhaForMaskedLM(config=_SCREAMING_SNAKE_CASE )
           model.to(_SCREAMING_SNAKE_CASE )
           model.eval()
           UpperCamelCase					       =  model(_SCREAMING_SNAKE_CASE   ,						attention_mask=_SCREAMING_SNAKE_CASE   ,						token_type_ids=_SCREAMING_SNAKE_CASE   ,						labels=_SCREAMING_SNAKE_CASE )
           self.parent.assertEqual(result.logits.shape   ,						(self.batch_size, self.seq_length, self.vocab_size) )







     def   _SCREAMING_SNAKE_CASE     ( self      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    str   ,						_SCREAMING_SNAKE_CASE      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    List[str]   ,						_SCREAMING_SNAKE_CASE      :    Union[str, Any]   ,						_SCREAMING_SNAKE_CASE      :    str   ,						_SCREAMING_SNAKE_CASE      :    List[Any] ):






           """simple docstring"""




           UpperCamelCase					       =  NezhaForNextSentencePrediction(config=_SCREAMING_SNAKE_CASE )
           model.to(_SCREAMING_SNAKE_CASE )
           model.eval()
           UpperCamelCase					       =  model(
               _SCREAMING_SNAKE_CASE   ,						attention_mask=_SCREAMING_SNAKE_CASE   ,						token_type_ids=_SCREAMING_SNAKE_CASE   ,						labels=_SCREAMING_SNAKE_CASE   ,						)
           self.parent.assertEqual(result.logits.shape   ,						(self.batch_size, 2) )







     def   _SCREAMING_SNAKE_CASE     ( self      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    Any   ,						_SCREAMING_SNAKE_CASE      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    List[str]   ,						_SCREAMING_SNAKE_CASE      :    List[str]   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any] ):






           """simple docstring"""




           UpperCamelCase					       =  NezhaForPreTraining(config=_SCREAMING_SNAKE_CASE )
           model.to(_SCREAMING_SNAKE_CASE )
           model.eval()
           UpperCamelCase					       =  model(
               _SCREAMING_SNAKE_CASE   ,						attention_mask=_SCREAMING_SNAKE_CASE   ,						token_type_ids=_SCREAMING_SNAKE_CASE   ,						labels=_SCREAMING_SNAKE_CASE   ,						next_sentence_label=_SCREAMING_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   _SCREAMING_SNAKE_CASE     ( self      :    Optional[int]   ,						_SCREAMING_SNAKE_CASE      :    Any   ,						_SCREAMING_SNAKE_CASE      :    int   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    Dict   ,						_SCREAMING_SNAKE_CASE      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    int   ,						_SCREAMING_SNAKE_CASE      :    List[Any] ):






           """simple docstring"""




           UpperCamelCase					       =  NezhaForQuestionAnswering(config=_SCREAMING_SNAKE_CASE )
           model.to(_SCREAMING_SNAKE_CASE )
           model.eval()
           UpperCamelCase					       =  model(
               _SCREAMING_SNAKE_CASE   ,						attention_mask=_SCREAMING_SNAKE_CASE   ,						token_type_ids=_SCREAMING_SNAKE_CASE   ,						start_positions=_SCREAMING_SNAKE_CASE   ,						end_positions=_SCREAMING_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   _SCREAMING_SNAKE_CASE     ( self      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    str   ,						_SCREAMING_SNAKE_CASE      :    Any   ,						_SCREAMING_SNAKE_CASE      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    List[str]   ,						_SCREAMING_SNAKE_CASE      :    int ):






           """simple docstring"""




           UpperCamelCase					       =  self.num_labels
           UpperCamelCase					       =  NezhaForSequenceClassification(_SCREAMING_SNAKE_CASE )
           model.to(_SCREAMING_SNAKE_CASE )
           model.eval()
           UpperCamelCase					       =  model(_SCREAMING_SNAKE_CASE   ,						attention_mask=_SCREAMING_SNAKE_CASE   ,						token_type_ids=_SCREAMING_SNAKE_CASE   ,						labels=_SCREAMING_SNAKE_CASE )
           self.parent.assertEqual(result.logits.shape   ,						(self.batch_size, self.num_labels) )







     def   _SCREAMING_SNAKE_CASE     ( self      :    int   ,						_SCREAMING_SNAKE_CASE      :    Dict   ,						_SCREAMING_SNAKE_CASE      :    Optional[int]   ,						_SCREAMING_SNAKE_CASE      :    Dict   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    Optional[int]   ,						_SCREAMING_SNAKE_CASE      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    Dict ):






           """simple docstring"""




           UpperCamelCase					       =  self.num_labels
           UpperCamelCase					       =  NezhaForTokenClassification(config=_SCREAMING_SNAKE_CASE )
           model.to(_SCREAMING_SNAKE_CASE )
           model.eval()
           UpperCamelCase					       =  model(_SCREAMING_SNAKE_CASE   ,						attention_mask=_SCREAMING_SNAKE_CASE   ,						token_type_ids=_SCREAMING_SNAKE_CASE   ,						labels=_SCREAMING_SNAKE_CASE )
           self.parent.assertEqual(result.logits.shape   ,						(self.batch_size, self.seq_length, self.num_labels) )







     def   _SCREAMING_SNAKE_CASE     ( self      :    List[str]   ,						_SCREAMING_SNAKE_CASE      :    Dict   ,						_SCREAMING_SNAKE_CASE      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    int   ,						_SCREAMING_SNAKE_CASE      :    str   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    int ):






           """simple docstring"""




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







     def   _SCREAMING_SNAKE_CASE     ( self      :    int ):






           """simple docstring"""




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





@require_torch
class 					A__     (     __snake_case   ,    __snake_case   ,    __snake_case   ,    unittest.TestCase ):

     '''simple docstring'''






     snake_case__					=					(
         (
             NezhaModel,
             NezhaForMaskedLM,
             NezhaForMultipleChoice,
             NezhaForNextSentencePrediction,
             NezhaForPreTraining,
             NezhaForQuestionAnswering,
             NezhaForSequenceClassification,
             NezhaForTokenClassification,
         )
         if is_torch_available()
         else ()
     )
     snake_case__					=					(
         {
             """feature-extraction""": NezhaModel,
             """fill-mask""": NezhaForMaskedLM,
             """question-answering""": NezhaForQuestionAnswering,
             """text-classification""": NezhaForSequenceClassification,
             """token-classification""": NezhaForTokenClassification,
             """zero-shot""": NezhaForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     snake_case__					=					True







     def   _SCREAMING_SNAKE_CASE     ( self      :    Dict   ,						_SCREAMING_SNAKE_CASE      :    str   ,						_SCREAMING_SNAKE_CASE      :    Any   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]=False ):






           """simple docstring"""




           UpperCamelCase					       =  super()._prepare_for_class(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						return_labels=_SCREAMING_SNAKE_CASE )

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







     def   _SCREAMING_SNAKE_CASE     ( self      :    List[str] ):






           """simple docstring"""




           UpperCamelCase					       =  NezhaModelTester(self )
           UpperCamelCase					       =  ConfigTester(self   ,						config_class=_SCREAMING_SNAKE_CASE   ,						hidden_size=37 )







     def   _SCREAMING_SNAKE_CASE     ( self      :    List[Any] ):






           """simple docstring"""




           self.config_tester.run_common_tests()







     def   _SCREAMING_SNAKE_CASE     ( self      :    Union[str, Any] ):






           """simple docstring"""




           UpperCamelCase					       =  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    Tuple ):






           """simple docstring"""




           UpperCamelCase					       =  self.model_tester.prepare_config_and_inputs_for_decoder()
           self.model_tester.create_and_check_model_as_decoder(*_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    List[Any] ):






           """simple docstring"""




           (
               (
               UpperCamelCase
           )       ,  (
               UpperCamelCase
           )       ,  (
               UpperCamelCase
           )       ,  (
               UpperCamelCase
           )       ,  (
               UpperCamelCase
           )       ,  (
               UpperCamelCase
           )       ,  (
               UpperCamelCase
           )       ,  (
               UpperCamelCase
           )       ,  (
               UpperCamelCase
           )       ,  
           )					       =  self.model_tester.prepare_config_and_inputs_for_decoder()

           UpperCamelCase					       =  None

           self.model_tester.create_and_check_model_as_decoder(
               _SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						)







     def   _SCREAMING_SNAKE_CASE     ( self      :    Tuple ):






           """simple docstring"""




           UpperCamelCase					       =  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_lm(*_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    int ):






           """simple docstring"""




           UpperCamelCase					       =  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_multiple_choice(*_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    Any ):






           """simple docstring"""




           UpperCamelCase					       =  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_next_sequence_prediction(*_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    Union[str, Any] ):






           """simple docstring"""




           UpperCamelCase					       =  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_pretraining(*_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    Dict ):






           """simple docstring"""




           UpperCamelCase					       =  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_question_answering(*_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    Union[str, Any] ):






           """simple docstring"""




           UpperCamelCase					       =  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_sequence_classification(*_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    str ):






           """simple docstring"""




           UpperCamelCase					       =  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_token_classification(*_SCREAMING_SNAKE_CASE )







     @slow
     def   _SCREAMING_SNAKE_CASE     ( self      :    Optional[Any] ):






           """simple docstring"""




           for model_name in NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 UpperCamelCase					       =  NezhaModel.from_pretrained(_SCREAMING_SNAKE_CASE )
                 self.assertIsNotNone(_SCREAMING_SNAKE_CASE )







     @slow
     @require_torch_gpu
     def   _SCREAMING_SNAKE_CASE     ( self      :    Dict ):






           """simple docstring"""




           UpperCamelCase       ,  UpperCamelCase					       =  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

                 UpperCamelCase					       =  True
                 UpperCamelCase					       =  model_class(config=_SCREAMING_SNAKE_CASE )

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

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





@require_torch
class 					A__     (     unittest.TestCase ):

     '''simple docstring'''






     @slow
     def   _SCREAMING_SNAKE_CASE     ( self      :    Tuple ):






           """simple docstring"""




           UpperCamelCase					       =  NezhaModel.from_pretrained('sijunhe/nezha-cn-base' )
           UpperCamelCase					       =  torch.tensor([[0, 1, 2, 3, 4, 5]] )
           UpperCamelCase					       =  torch.tensor([[0, 1, 1, 1, 1, 1]] )
           with torch.no_grad():
                 UpperCamelCase					       =  model(_SCREAMING_SNAKE_CASE   ,						attention_mask=_SCREAMING_SNAKE_CASE )[0]
           UpperCamelCase					       =  torch.Size((1, 6, 768) )
           self.assertEqual(output.shape   ,						_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  torch.tensor([[[0.0_6_8_5, 0.2_4_4_1, 0.1_1_0_2], [0.0_6_0_0, 0.1_9_0_6, 0.1_3_4_9], [0.0_2_2_1, 0.0_8_1_9, 0.0_5_8_6]]] )

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







     @slow
     def   _SCREAMING_SNAKE_CASE     ( self      :    Any ):






           """simple docstring"""




           UpperCamelCase					       =  NezhaForMaskedLM.from_pretrained('sijunhe/nezha-cn-base' )
           UpperCamelCase					       =  torch.tensor([[0, 1, 2, 3, 4, 5]] )
           UpperCamelCase					       =  torch.tensor([[1, 1, 1, 1, 1, 1]] )
           with torch.no_grad():
                 UpperCamelCase					       =  model(_SCREAMING_SNAKE_CASE   ,						attention_mask=_SCREAMING_SNAKE_CASE )[0]
           UpperCamelCase					       =  torch.Size((1, 6, 2_1128) )
           self.assertEqual(output.shape   ,						_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  torch.tensor(
               [[-2.7_9_3_9, -1.7_9_0_2, -2.2_1_8_9], [-2.8_5_8_5, -1.8_9_0_8, -2.3_7_2_3], [-2.6_4_9_9, -1.7_7_5_0, -2.2_5_5_8]] )

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











import inspect
from typing import Callable, List, Optional, Union

import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextModel,
    CLIPTokenizer,
    WhisperForConditionalGeneration,
    WhisperProcessor,
)

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipelineOutput
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from diffusers.utils import logging


__magic_name__   :  Any        =							logging.get_logger(__name__)  # pylint: disable=invalid-name





class 					A__     (     __snake_case ):

     '''simple docstring'''






     def __init__( self      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    WhisperForConditionalGeneration   ,						_SCREAMING_SNAKE_CASE      :    WhisperProcessor   ,						_SCREAMING_SNAKE_CASE      :    AutoencoderKL   ,						_SCREAMING_SNAKE_CASE      :    CLIPTextModel   ,						_SCREAMING_SNAKE_CASE      :    CLIPTokenizer   ,						_SCREAMING_SNAKE_CASE      :    UNetaDConditionModel   ,						_SCREAMING_SNAKE_CASE      :    Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler]   ,						_SCREAMING_SNAKE_CASE      :    StableDiffusionSafetyChecker   ,						_SCREAMING_SNAKE_CASE      :    CLIPImageProcessor   ,						):






           """simple docstring"""




           super().__init__()

           if safety_checker is None:
                 logger.warning(
                     f'You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure'
                     ' that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered'
                     ' results in services or applications open to the public. Both the diffusers team and Hugging Face'
                     ' strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling'
                     ' it only for use-cases that involve analyzing network behavior or auditing its results. For more'
                     ' information, please have a look at https://github.com/huggingface/diffusers/pull/254 .' )

           self.register_modules(
               speech_model=_SCREAMING_SNAKE_CASE   ,						speech_processor=_SCREAMING_SNAKE_CASE   ,						vae=_SCREAMING_SNAKE_CASE   ,						text_encoder=_SCREAMING_SNAKE_CASE   ,						tokenizer=_SCREAMING_SNAKE_CASE   ,						unet=_SCREAMING_SNAKE_CASE   ,						scheduler=_SCREAMING_SNAKE_CASE   ,						feature_extractor=_SCREAMING_SNAKE_CASE   ,						)







     def   _SCREAMING_SNAKE_CASE     ( self      :    str   ,						_SCREAMING_SNAKE_CASE      :    Optional[Union[str, int]] = "auto" ):






           """simple docstring"""




           if slice_size == "auto":
                 UpperCamelCase					       =  self.unet.config.attention_head_dim // 2
           self.unet.set_attention_slice(_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    int ):






           """simple docstring"""




           self.enable_attention_slicing(_SCREAMING_SNAKE_CASE )







     @torch.no_grad()
     def __call__( self      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    Dict   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]=1_6000   ,						_SCREAMING_SNAKE_CASE      :    int = 512   ,						_SCREAMING_SNAKE_CASE      :    int = 512   ,						_SCREAMING_SNAKE_CASE      :    int = 50   ,						_SCREAMING_SNAKE_CASE      :    float = 7.5   ,						_SCREAMING_SNAKE_CASE      :    Optional[Union[str, List[str]]] = None   ,						_SCREAMING_SNAKE_CASE      :    Optional[int] = 1   ,						_SCREAMING_SNAKE_CASE      :    float = 0.0   ,						_SCREAMING_SNAKE_CASE      :    Optional[torch.Generator] = None   ,						_SCREAMING_SNAKE_CASE      :    Optional[torch.FloatTensor] = None   ,						_SCREAMING_SNAKE_CASE      :    Optional[str] = "pil"   ,						_SCREAMING_SNAKE_CASE      :    bool = True   ,						_SCREAMING_SNAKE_CASE      :    Optional[Callable[[int, int, torch.FloatTensor], None]] = None   ,						_SCREAMING_SNAKE_CASE      :    int = 1   ,						**_SCREAMING_SNAKE_CASE      :    str   ,						):






           """simple docstring"""




           UpperCamelCase					       =  self.speech_processor.feature_extractor(
               _SCREAMING_SNAKE_CASE   ,						return_tensors='pt'   ,						sampling_rate=_SCREAMING_SNAKE_CASE ).input_features.to(self.device )
           UpperCamelCase					       =  self.speech_model.generate(_SCREAMING_SNAKE_CASE   ,						max_length=48_0000 )

           UpperCamelCase					       =  self.speech_processor.tokenizer.batch_decode(_SCREAMING_SNAKE_CASE   ,						skip_special_tokens=_SCREAMING_SNAKE_CASE   ,						normalize=_SCREAMING_SNAKE_CASE )[
               0
           ]

           if isinstance(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE ):
                 UpperCamelCase					       =  1
           elif isinstance(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE ):
                 UpperCamelCase					       =  len(_SCREAMING_SNAKE_CASE )
           else:
                 raise ValueError(f'`prompt` has to be of type `str` or `list` but is {type(_SCREAMING_SNAKE_CASE )}' )

           if height % 8 != 0 or width % 8 != 0:
                 raise ValueError(f'`height` and `width` have to be divisible by 8 but are {height} and {width}.' )

           if (callback_steps is None) or (
               callback_steps is not None and (not isinstance(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE ) or callback_steps <= 0)
           ):
                 raise ValueError(
                     f'`callback_steps` has to be a positive integer but is {callback_steps} of type'
                     f' {type(_SCREAMING_SNAKE_CASE )}.' )

           # get prompt text embeddings
           UpperCamelCase					       =  self.tokenizer(
               _SCREAMING_SNAKE_CASE   ,						padding='max_length'   ,						max_length=self.tokenizer.model_max_length   ,						return_tensors='pt'   ,						)
           UpperCamelCase					       =  text_inputs.input_ids

           if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
                 UpperCamelCase					       =  self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :] )
                 logger.warning(
                     'The following part of your input was truncated because CLIP can only handle sequences up to'
                     f' {self.tokenizer.model_max_length} tokens: {removed_text}' )
                 UpperCamelCase					       =  text_input_ids[:, : self.tokenizer.model_max_length]
           UpperCamelCase					       =  self.text_encoder(text_input_ids.to(self.device ) )[0]

           # duplicate text embeddings for each generation per prompt, using mps friendly method
           UpperCamelCase       ,  UpperCamelCase       ,  UpperCamelCase					       =  text_embeddings.shape
           UpperCamelCase					       =  text_embeddings.repeat(1   ,						_SCREAMING_SNAKE_CASE   ,						1 )
           UpperCamelCase					       =  text_embeddings.view(bs_embed * num_images_per_prompt   ,						_SCREAMING_SNAKE_CASE   ,						-1 )

           # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
           # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
           # corresponds to doing no classifier free guidance.
           UpperCamelCase					       =  guidance_scale > 1.0
           # get unconditional embeddings for classifier free guidance
           if do_classifier_free_guidance:
                 UpperCamelCase					       =  42
                 if negative_prompt is None:
                       UpperCamelCase					       =  [''] * batch_size
                 elif type(_SCREAMING_SNAKE_CASE ) is not type(_SCREAMING_SNAKE_CASE ):
                       raise TypeError(
                           f'`negative_prompt` should be the same type to `prompt`, but got {type(_SCREAMING_SNAKE_CASE )} !='
                           f' {type(_SCREAMING_SNAKE_CASE )}.' )
                 elif isinstance(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE ):
                       UpperCamelCase					       =  [negative_prompt]
                 elif batch_size != len(_SCREAMING_SNAKE_CASE ):
                       raise ValueError(
                           f'`negative_prompt`: {negative_prompt} has batch size {len(_SCREAMING_SNAKE_CASE )}, but `prompt`:'
                           f' {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches'
                           ' the batch size of `prompt`.' )
                 else:
                       UpperCamelCase					       =  negative_prompt

                 UpperCamelCase					       =  text_input_ids.shape[-1]
                 UpperCamelCase					       =  self.tokenizer(
                     _SCREAMING_SNAKE_CASE   ,						padding='max_length'   ,						max_length=_SCREAMING_SNAKE_CASE   ,						truncation=_SCREAMING_SNAKE_CASE   ,						return_tensors='pt'   ,						)
                 UpperCamelCase					       =  self.text_encoder(uncond_input.input_ids.to(self.device ) )[0]

                 # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
                 UpperCamelCase					       =  uncond_embeddings.shape[1]
                 UpperCamelCase					       =  uncond_embeddings.repeat(1   ,						_SCREAMING_SNAKE_CASE   ,						1 )
                 UpperCamelCase					       =  uncond_embeddings.view(batch_size * num_images_per_prompt   ,						_SCREAMING_SNAKE_CASE   ,						-1 )

                 # For classifier free guidance, we need to do two forward passes.
                 # Here we concatenate the unconditional and text embeddings into a single batch
                 # to avoid doing two forward passes
                 UpperCamelCase					       =  torch.cat([uncond_embeddings, text_embeddings] )

           # get the initial random noise unless the user supplied it

           # Unlike in other pipelines, latents need to be generated in the target device
           # for 1-to-1 results reproducibility with the CompVis implementation.
           # However this currently doesn't work in `mps`.
           UpperCamelCase					       =  (batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
           UpperCamelCase					       =  text_embeddings.dtype
           if latents is None:
                 if self.device.type == "mps":
                       # randn does not exist on mps
                       UpperCamelCase					       =  torch.randn(_SCREAMING_SNAKE_CASE   ,						generator=_SCREAMING_SNAKE_CASE   ,						device='cpu'   ,						dtype=_SCREAMING_SNAKE_CASE ).to(
                           self.device )
                 else:
                       UpperCamelCase					       =  torch.randn(_SCREAMING_SNAKE_CASE   ,						generator=_SCREAMING_SNAKE_CASE   ,						device=self.device   ,						dtype=_SCREAMING_SNAKE_CASE )
           else:
                 if latents.shape != latents_shape:
                       raise ValueError(f'Unexpected latents shape, got {latents.shape}, expected {latents_shape}' )
                 UpperCamelCase					       =  latents.to(self.device )

           # set timesteps
           self.scheduler.set_timesteps(_SCREAMING_SNAKE_CASE )

           # Some schedulers like PNDM have timesteps as arrays
           # It's more optimized to move all timesteps to correct device beforehand
           UpperCamelCase					       =  self.scheduler.timesteps.to(self.device )

           # scale the initial noise by the standard deviation required by the scheduler
           UpperCamelCase					       =  latents * self.scheduler.init_noise_sigma

           # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
           # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
           # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
           # and should be between [0, 1]
           UpperCamelCase					       =  'eta' in set(inspect.signature(self.scheduler.step ).parameters.keys() )
           UpperCamelCase					       =  {}
           if accepts_eta:
                 UpperCamelCase					       =  eta

           for i, t in enumerate(self.progress_bar(_SCREAMING_SNAKE_CASE ) ):
                 # expand the latents if we are doing classifier free guidance
                 UpperCamelCase					       =  torch.cat([latents] * 2 ) if do_classifier_free_guidance else latents
                 UpperCamelCase					       =  self.scheduler.scale_model_input(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )

                 # predict the noise residual
                 UpperCamelCase					       =  self.unet(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						encoder_hidden_states=_SCREAMING_SNAKE_CASE ).sample

                 # perform guidance
                 if do_classifier_free_guidance:
                       UpperCamelCase       ,  UpperCamelCase					       =  noise_pred.chunk(2 )
                       UpperCamelCase					       =  noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

                 # compute the previous noisy sample x_t -> x_t-1
                 UpperCamelCase					       =  self.scheduler.step(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						**_SCREAMING_SNAKE_CASE ).prev_sample

                 # call the callback, if provided
                 if callback is not None and i % callback_steps == 0:
                       callback(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )

           UpperCamelCase					       =  1 / 0.1_8_2_1_5 * latents
           UpperCamelCase					       =  self.vae.decode(_SCREAMING_SNAKE_CASE ).sample

           UpperCamelCase					       =  (image / 2 + 0.5).clamp(0   ,						1 )

           # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
           UpperCamelCase					       =  image.cpu().permute(0   ,						2   ,						3   ,						1 ).float().numpy()

           if output_type == "pil":
                 UpperCamelCase					       =  self.numpy_to_pil(_SCREAMING_SNAKE_CASE )

           if not return_dict:
                 return image

           return StableDiffusionPipelineOutput(images=_SCREAMING_SNAKE_CASE   ,						nsfw_content_detected=_SCREAMING_SNAKE_CASE )



