'''simple docstring'''




from __future__ import annotations

import unittest

from transformers import AutoTokenizer, GPTJConfig, is_tf_available
from transformers.testing_utils import require_tf, slow, tooslow

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


if is_tf_available():
  import tensorflow as tf

  from transformers.models.gptj.modeling_tf_gptj import (
      TFGPTJForCausalLM,
      TFGPTJForQuestionAnswering,
      TFGPTJForSequenceClassification,
      TFGPTJModel,
      shape_list,
  )

class 			lowerCamelCase_			:
      '''simple docstring'''

      def __init__( self		:       Tuple		,      A		:       Any ):
             _UpperCAmelCase  :						List[Any]							   =					parent
             _UpperCAmelCase  :						str							   =					13
             _UpperCAmelCase  :						Optional[int]							   =					7
             _UpperCAmelCase  :						Optional[Any]							   =					True
             _UpperCAmelCase  :						Dict							   =					True
             _UpperCAmelCase  :						List[Any]							   =					True
             _UpperCAmelCase  :						Optional[int]							   =					True
             _UpperCAmelCase  :						List[str]							   =					True
             _UpperCAmelCase  :						List[str]							   =					99
             _UpperCAmelCase  :						List[str]							   =					32
             _UpperCAmelCase  :						Optional[int]							   =					4
             _UpperCAmelCase  :						Optional[int]							   =					2
             _UpperCAmelCase  :						Optional[int]							   =					4
             _UpperCAmelCase  :						Optional[Any]							   =					37
             _UpperCAmelCase  :						int							   =					"gelu"
             _UpperCAmelCase  :						Union[str, Any]							   =					0.1
             _UpperCAmelCase  :						List[Any]							   =					0.1
             _UpperCAmelCase  :						Optional[int]							   =					512
             _UpperCAmelCase  :						Any							   =					16
             _UpperCAmelCase  :						List[str]							   =					2
             _UpperCAmelCase  :						List[str]							   =					0.02
             _UpperCAmelCase  :						List[Any]							   =					3
             _UpperCAmelCase  :						List[str]							   =					4
             _UpperCAmelCase  :						int							   =					None
             _UpperCAmelCase  :						Optional[Any]							   =					self.vocab_size - 1
             _UpperCAmelCase  :						List[str]							   =					self.vocab_size - 1
             _UpperCAmelCase  :						List[Any]							   =					self.vocab_size - 1




      def   _A			( self		:       List[str] ):
             _UpperCAmelCase  :						Optional[int]							   =					ids_tensor([self.batch_size, self.seq_length]		,      self.vocab_size )

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

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

             _UpperCAmelCase  :						Dict							   =					None
             if self.use_mc_token_ids:
                    _UpperCAmelCase  :						Dict							   =					ids_tensor([self.batch_size, self.num_choices]		,      self.seq_length )

             _UpperCAmelCase  :						Optional[Any]							   =					None
             _UpperCAmelCase  :						Optional[Any]							   =					None
             _UpperCAmelCase  :						int							   =					None
             if self.use_labels:
                    _UpperCAmelCase  :						Any							   =					ids_tensor([self.batch_size]		,      self.type_sequence_label_size )
                    _UpperCAmelCase  :						Optional[int]							   =					ids_tensor([self.batch_size, self.seq_length]		,      self.num_labels )
                    _UpperCAmelCase  :						int							   =					ids_tensor([self.batch_size]		,      self.num_choices )

             _UpperCAmelCase  :						Optional[Any]							   =					GPTJConfig(
                 vocab_size=self.vocab_size		,      n_embd=self.hidden_size		,      n_layer=self.num_hidden_layers		,      n_head=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		,      n_positions=self.max_position_embeddings		,      type_vocab_size=self.type_vocab_size		,      initializer_range=self.initializer_range		,      bos_token_id=self.bos_token_id		,      eos_token_id=self.eos_token_id		,      pad_token_id=self.pad_token_id		,      rotary_dim=self.rotary_dim		,      return_dict=A		,      )

             _UpperCAmelCase  :						Union[str, Any]							   =					ids_tensor([self.num_hidden_layers, self.num_attention_heads]		,      2 )

             return (
                 config,
                 input_ids,
                 input_mask,
                 head_mask,
                 token_type_ids,
                 mc_token_ids,
                 sequence_labels,
                 token_labels,
                 choice_labels,
             )




      def   _A			( self		:       Union[str, Any]		,      A		:       str		,      A		:       Tuple		,      A		:       Union[str, Any]		,      A		:       int		,      A		:       List[Any]		,      *A		:       List[Any] ):
             _UpperCAmelCase  :						Optional[Any]							   =					TFGPTJModel(config=A )
             _UpperCAmelCase  :						str							   =					{
                 "input_ids": input_ids,
                 "attention_mask": input_mask,
                 "token_type_ids": token_type_ids,
             }
             _UpperCAmelCase  :						Optional[Any]							   =					model(A )

             _UpperCAmelCase  :						Union[str, Any]							   =					[input_ids, None, input_mask]  # None is the input for 'past'
             _UpperCAmelCase  :						Union[str, Any]							   =					model(A )

             _UpperCAmelCase  :						Tuple							   =					model(A )

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




      def   _A			( self		:       str		,      A		:       Optional[Any]		,      A		:       Optional[int]		,      A		:       str		,      A		:       int		,      A		:       Any		,      *A		:       Optional[Any] ):
             _UpperCAmelCase  :						Tuple							   =					TFGPTJModel(config=A )

             # first forward pass
             _UpperCAmelCase  :						Union[str, Any]							   =					model(A		,      token_type_ids=A		,      use_cache=A )
             _UpperCAmelCase  :						Any							   =					model(A		,      token_type_ids=A )
             _UpperCAmelCase  :						int							   =					model(A		,      token_type_ids=A		,      use_cache=A )

             self.parent.assertTrue(len(A ) == len(A ) )
             self.parent.assertTrue(len(A ) == len(A ) + 1 )

             _UpperCAmelCase   ,						_UpperCAmelCase  :						int							   =					outputs.to_tuple()

             # create hypothetical next token and extent to next_input_ids
             _UpperCAmelCase  :						Optional[int]							   =					ids_tensor((self.batch_size, 1)		,      config.vocab_size )
             _UpperCAmelCase  :						int							   =					ids_tensor([self.batch_size, 1]		,      self.type_vocab_size )

             # append to next input_ids and token_type_ids
             _UpperCAmelCase  :						int							   =					tf.concat([input_ids, next_tokens]		,      axis=-1 )
             _UpperCAmelCase  :						Optional[Any]							   =					tf.concat([token_type_ids, next_token_types]		,      axis=-1 )

             _UpperCAmelCase  :						Dict							   =					model(A		,      token_type_ids=A )["last_hidden_state"]
             _UpperCAmelCase  :						List[Any]							   =					model(A		,      token_type_ids=A		,      past_key_values=A )[
                 "last_hidden_state"
             ]

             # select random slice
             _UpperCAmelCase  :						int							   =					int(ids_tensor((1,)		,      shape_list(A )[-1] ) )
             _UpperCAmelCase  :						Any							   =					output_from_no_past[:, -1, random_slice_idx]
             _UpperCAmelCase  :						str							   =					output_from_past[:, 0, random_slice_idx]

             # test that outputs are equal for slice
             tf.debugging.assert_near(A		,      A		,      rtol=1E-6 )




      def   _A			( self		:       Any		,      A		:       int		,      A		:       Optional[int]		,      A		:       Union[str, Any]		,      A		:       Optional[Any]		,      A		:       List[str]		,      *A		:       Optional[Any] ):
             _UpperCAmelCase  :						Optional[int]							   =					TFGPTJModel(config=A )

             # create attention mask
             _UpperCAmelCase  :						Union[str, Any]							   =					self.seq_length // 2
             _UpperCAmelCase  :						Tuple							   =					tf.ones((self.batch_size, half_seq_length)		,      dtype=tf.intaa )
             _UpperCAmelCase  :						Dict							   =					tf.zeros((self.batch_size, self.seq_length - half_seq_length)		,      dtype=tf.intaa )
             _UpperCAmelCase  :						str							   =					tf.concat([attn_mask_begin, attn_mask_end]		,      axis=1 )

             # first forward pass
             _UpperCAmelCase   ,						_UpperCAmelCase  :						Tuple							   =					model(A		,      attention_mask=A ).to_tuple()

             # create hypothetical next token and extent to next_input_ids
             _UpperCAmelCase  :						List[Any]							   =					ids_tensor((self.batch_size, 1)		,      config.vocab_size )

             # change a random masked slice from input_ids
             _UpperCAmelCase  :						Optional[Any]							   =					ids_tensor((1,)		,      A ).numpy() + 1
             _UpperCAmelCase  :						Dict							   =					ids_tensor((self.batch_size, self.seq_length)		,      config.vocab_size )
             _UpperCAmelCase  :						List[Any]							   =					tf.range(self.seq_length ) == (self.seq_length - random_seq_idx_to_change)
             _UpperCAmelCase  :						str							   =					tf.transpose(
                 tf.broadcast_to(tf.expand_dims(A		,      -1 )		,      (self.seq_length, self.batch_size) ) )
             _UpperCAmelCase  :						Dict							   =					tf.where(A		,      A		,      A )

             # append to next input_ids and attn_mask
             _UpperCAmelCase  :						List[str]							   =					tf.concat([input_ids, next_tokens]		,      axis=-1 )
             _UpperCAmelCase  :						Optional[int]							   =					tf.concat([attn_mask, tf.ones((shape_list(A )[0], 1)		,      dtype=tf.intaa )]		,      axis=1 )

             # get two different outputs
             _UpperCAmelCase  :						str							   =					model(A		,      attention_mask=A )["last_hidden_state"]
             _UpperCAmelCase  :						Dict							   =					model(A		,      past_key_values=A		,      attention_mask=A )[
                 "last_hidden_state"
             ]

             # select random slice
             _UpperCAmelCase  :						Tuple							   =					int(ids_tensor((1,)		,      shape_list(A )[-1] ) )
             _UpperCAmelCase  :						Dict							   =					output_from_no_past[:, -1, random_slice_idx]
             _UpperCAmelCase  :						Optional[Any]							   =					output_from_past[:, 0, random_slice_idx]

             # test that outputs are equal for slice
             tf.debugging.assert_near(A		,      A		,      rtol=1E-12 )




      def   _A			( self		:       Union[str, Any]		,      A		:       Dict		,      A		:       Dict		,      A		:       Optional[int]		,      A		:       str		,      A		:       str		,      *A		:       Union[str, Any] ):
             _UpperCAmelCase  :						Optional[Any]							   =					TFGPTJModel(config=A )

             _UpperCAmelCase  :						Tuple							   =					input_ids[:1, :]
             _UpperCAmelCase  :						Tuple							   =					input_mask[:1, :]
             _UpperCAmelCase  :						Any							   =					token_type_ids[:1, :]
             _UpperCAmelCase  :						int							   =					1

             # first forward pass
             _UpperCAmelCase  :						List[Any]							   =					model(A		,      attention_mask=A		,      token_type_ids=A		,      use_cache=A )

             _UpperCAmelCase   ,						_UpperCAmelCase  :						Optional[Any]							   =					outputs.to_tuple()

             # create hypothetical next token and extent to next_input_ids
             _UpperCAmelCase  :						Optional[int]							   =					ids_tensor((self.batch_size, 3)		,      config.vocab_size )
             _UpperCAmelCase  :						List[Any]							   =					ids_tensor((self.batch_size, 3)		,      2 )
             _UpperCAmelCase  :						Union[str, Any]							   =					ids_tensor((self.batch_size, 3)		,      self.type_vocab_size )

             # append to next input_ids and token_type_ids
             _UpperCAmelCase  :						List[str]							   =					tf.concat([input_ids, next_tokens]		,      axis=-1 )
             _UpperCAmelCase  :						int							   =					tf.concat([input_mask, next_attn_mask]		,      axis=-1 )
             _UpperCAmelCase  :						Tuple							   =					tf.concat([token_type_ids, next_token_types]		,      axis=-1 )

             _UpperCAmelCase  :						Optional[int]							   =					model(
                 A		,      token_type_ids=A		,      attention_mask=A )["last_hidden_state"]
             _UpperCAmelCase  :						Optional[Any]							   =					model(
                 A		,      token_type_ids=A		,      attention_mask=A		,      past_key_values=A		,      )["last_hidden_state"]
             self.parent.assertTrue(output_from_past.shape[1] == next_tokens.shape[1] )

             # select random slice
             _UpperCAmelCase  :						Union[str, Any]							   =					int(ids_tensor((1,)		,      shape_list(A )[-1] ) )
             _UpperCAmelCase  :						Optional[Any]							   =					output_from_no_past[:, -3:, random_slice_idx]
             _UpperCAmelCase  :						Tuple							   =					output_from_past[:, :, random_slice_idx]

             # test that outputs are equal for slice
             tf.debugging.assert_near(A		,      A		,      rtol=1E-3 )




      def   _A			( self		:       Optional[int]		,      A		:       Any		,      A		:       List[Any]		,      A		:       List[Any]		,      A		:       Optional[int]		,      A		:       Tuple		,      *A		:       int ):
             _UpperCAmelCase  :						List[str]							   =					TFGPTJForCausalLM(config=A )
             _UpperCAmelCase  :						Tuple							   =					{
                 "input_ids": input_ids,
                 "attention_mask": input_mask,
                 "token_type_ids": token_type_ids,
             }
             _UpperCAmelCase  :						Tuple							   =					model(A )
             self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.seq_length, self.vocab_size) )





      def   _A			( self		:       List[Any] ):
             _UpperCAmelCase  :						Optional[int]							   =					self.prepare_config_and_inputs()

             (
                 (
                 _UpperCAmelCase
             )   ,						(
                 _UpperCAmelCase
             )   ,						(
                 _UpperCAmelCase
             )   ,						(
                 _UpperCAmelCase
             )   ,						(
                 _UpperCAmelCase
             )   ,						(
                 _UpperCAmelCase
             )   ,						(
                 _UpperCAmelCase
             )   ,						(
                 _UpperCAmelCase
             )   ,						(
                 _UpperCAmelCase
             )   ,						
             )  :						Any							   =					config_and_inputs

             _UpperCAmelCase  :						str							   =					{
                 "input_ids": input_ids,
                 "token_type_ids": token_type_ids,
                 "attention_mask": input_mask,
             }
             return config, inputs_dict

