
from __future__ import annotations

import os
import tempfile
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import is_tensorflow_text_available, is_tf_available
from transformers.testing_utils import require_tensorflow_text, require_tf, slow

from ..test_modeling_tf_common import floats_tensor
from .test_framework_agnostic import GenerationIntegrationTestsMixin


if is_tf_available():
  import tensorflow as tf

  from transformers import (
      AutoTokenizer,
      TFAutoModelForCausalLM,
      TFAutoModelForSeqaSeqLM,
      TFAutoModelForSpeechSeqaSeq,
      TFAutoModelForVisionaSeq,
      TFBartForConditionalGeneration,
      TFLogitsProcessorList,
      TFMinLengthLogitsProcessor,
      tf_top_k_top_p_filtering,
  )

if is_tensorflow_text_available():
  import tensorflow_text as text







@require_tf
class        lowerCAmelCase__      (   unittest.TestCase	):

      def  A_				(			self				)					->      int:


        '''simple docstring'''

        _UpperCamelCase						    =	tf.convert_to_tensor(
            [
                [
                    8.222_0991,  # 3rd highest value; idx. 0
                    -0.562_0044,
                    5.2322_9752,
                    4.038_6393,
                    -6.879_8378,
                    -0.5478_5802,
                    -3.201_2153,
                    2.9277_7176,
                    1.8817_1953,
                    7.3534_1276,  # 5th highest value; idx. 9
                    8.4320_7833,  # 2nd highest value; idx. 10
                    -9.8571_1836,
                    -5.9620_9236,
                    -1.1303_9161,
                    -7.111_5294,
                    -0.836_9633,
                    -5.318_6408,
                    7.0642_7407,
                    0.8136_9344,
                    -0.8202_3817,
                    -5.917_9796,
                    0.5881_3443,
                    -6.9977_8438,
                    4.7155_1189,
                    -0.1877_1637,
                    7.4402_0759,  # 4th highest value; idx. 25
                    9.3845_0987,  # 1st highest value; idx. 26
                    2.1266_2941,
                    -9.3256_2038,
                    2.3565_2522,
                ],  # cummulative prob of 5 highest values <= 0.6
                [
                    0.5842_5518,
                    4.5313_9238,
                    -5.5751_0464,
                    -6.2803_0699,
                    -7.1952_9503,
                    -4.0212_2551,
                    1.3933_7037,
                    -6.0670_7057,
                    1.5948_0517,
                    -9.64_3119,
                    0.0390_7799,
                    0.6723_1762,
                    -8.8820_6726,
                    6.2711_5922,  # 4th highest value; idx. 13
                    2.2852_0723,
                    4.8276_7506,
                    4.3042_1368,
                    8.827_5313,  # 2nd highest value; idx. 17
                    5.4402_9958,  # 5th highest value; idx. 18
                    -4.473_5794,
                    7.3857_9536,  # 3rd highest value; idx. 20
                    -2.9105_1663,
                    2.6194_6077,
                    -2.567_4762,
                    -9.4895_9302,
                    -4.0292_2645,
                    -1.3541_6918,
                    9.6770_2323,  # 1st highest value; idx. 27
                    -5.8947_8553,
                    1.8537_0467,
                ],  # cummulative prob of 5 highest values <= 0.6
            ]       ,							dtype=tf.floataa       ,							)

        _UpperCamelCase						    =	tf.convert_to_tensor(
            [[0, 0], [0, 9], [0, 10], [0, 25], [0, 26], [1, 13], [1, 17], [1, 18], [1, 20], [1, 27]]       ,							dtype=tf.intaa       ,							)  # expected non filtered idx as noted above

        _UpperCamelCase						    =	tf.convert_to_tensor(
            [8.22_2099, 7.353_4126, 8.43_2078, 7.440_2075, 9.3_8451, 6.27_1159, 8.82_7531, 5.440_2995, 7.385_7956, 9.67_7023]       ,							dtype=tf.floataa       ,							)  # expected non filtered values as noted above

        _UpperCamelCase						    =	tf_top_k_top_p_filtering(_a       ,							top_k=10       ,							top_p=0.6       ,							min_tokens_to_keep=4				)

        _UpperCamelCase						    =	output[output != -float("""inf"""				)]
        _UpperCamelCase						    =	tf.cast(
            tf.where(tf.not_equal(_a       ,							tf.constant(-float("""inf"""				)       ,							dtype=tf.floataa				)				)				)       ,							dtype=tf.intaa       ,							)

        tf.debugging.assert_near(_a       ,							_a       ,							rtol=1e-12				)
        tf.debugging.assert_equal(_a       ,							_a				)







@require_tf
class        lowerCAmelCase__      (   unittest.TestCase		,	UpperCamelCase__	):
      # setting framework_dependent_parameters needs to be gated, just like its contents' imports
      if is_tf_available():
            UpperCamelCase_							: str	=		{
                "AutoModelForCausalLM": TFAutoModelForCausalLM,
                "AutoModelForSpeechSeq2Seq": TFAutoModelForSpeechSeqaSeq,
                "AutoModelForSeq2SeqLM": TFAutoModelForSeqaSeqLM,
                "AutoModelForVision2Seq": TFAutoModelForVisionaSeq,
                "LogitsProcessorList": TFLogitsProcessorList,
                "MinLengthLogitsProcessor": TFMinLengthLogitsProcessor,
                "create_tensor_fn": tf.convert_to_tensor,
                "floats_tensor": floats_tensor,
                "return_tensors": "tf",
            }

      @slow
      def  A_				(			self				)					->      List[Any]:


        '''simple docstring'''

        _UpperCamelCase						    =	TFAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""				)
        _UpperCamelCase						    =	2
        _UpperCamelCase						    =	2







        class        lowerCAmelCase__      (   tf.Module	):

              def __init__(			self       ,							a				)					->      List[str]:


                '''simple docstring'''

                super(_a       ,							self				).__init__()
                _UpperCamelCase						    =	model







              @tf.function(
                  input_signature=(
                      tf.TensorSpec((None, input_length)       ,							tf.intaa       ,							name="""input_ids"""				),
                      tf.TensorSpec((None, input_length)       ,							tf.intaa       ,							name="""attention_mask"""				),
                  )       ,							jit_compile=_a       ,							)
              def  A_				(			self       ,							a       ,							a				)					->      Any:


                '''simple docstring'''

                _UpperCamelCase						    =	self.model.generate(
                    input_ids=_a       ,							attention_mask=_a       ,							max_new_tokens=_a       ,							return_dict_in_generate=_a       ,							)
                return {"sequences": outputs["sequences"]}




        _UpperCamelCase						    =	[[2, 0], [1_02, 1_03]]
        _UpperCamelCase						    =	[[1, 0], [1, 1]]
        _UpperCamelCase						    =	DummyModel(model=_a				)
        with tempfile.TemporaryDirectory() as tmp_dir:
          tf.saved_model.save(_a       ,							_a       ,							signatures={"""serving_default""": dummy_model.serving}				)
          _UpperCamelCase						    =	tf.saved_model.load(_a				).signatures["""serving_default"""]
          for batch_size in range(1       ,							len(_a				) + 1				):
            _UpperCamelCase						    =	{
                """input_ids""": tf.constant(dummy_input_ids[:batch_size]				),
                """attention_mask""": tf.constant(dummy_attention_masks[:batch_size]				),
            }
            _UpperCamelCase						    =	serving_func(**_a				)["""sequences"""]
            _UpperCamelCase						    =	test_model.generate(**_a       ,							max_new_tokens=_a				)
            tf.debugging.assert_equal(_a       ,							_a				)

      @slow
      def  A_				(			self				)					->      str:


        '''simple docstring'''

        _UpperCamelCase						    =	TFAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""				)
        _UpperCamelCase						    =	1
        _UpperCamelCase						    =	2







        class        lowerCAmelCase__      (   tf.Module	):

              def __init__(			self       ,							a				)					->      str:


                '''simple docstring'''

                super(_a       ,							self				).__init__()
                _UpperCamelCase						    =	model







              @tf.function(
                  input_signature=(
                      tf.TensorSpec((batch_size, None)       ,							tf.intaa       ,							name="""input_ids"""				),
                      tf.TensorSpec((batch_size, None)       ,							tf.intaa       ,							name="""attention_mask"""				),
                  )       ,							jit_compile=_a       ,							)
              def  A_				(			self       ,							a       ,							a				)					->      Any:


                '''simple docstring'''

                _UpperCamelCase						    =	self.model.generate(
                    input_ids=_a       ,							attention_mask=_a       ,							max_new_tokens=_a       ,							return_dict_in_generate=_a       ,							)
                return {"sequences": outputs["sequences"]}




        _UpperCamelCase						    =	[[2], [1_02, 1_03]]
        _UpperCamelCase						    =	[[1], [1, 1]]
        _UpperCamelCase						    =	DummyModel(model=_a				)
        with tempfile.TemporaryDirectory() as tmp_dir:
          tf.saved_model.save(_a       ,							_a       ,							signatures={"""serving_default""": dummy_model.serving}				)
          _UpperCamelCase						    =	tf.saved_model.load(_a				).signatures["""serving_default"""]
          for input_row in range(len(_a				)				):
            _UpperCamelCase						    =	{
                """input_ids""": tf.constant([dummy_input_ids[input_row]]				),
                """attention_mask""": tf.constant([dummy_attention_masks[input_row]]				),
            }
            _UpperCamelCase						    =	serving_func(**_a				)["""sequences"""]
            _UpperCamelCase						    =	test_model.generate(**_a       ,							max_new_tokens=_a				)
            tf.debugging.assert_equal(_a       ,							_a				)

      @slow
      @require_tensorflow_text
      def  A_				(			self				)					->      Tuple:


        '''simple docstring'''

        with tempfile.TemporaryDirectory() as tmp_dir:
          # file needed to load the TF tokenizer
          hf_hub_download(repo_id="""google/flan-t5-small"""       ,							filename="""spiece.model"""       ,							local_dir=_a				)







          class        lowerCAmelCase__      (   tf.keras.layers.Layer	):

                def __init__(			self				)					->      Tuple:


                  '''simple docstring'''

                  super().__init__()
                  _UpperCamelCase						    =	text.SentencepieceTokenizer(
                      model=tf.io.gfile.GFile(os.path.join(_a       ,							"""spiece.model"""				)       ,							"""rb"""				).read()				)
                  _UpperCamelCase						    =	TFAutoModelForSeqaSeqLM.from_pretrained("""hf-internal-testing/tiny-random-t5"""				)







                def  A_				(			self       ,							a       ,							*a       ,							**a				)					->      str:


                  '''simple docstring'''

                  _UpperCamelCase						    =	self.tokenizer.tokenize(_a				)
                  _UpperCamelCase						    =	text.pad_model_inputs(
                      _a       ,							max_seq_length=64       ,							pad_value=self.model.config.pad_token_id				)
                  _UpperCamelCase						    =	self.model.generate(input_ids=_a       ,							attention_mask=_a				)
                  return self.tokenizer.detokenize(_a				)




          _UpperCamelCase						    =	CompleteSentenceTransformer()
          _UpperCamelCase						    =	tf.keras.layers.Input(shape=(1,)       ,							dtype=tf.string       ,							name="""inputs"""				)
          _UpperCamelCase						    =	complete_model(_a				)
          _UpperCamelCase						    =	tf.keras.Model(_a       ,							_a				)
          keras_model.save(_a				)

      def  A_				(			self				)					->      Optional[Any]:


        '''simple docstring'''

        _UpperCamelCase						    =	{
            """do_sample""": True,
            """num_beams""": 1,
            """top_p""": 0.7,
            """top_k""": 10,
            """temperature""": 0.7,
        }
        _UpperCamelCase						    =	14

        _UpperCamelCase						    =	AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""				)
        _UpperCamelCase						    =	"""Hello, my dog is cute and"""
        _UpperCamelCase						    =	tokenizer(_a       ,							return_tensors="""tf"""				)
        _UpperCamelCase						    =	TFAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""				)

        _UpperCamelCase						    =	6_38
        # forces the generation to happen on CPU, to avoid GPU-related quirks
        with tf.device(""":/CPU:0"""				):
          tf.random.set_seed(0				)
          _UpperCamelCase						    =	model.generate(**_a       ,							eos_token_id=_a       ,							**_a				)
        self.assertTrue(expectation == len(generated_tokens[0]				)				)

        _UpperCamelCase						    =	[6_38, 1_98]
        with tf.device(""":/CPU:0"""				):
          tf.random.set_seed(0				)
          _UpperCamelCase						    =	model.generate(**_a       ,							eos_token_id=_a       ,							**_a				)
        self.assertTrue(expectation == len(generated_tokens[0]				)				)







      def  A_				(			self				)					->      Optional[Any]:


        '''simple docstring'''

        _UpperCamelCase						    =	AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-bart"""				)
        _UpperCamelCase						    =	"""Hugging Face is a technology company based in New York and Paris."""
        _UpperCamelCase						    =	bart_tokenizer(_a       ,							return_tensors="""tf"""				).input_ids
        _UpperCamelCase						    =	TFBartForConditionalGeneration.from_pretrained("""hf-internal-testing/tiny-random-bart"""				)
        _UpperCamelCase						    =	bart_model.generate(_a				).numpy()







        class        lowerCAmelCase__      (   UpperCamelCase__	):

              def  A_				(			self       ,							a       ,							a=None       ,							**a				)					->      Dict:


                '''simple docstring'''

                return super().call(_a       ,							**_a				)




        _UpperCamelCase						    =	FakeBart.from_pretrained("""hf-internal-testing/tiny-random-bart"""				)
        _UpperCamelCase						    =	bart_model.generate(_a       ,							foo="""bar"""				).numpy()
        self.assertTrue(np.array_equal(_a       ,							_a				)				)







        class        lowerCAmelCase__      (   bart_model.model.encoder.__class__	):

              def  A_				(			self       ,							a       ,							**a				)					->      Union[str, Any]:


                '''simple docstring'''

                return super().call(_a       ,							**_a				)




        _UpperCamelCase						    =	FakeEncoder(bart_model.config       ,							bart_model.model.shared				)
        _UpperCamelCase						    =	fake_encoder

        # Normal generation still works (the output will be different because the encoder weights are different)
        _UpperCamelCase						    =	bart_model.generate(_a				).numpy()
        with self.assertRaises(_a				):
          # FakeEncoder.call() accepts **kwargs -> no filtering -> value error due to unexpected input "foo"
          bart_model.generate(_a       ,							foo="""bar"""				)





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

						def __init__(			self       ,							a       ,							a=13       ,							a=7       ,							a=True       ,							a=True       ,							a=True       ,							a=True       ,							a=99       ,							a=32       ,							a=5       ,							a=4       ,							a=37       ,							a="gelu"       ,							a=0.1       ,							a=0.1       ,							a=1_28       ,							a=32       ,							a=16       ,							a=2       ,							a=0.02       ,							a=3       ,							a=4       ,							a=None       ,							)					->      Optional[int]:


								'''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  A_				(			self				)					->      List[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  A_				(			self				)					->      Optional[Any]:


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

						def  A_				(			self				)					->      int:


								'''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  A_				(			self       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a				)					->      Union[str, Any]:


								'''simple docstring'''

								_UpperCamelCase						    =	NezhaModel(config=a				)
								model.to(a				)
								model.eval()
								_UpperCamelCase						    =	model(a       ,							attention_mask=a       ,							token_type_ids=a				)
								_UpperCamelCase						    =	model(a       ,							token_type_ids=a				)
								_UpperCamelCase						    =	model(a				)
								self.parent.assertEqual(result.last_hidden_state.shape       ,							(self.batch_size, self.seq_length, self.hidden_size)				)
								self.parent.assertEqual(result.pooler_output.shape       ,							(self.batch_size, self.hidden_size)				)

						def  A_				(			self       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							)					->      List[Any]:


								'''simple docstring'''

								_UpperCamelCase						    =	True
								_UpperCamelCase						    =	NezhaModel(a				)
								model.to(a				)
								model.eval()
								_UpperCamelCase						    =	model(
								    a       ,							attention_mask=a       ,							token_type_ids=a       ,							encoder_hidden_states=a       ,							encoder_attention_mask=a       ,							)
								_UpperCamelCase						    =	model(
								    a       ,							attention_mask=a       ,							token_type_ids=a       ,							encoder_hidden_states=a       ,							)
								_UpperCamelCase						    =	model(a       ,							attention_mask=a       ,							token_type_ids=a				)
								self.parent.assertEqual(result.last_hidden_state.shape       ,							(self.batch_size, self.seq_length, self.hidden_size)				)
								self.parent.assertEqual(result.pooler_output.shape       ,							(self.batch_size, self.hidden_size)				)

						def  A_				(			self       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a				)					->      Dict:


								'''simple docstring'''

								_UpperCamelCase						    =	NezhaForMaskedLM(config=a				)
								model.to(a				)
								model.eval()
								_UpperCamelCase						    =	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  A_				(			self       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a				)					->      Optional[int]:


								'''simple docstring'''

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

						def  A_				(			self       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a				)					->      Dict:


								'''simple docstring'''

								_UpperCamelCase						    =	NezhaForPreTraining(config=a				)
								model.to(a				)
								model.eval()
								_UpperCamelCase						    =	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  A_				(			self       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a				)					->      List[str]:


								'''simple docstring'''

								_UpperCamelCase						    =	NezhaForQuestionAnswering(config=a				)
								model.to(a				)
								model.eval()
								_UpperCamelCase						    =	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  A_				(			self       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a				)					->      Optional[int]:


								'''simple docstring'''

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

						def  A_				(			self       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a				)					->      Dict:


								'''simple docstring'''

								_UpperCamelCase						    =	self.num_labels
								_UpperCamelCase						    =	NezhaForTokenClassification(config=a				)
								model.to(a				)
								model.eval()
								_UpperCamelCase						    =	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  A_				(			self       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a				)					->      Any:


								'''simple docstring'''

								_UpperCamelCase						    =	self.num_choices
								_UpperCamelCase						    =	NezhaForMultipleChoice(config=a				)
								model.to(a				)
								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(
								    a       ,							attention_mask=a       ,							token_type_ids=a       ,							labels=a       ,							)
								self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.num_choices)				)







						def  A_				(			self				)					->      Tuple:


								'''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        lowerCAmelCase__      (   __lowercase		,	__lowercase		,	__lowercase		,	unittest.TestCase	):
						UpperCamelCase_							: int	=		(
						    (
						        NezhaModel,
						        NezhaForMaskedLM,
						        NezhaForMultipleChoice,
						        NezhaForNextSentencePrediction,
						        NezhaForPreTraining,
						        NezhaForQuestionAnswering,
						        NezhaForSequenceClassification,
						        NezhaForTokenClassification,
						    )
						    if is_torch_available()
						    else ()
						)
						UpperCamelCase_							: Dict	=		(
						    {
						        "feature-extraction": NezhaModel,
						        "fill-mask": NezhaForMaskedLM,
						        "question-answering": NezhaForQuestionAnswering,
						        "text-classification": NezhaForSequenceClassification,
						        "token-classification": NezhaForTokenClassification,
						        "zero-shot": NezhaForSequenceClassification,
						    }
						    if is_torch_available()
						    else {}
						)
						UpperCamelCase_							: Dict	=		True

						def  A_				(			self       ,							a       ,							a       ,							a=False				)					->      List[Any]:


								'''simple docstring'''

								_UpperCamelCase						    =	super()._prepare_for_class(a       ,							a       ,							return_labels=a				)

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

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


								'''simple docstring'''

								_UpperCamelCase						    =	NezhaModelTester(self				)
								_UpperCamelCase						    =	ConfigTester(self       ,							config_class=a       ,							hidden_size=37				)

						def  A_				(			self				)					->      List[str]:


								'''simple docstring'''

								self.config_tester.run_common_tests()

						def  A_				(			self				)					->      Any:


								'''simple docstring'''

								_UpperCamelCase						    =	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*a				)

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


								'''simple docstring'''

								_UpperCamelCase						    =	self.model_tester.prepare_config_and_inputs_for_decoder()
								self.model_tester.create_and_check_model_as_decoder(*a				)

						def  A_				(			self				)					->      Dict:


								'''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(
								    a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							a       ,							)

						def  A_				(			self				)					->      str:


								'''simple docstring'''

								_UpperCamelCase						    =	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_lm(*a				)

						def  A_				(			self				)					->      Union[str, Any]:


								'''simple docstring'''

								_UpperCamelCase						    =	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_multiple_choice(*a				)

						def  A_				(			self				)					->      str:


								'''simple docstring'''

								_UpperCamelCase						    =	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_next_sequence_prediction(*a				)

						def  A_				(			self				)					->      int:


								'''simple docstring'''

								_UpperCamelCase						    =	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_pretraining(*a				)

						def  A_				(			self				)					->      List[str]:


								'''simple docstring'''

								_UpperCamelCase						    =	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_question_answering(*a				)

						def  A_				(			self				)					->      Any:


								'''simple docstring'''

								_UpperCamelCase						    =	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_sequence_classification(*a				)

						def  A_				(			self				)					->      List[str]:


								'''simple docstring'''

								_UpperCamelCase						    =	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_token_classification(*a				)

						@slow
						def  A_				(			self				)					->      List[Any]:


								'''simple docstring'''

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







						@slow
						@require_torch_gpu
						def  A_				(			self				)					->      Optional[Any]:


								'''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=a				)

										_UpperCamelCase						    =	self._prepare_for_class(a       ,							a				)
										_UpperCamelCase						    =	torch.jit.trace(
										    a       ,							(inputs_dict["""input_ids"""].to("""cpu"""				), inputs_dict["""attention_mask"""].to("""cpu"""				))				)

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







@require_torch
class        lowerCAmelCase__      (   unittest.TestCase	):

						@slow
						def  A_				(			self				)					->      List[str]:


								'''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(a       ,							attention_mask=a				)[0]
								_UpperCamelCase						    =	torch.Size((1, 6, 7_68)				)
								self.assertEqual(output.shape       ,							a				)
								_UpperCamelCase						    =	torch.tensor([[[0.0685, 0.2441, 0.1102], [0.0600, 0.1906, 0.1349], [0.0221, 0.0819, 0.0586]]]				)

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







						@slow
						def  A_				(			self				)					->      List[str]:


								'''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(a       ,							attention_mask=a				)[0]
								_UpperCamelCase						    =	torch.Size((1, 6, 2_11_28)				)
								self.assertEqual(output.shape       ,							a				)
								_UpperCamelCase						    =	torch.tensor(
								    [[-2.7939, -1.7902, -2.2189], [-2.8585, -1.8908, -2.3723], [-2.6499, -1.7750, -2.2558]]				)

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