@require_tf
class 			lowerCamelCase_			(snake_case__							,				snake_case__							,				snake_case__							,				unittest.TestCase					):
      '''simple docstring'''

      __UpperCamelCase:						List[Any]      			=		(
          (TFGPTJForCausalLM, TFGPTJForSequenceClassification, TFGPTJForQuestionAnswering, TFGPTJModel)
          if is_tf_available()
          else ()
      )

      __UpperCamelCase:						Dict      			=		(TFGPTJForCausalLM,) if is_tf_available() else ()
      __UpperCamelCase:						Optional[int]      			=		(
          {
              "feature-extraction": TFGPTJModel,
              "question-answering": TFGPTJForQuestionAnswering,
              "text-classification": TFGPTJForSequenceClassification,
              "text-generation": TFGPTJForCausalLM,
              "zero-shot": TFGPTJForSequenceClassification,
          }
          if is_tf_available()
          else {}
      )
      __UpperCamelCase:						Optional[Any]      			=		False
      __UpperCamelCase:						Optional[int]      			=		False
      __UpperCamelCase:						Dict      			=		False
      __UpperCamelCase:						Any      			=		False




      def   _A			( self		:       str		,      A		:       Union[str, Any]		,      A		:       Optional[int]		,      A		:       Any		,      A		:       Any		,      A		:       Dict ):
             if (
                 pipeline_test_casse_name == "QAPipelineTests"
                 and tokenizer_name is not None
                 and not tokenizer_name.endswith("Fast" )
             ):
                    # `QAPipelineTests` fails for a few models when the slower tokenizer are used.
                    # (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
                    # TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
                    return True

             return False




      def   _A			( self		:       int ):
             _UpperCAmelCase  :						Tuple							   =					TFGPTJModelTester(self )
             _UpperCAmelCase  :						Optional[int]							   =					ConfigTester(self		,      config_class=A		,      n_embd=37 )




      def   _A			( self		:       Union[str, Any] ):
             self.config_tester.run_common_tests()




      def   _A			( self		:       List[Any] ):
             _UpperCAmelCase  :						Any							   =					self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_gptj_model(*A )




      def   _A			( self		:       str ):
             _UpperCAmelCase  :						Tuple							   =					self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_gptj_model_past(*A )




      def   _A			( self		:       Optional[Any] ):
             _UpperCAmelCase  :						str							   =					self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_gptj_model_attention_mask_past(*A )




      def   _A			( self		:       Tuple ):
             _UpperCAmelCase  :						Optional[int]							   =					self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_gptj_model_past_large_inputs(*A )




      def   _A			( self		:       int ):
             _UpperCAmelCase  :						Dict							   =					self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_gptj_lm_head_model(*A )




      @slow
      @unittest.skipIf(
          not is_tf_available() or len(tf.config.list_physical_devices("GPU" ) ) > 0		,      "skip testing on GPU for now to avoid GPU OOM."		,      )
      def   _A			( self		:       Dict ):
             _UpperCAmelCase  :						Any							   =					TFGPTJModel.from_pretrained("EleutherAI/gpt-j-6B"		,      from_pt=A )
             self.assertIsNotNone(A )





      @unittest.skip(reason="Currently, model embeddings are going to undergo a major refactor." )
      def   _A			( self		:       Optional[Any] ):
             super().test_resize_token_embeddings()

@require_tf
@tooslow
# Marked as @tooslow due to GPU OOM -- but still useful to run locally. Requires ~39GB of RAM.
class 			lowerCamelCase_			(unittest.TestCase					):
      '''simple docstring'''

      def   _A			( self		:       Tuple ):
             _UpperCAmelCase  :						str							   =					TFGPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B"		,      from_pt=A )
             _UpperCAmelCase  :						Union[str, Any]							   =					tf.convert_to_tensor([[464, 3290]]		,      dtype=tf.intaa )  # The dog
             # fmt: off
             # The dog is a man's best friend. It is a loyal companion, and it is a friend
             _UpperCAmelCase  :						str							   =					[464, 3290, 318, 257, 582, 338, 1266, 1545, 13, 632, 318, 257, 9112, 15185, 11, 290, 340, 318, 257, 1545]
             # fmt: on
             _UpperCAmelCase  :						Any							   =					model.generate(A		,      do_sample=A )
             self.assertListEqual(output_ids[0].numpy().tolist()		,      A )




      def   _A			( self		:       Optional[int] ):
             _UpperCAmelCase  :						int							   =					AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B"		,      revision="float16" )
             _UpperCAmelCase  :						Dict							   =					TFGPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B"		,      revision="float16"		,      from_pt=A )

             _UpperCAmelCase  :						str							   =					tokenizer("Today is a nice day and"		,      return_tensors="tf" )
             # forces the generation to happen on CPU, to avoid GPU-related quirks
             with tf.device(":/CPU:0" ):
                    _UpperCAmelCase  :						str							   =					model.generate(**A		,      do_sample=A		,      seed=[42, 0] )
             _UpperCAmelCase  :						Tuple							   =					tokenizer.decode(output_ids[0]		,      skip_special_tokens=A )

             _UpperCAmelCase  :						Union[str, Any]							   =					"Today is a nice day and I’m going to go for a walk. I’"
             self.assertEqual(A		,      A )




      def   _A			( self		:       str ):
             _UpperCAmelCase  :						str							   =					TFGPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B"		,      revision="float16"		,      from_pt=A )
             _UpperCAmelCase  :						Union[str, Any]							   =					AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B"		,      revision="float16" )

             _UpperCAmelCase  :						Any							   =					"left"

             # Define PAD Token = EOS Token = 50256
             _UpperCAmelCase  :						Optional[Any]							   =					tokenizer.eos_token
             _UpperCAmelCase  :						List[Any]							   =					model.config.eos_token_id

             # use different length sentences to test batching
             _UpperCAmelCase  :						List[Any]							   =					[
                 "Hello, my dog is a little",
                 "Today, I",
             ]
             _UpperCAmelCase  :						Tuple							   =					[
                 "Hello, my dog is a little over a year old and has been diagnosed with hip dysplasia",
                 "Today, I’m going to be talking about a topic that’",
             ]
             return model, tokenizer, sentences, expected_output_sentences




      def   _A			( self		:       List[Any] ):
             # Confirms that we get the expected results with left-padded beam search
             _UpperCAmelCase   ,						_UpperCAmelCase   ,						_UpperCAmelCase   ,						_UpperCAmelCase  :						Optional[Any]							   =					self._get_beam_search_test_objects()

             _UpperCAmelCase  :						List[Any]							   =					tokenizer(A		,      return_tensors="tf"		,      padding=A )
             _UpperCAmelCase  :						int							   =					model.generate(**A		,      do_sample=A		,      num_beams=2 )
             _UpperCAmelCase  :						Tuple							   =					tokenizer.batch_decode(A		,      skip_special_tokens=A )
             self.assertListEqual(A		,      A )




      def   _A			( self		:       Tuple ):
             # Confirms that left-padding is working properly
             _UpperCAmelCase   ,						_UpperCAmelCase   ,						_UpperCAmelCase   ,						_UpperCAmelCase  :						Union[str, Any]							   =					self._get_beam_search_test_objects()

             _UpperCAmelCase  :						int							   =					tokenizer(A		,      return_tensors="tf"		,      padding=A )
             _UpperCAmelCase  :						Optional[Any]							   =					tokenizer(sentences[0]		,      return_tensors="tf" )
             _UpperCAmelCase  :						Tuple							   =					model.generate(**A		,      do_sample=A		,      num_beams=2 )
             _UpperCAmelCase  :						List[str]							   =					(
                 shape_list(inputs_non_padded["input_ids"] )[-1]
                 - tf.reduce_sum(tf.cast(inputs["attention_mask"][-1]		,      tf.intaa ) ).numpy()
             )
             _UpperCAmelCase  :						str							   =					tokenizer(sentences[1]		,      return_tensors="tf" )
             _UpperCAmelCase  :						int							   =					model.generate(
                 **A		,      do_sample=A		,      num_beams=2		,      max_length=model.config.max_length - num_paddings )
             _UpperCAmelCase  :						str							   =					tokenizer.decode(output_non_padded[0]		,      skip_special_tokens=A )
             _UpperCAmelCase  :						int							   =					tokenizer.decode(output_padded[0]		,      skip_special_tokens=A )
             self.assertListEqual(A		,      [non_padded_sentence, padded_sentence] )





      def   _A			( self		:       Optional[Any] ):
             # Confirms that XLA is working properly
             _UpperCAmelCase   ,						_UpperCAmelCase   ,						_UpperCAmelCase   ,						_UpperCAmelCase  :						str							   =					self._get_beam_search_test_objects()

             _UpperCAmelCase  :						Any							   =					tokenizer(A		,      return_tensors="tf"		,      padding=A )
             _UpperCAmelCase  :						Union[str, Any]							   =					tf.function(model.generate		,      jit_compile=A )
             _UpperCAmelCase  :						Optional[Any]							   =					xla_generate(**A		,      do_sample=A		,      num_beams=2 )
             _UpperCAmelCase  :						str							   =					tokenizer.batch_decode(A		,      skip_special_tokens=A )
             self.assertListEqual(A		,      A )


'''simple docstring'''




import unittest

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

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,
      ElectraForCausalLM,
      ElectraForMaskedLM,
      ElectraForMultipleChoice,
      ElectraForPreTraining,
      ElectraForQuestionAnswering,
      ElectraForSequenceClassification,
      ElectraForTokenClassification,
      ElectraModel,
  )
  from transformers.models.electra.modeling_electra import ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST

class 			lowerCamelCase_			:
      '''simple docstring'''

      def __init__( self		:       Dict		,      A		:       Union[str, Any]		,      A		:       int=13		,      A		:       List[str]=7		,      A		:       Tuple=True		,      A		:       Optional[Any]=True		,      A		:       Tuple=True		,      A		:       List[str]=True		,      A		:       List[str]=99		,      A		:       List[Any]=32		,      A		:       List[str]=5		,      A		:       Dict=4		,      A		:       str=37		,      A		:       str="gelu"		,      A		:       int=0.1		,      A		:       List[str]=0.1		,      A		:       List[str]=512		,      A		:       Optional[int]=16		,      A		:       List[Any]=2		,      A		:       Optional[Any]=0.02		,      A		:       Optional[int]=3		,      A		:       List[Any]=4		,      A		:       int=None		,      ):
             _UpperCAmelCase  :						int							   =					parent
             _UpperCAmelCase  :						Any							   =					batch_size
             _UpperCAmelCase  :						Optional[int]							   =					seq_length
             _UpperCAmelCase  :						Any							   =					is_training
             _UpperCAmelCase  :						Union[str, Any]							   =					use_input_mask
             _UpperCAmelCase  :						Optional[Any]							   =					use_token_type_ids
             _UpperCAmelCase  :						Dict							   =					use_labels
             _UpperCAmelCase  :						Tuple							   =					vocab_size
             _UpperCAmelCase  :						Optional[int]							   =					hidden_size
             _UpperCAmelCase  :						List[Any]							   =					num_hidden_layers
             _UpperCAmelCase  :						Union[str, Any]							   =					num_attention_heads
             _UpperCAmelCase  :						List[str]							   =					intermediate_size
             _UpperCAmelCase  :						Optional[int]							   =					hidden_act
             _UpperCAmelCase  :						List[Any]							   =					hidden_dropout_prob
             _UpperCAmelCase  :						Union[str, Any]							   =					attention_probs_dropout_prob
             _UpperCAmelCase  :						str							   =					max_position_embeddings
             _UpperCAmelCase  :						Any							   =					type_vocab_size
             _UpperCAmelCase  :						str							   =					type_sequence_label_size
             _UpperCAmelCase  :						List[str]							   =					initializer_range
             _UpperCAmelCase  :						Union[str, Any]							   =					num_labels
             _UpperCAmelCase  :						Dict							   =					num_choices
             _UpperCAmelCase  :						List[str]							   =					scope




      def   _A			( self		:       Any ):
             _UpperCAmelCase  :						Union[str, Any]							   =					ids_tensor([self.batch_size, self.seq_length]		,      self.vocab_size )

             _UpperCAmelCase  :						int							   =					None
             if self.use_input_mask:
                    _UpperCAmelCase  :						List[str]							   =					random_attention_mask([self.batch_size, self.seq_length] )

             _UpperCAmelCase  :						List[str]							   =					None
             if self.use_token_type_ids:
                    _UpperCAmelCase  :						Optional[Any]							   =					ids_tensor([self.batch_size, self.seq_length]		,      self.type_vocab_size )

             _UpperCAmelCase  :						Optional[int]							   =					None
             _UpperCAmelCase  :						int							   =					None
             _UpperCAmelCase  :						Tuple							   =					None
             if self.use_labels:
                    _UpperCAmelCase  :						str							   =					ids_tensor([self.batch_size]		,      self.type_sequence_label_size )
                    _UpperCAmelCase  :						Any							   =					ids_tensor([self.batch_size, self.seq_length]		,      self.num_labels )
                    _UpperCAmelCase  :						Optional[Any]							   =					ids_tensor([self.batch_size]		,      self.num_choices )
                    _UpperCAmelCase  :						Tuple							   =					ids_tensor([self.batch_size, self.seq_length]		,      1 )

             _UpperCAmelCase  :						Any							   =					self.get_config()

             return (
                 config,
                 input_ids,
                 token_type_ids,
                 input_mask,
                 sequence_labels,
                 token_labels,
                 choice_labels,
                 fake_token_labels,
             )




      def   _A			( self		:       List[str] ):
             return ElectraConfig(
                 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		:       Union[str, Any] ):
             (
                 (
                 _UpperCAmelCase
             )   ,						(
                 _UpperCAmelCase
             )   ,						(
                 _UpperCAmelCase
             )   ,						(
                 _UpperCAmelCase
             )   ,						(
                 _UpperCAmelCase
             )   ,						(
                 _UpperCAmelCase
             )   ,						(
                 _UpperCAmelCase
             )   ,						(
                 _UpperCAmelCase
             )   ,						
             )  :						Dict							   =					self.prepare_config_and_inputs()

             _UpperCAmelCase  :						int							   =					True
             _UpperCAmelCase  :						Any							   =					floats_tensor([self.batch_size, self.seq_length, self.hidden_size] )
             _UpperCAmelCase  :						List[Any]							   =					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		:       Union[str, Any]		,      A		:       Dict		,      A		:       int		,      A		:       Any		,      A		:       Tuple		,      A		:       Union[str, Any]		,      A		:       int		,      A		:       str		,      A		:       Union[str, Any]		,      ):
             _UpperCAmelCase  :						Any							   =					ElectraModel(config=A )
             model.to(A )
             model.eval()
             _UpperCAmelCase  :						Optional[int]							   =					model(A		,      attention_mask=A		,      token_type_ids=A )
             _UpperCAmelCase  :						Tuple							   =					model(A		,      token_type_ids=A )
             _UpperCAmelCase  :						List[str]							   =					model(A )
             self.parent.assertEqual(result.last_hidden_state.shape		,      (self.batch_size, self.seq_length, self.hidden_size) )




      def   _A			( self		:       List[Any]		,      A		:       Any		,      A		:       Optional[int]		,      A		:       List[Any]		,      A		:       Dict		,      A		:       List[Any]		,      A		:       Optional[int]		,      A		:       List[str]		,      A		:       Dict		,      A		:       str		,      ):
             _UpperCAmelCase  :						Optional[Any]							   =					True
             _UpperCAmelCase  :						List[str]							   =					ElectraModel(A )
             model.to(A )
             model.eval()
             _UpperCAmelCase  :						int							   =					model(
                 A		,      attention_mask=A		,      token_type_ids=A		,      encoder_hidden_states=A		,      encoder_attention_mask=A		,      )
             _UpperCAmelCase  :						Tuple							   =					model(
                 A		,      attention_mask=A		,      token_type_ids=A		,      encoder_hidden_states=A		,      )
             _UpperCAmelCase  :						Union[str, Any]							   =					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) )




      def   _A			( self		:       Optional[int]		,      A		:       int		,      A		:       Optional[Any]		,      A		:       int		,      A		:       Tuple		,      A		:       List[str]		,      A		:       Optional[Any]		,      A		:       str		,      A		:       Dict		,      ):
             _UpperCAmelCase  :						Optional[Any]							   =					ElectraForMaskedLM(config=A )
             model.to(A )
             model.eval()
             _UpperCAmelCase  :						int							   =					model(A		,      attention_mask=A		,      token_type_ids=A		,      labels=A )
             self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.seq_length, self.vocab_size) )




      def   _A			( self		:       Optional[Any]		,      A		:       Optional[int]		,      A		:       Optional[int]		,      A		:       Any		,      A		:       str		,      A		:       Optional[int]		,      A		:       Any		,      A		:       str		,      A		:       int		,      A		:       Dict		,      ):
             _UpperCAmelCase  :						Tuple							   =					ElectraForCausalLM(config=A )
             model.to(A )
             model.eval()
             _UpperCAmelCase  :						List[Any]							   =					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		:       List[str]		,      A		:       Tuple		,      A		:       Dict		,      A		:       Tuple		,      A		:       List[str]		,      A		:       List[Any]		,      A		:       Optional[Any]		,      A		:       str		,      A		:       Dict		,      ):
             _UpperCAmelCase  :						Any							   =					self.num_labels
             _UpperCAmelCase  :						Tuple							   =					ElectraForTokenClassification(config=A )
             model.to(A )
             model.eval()
             _UpperCAmelCase  :						Union[str, Any]							   =					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		:       Union[str, Any]		,      A		:       Optional[Any]		,      A		:       Optional[int]		,      A		:       Union[str, Any]		,      A		:       str		,      A		:       Union[str, Any]		,      A		:       Tuple		,      A		:       Optional[int]		,      A		:       Optional[int]		,      ):
             _UpperCAmelCase  :						Union[str, Any]							   =					self.num_labels
             _UpperCAmelCase  :						int							   =					ElectraForPreTraining(config=A )
             model.to(A )
             model.eval()
             _UpperCAmelCase  :						Dict							   =					model(A		,      attention_mask=A		,      token_type_ids=A		,      labels=A )
             self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.seq_length) )




      def   _A			( self		:       Union[str, Any]		,      A		:       List[Any]		,      A		:       int		,      A		:       int		,      A		:       Optional[int]		,      A		:       Optional[Any]		,      A		:       Dict		,      A		:       str		,      A		:       Any		,      ):
             _UpperCAmelCase  :						Union[str, Any]							   =					self.num_labels
             _UpperCAmelCase  :						List[str]							   =					ElectraForSequenceClassification(A )
             model.to(A )
             model.eval()
             _UpperCAmelCase  :						str							   =					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		:       List[str]		,      A		:       Optional[Any]		,      A		:       Optional[int]		,      A		:       Tuple		,      A		:       Optional[Any]		,      A		:       int		,      A		:       List[str]		,      A		:       Tuple		,      A		:       Optional[int]		,      ):
             _UpperCAmelCase  :						List[str]							   =					ElectraForQuestionAnswering(config=A )
             model.to(A )
             model.eval()
             _UpperCAmelCase  :						Any							   =					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		:       int		,      A		:       Optional[Any]		,      A		:       Optional[Any]		,      A		:       Tuple		,      A		:       Dict		,      A		:       Optional[Any]		,      A		:       Union[str, Any]		,      A		:       Optional[Any]		,      A		:       str		,      ):
             _UpperCAmelCase  :						Tuple							   =					self.num_choices
             _UpperCAmelCase  :						int							   =					ElectraForMultipleChoice(config=A )
             model.to(A )
             model.eval()
             _UpperCAmelCase  :						str							   =					input_ids.unsqueeze(1 ).expand(-1		,      self.num_choices		,      -1 ).contiguous()
             _UpperCAmelCase  :						Dict							   =					token_type_ids.unsqueeze(1 ).expand(-1		,      self.num_choices		,      -1 ).contiguous()
             _UpperCAmelCase  :						Union[str, Any]							   =					input_mask.unsqueeze(1 ).expand(-1		,      self.num_choices		,      -1 ).contiguous()
             _UpperCAmelCase  :						Union[str, Any]							   =					model(
                 A		,      attention_mask=A		,      token_type_ids=A		,      labels=A		,      )
             self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.num_choices) )





      def   _A			( self		:       Any ):
             _UpperCAmelCase  :						List[Any]							   =					self.prepare_config_and_inputs()
             (
                 (
                 _UpperCAmelCase
             )   ,						(
                 _UpperCAmelCase
             )   ,						(
                 _UpperCAmelCase
             )   ,						(
                 _UpperCAmelCase
             )   ,						(
                 _UpperCAmelCase
             )   ,						(
                 _UpperCAmelCase
             )   ,						(
                 _UpperCAmelCase
             )   ,						(
                 _UpperCAmelCase
             )   ,						
             )  :						Tuple							   =					config_and_inputs
             _UpperCAmelCase  :						List[Any]							   =					{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
             return config, inputs_dict

@require_torch
class 			lowerCamelCase_			(snake_case__							,				snake_case__							,				unittest.TestCase					):
      '''simple docstring'''

      __UpperCamelCase:						Tuple      			=		(
          (
              ElectraModel,
              ElectraForPreTraining,
              ElectraForMaskedLM,
              ElectraForCausalLM,
              ElectraForMultipleChoice,
              ElectraForTokenClassification,
              ElectraForSequenceClassification,
              ElectraForQuestionAnswering,
          )
          if is_torch_available()
          else ()
      )
      __UpperCamelCase:						Optional[Any]      			=		(
          {
              "feature-extraction": ElectraModel,
              "fill-mask": ElectraForMaskedLM,
              "question-answering": ElectraForQuestionAnswering,
              "text-classification": ElectraForSequenceClassification,
              "text-generation": ElectraForCausalLM,
              "token-classification": ElectraForTokenClassification,
              "zero-shot": ElectraForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      __UpperCamelCase:						str      			=		True




      def   _A			( self		:       Optional[Any]		,      A		:       Tuple		,      A		:       Optional[Any]		,      A		:       List[Any]=False ):
             _UpperCAmelCase  :						int							   =					super()._prepare_for_class(A		,      A		,      return_labels=A )

             if return_labels:
                    if model_class in get_values(A ):
                           _UpperCAmelCase  :						List[str]							   =					torch.zeros(
                               (self.model_tester.batch_size, self.model_tester.seq_length)		,      dtype=torch.long		,      device=A )
             return inputs_dict




      def   _A			( self		:       Optional[int] ):
             _UpperCAmelCase  :						Any							   =					ElectraModelTester(self )
             _UpperCAmelCase  :						Tuple							   =					ConfigTester(self		,      config_class=A		,      hidden_size=37 )




      def   _A			( self		:       Tuple ):
             self.config_tester.run_common_tests()




      def   _A			( self		:       Dict ):
             _UpperCAmelCase  :						Dict							   =					self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_electra_model(*A )




      def   _A			( self		:       int ):
             _UpperCAmelCase  :						List[Any]							   =					self.model_tester.prepare_config_and_inputs_for_decoder()
             self.model_tester.create_and_check_electra_model_as_decoder(*A )




      def   _A			( self		:       Union[str, Any] ):
             _UpperCAmelCase  :						Optional[Any]							   =					self.model_tester.prepare_config_and_inputs()
             for type in ["absolute", "relative_key", "relative_key_query"]:
                    _UpperCAmelCase  :						Optional[int]							   =					type
                    self.model_tester.create_and_check_electra_model(*A )




      def   _A			( self		:       Optional[Any] ):
             _UpperCAmelCase  :						Optional[Any]							   =					self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_electra_for_masked_lm(*A )




      def   _A			( self		:       Any ):
             _UpperCAmelCase  :						Optional[Any]							   =					self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_electra_for_token_classification(*A )




      def   _A			( self		:       str ):
             _UpperCAmelCase  :						int							   =					self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_electra_for_pretraining(*A )




      def   _A			( self		:       Tuple ):
             _UpperCAmelCase  :						List[str]							   =					self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_electra_for_sequence_classification(*A )




      def   _A			( self		:       Dict ):
             _UpperCAmelCase  :						List[Any]							   =					self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_electra_for_question_answering(*A )




      def   _A			( self		:       Dict ):
             _UpperCAmelCase  :						Optional[int]							   =					self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_electra_for_multiple_choice(*A )




      @slow
      def   _A			( self		:       Union[str, Any] ):
             for model_name in ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    _UpperCAmelCase  :						List[Any]							   =					ElectraModel.from_pretrained(A )
                    self.assertIsNotNone(A )





      def   _A			( self		:       List[Any] ):
             _UpperCAmelCase  :						Union[str, Any]							   =					self.model_tester.prepare_config_and_inputs_for_decoder()
             self.model_tester.create_and_check_electra_for_causal_lm(*A )

@require_torch
class 			lowerCamelCase_			(unittest.TestCase					):
      '''simple docstring'''

      @slow
      def   _A			( self		:       List[Any] ):
             _UpperCAmelCase  :						List[Any]							   =					ElectraModel.from_pretrained("google/electra-small-discriminator" )
             _UpperCAmelCase  :						List[Any]							   =					torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]] )
             _UpperCAmelCase  :						List[str]							   =					torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]] )
             _UpperCAmelCase  :						str							   =					model(A		,      attention_mask=A )[0]
             _UpperCAmelCase  :						Union[str, Any]							   =					torch.Size((1, 11, 256) )
             self.assertEqual(output.shape		,      A )
             _UpperCAmelCase  :						Optional[Any]							   =					torch.tensor(
                 [[[0.4_471, 0.6_821, -0.3_265], [0.4_627, 0.5_255, -0.3_668], [0.4_532, 0.3_313, -0.4_344]]] )

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

