




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      A     :
     """simple docstring"""

     def __init__(      self					:				List[Any]		,					A__					:				Optional[Any]     )			->    Any:
           lowerCamelCase__ 					=			parent
           lowerCamelCase__ 					=			1_3
           lowerCamelCase__ 					=			7
           lowerCamelCase__ 					=			True
           lowerCamelCase__ 					=			True
           lowerCamelCase__ 					=			True
           lowerCamelCase__ 					=			True
           lowerCamelCase__ 					=			True
           lowerCamelCase__ 					=			9_9
           lowerCamelCase__ 					=			3_2
           lowerCamelCase__ 					=			4
           lowerCamelCase__ 					=			2
           lowerCamelCase__ 					=			4
           lowerCamelCase__ 					=			3_7
           lowerCamelCase__ 					=			"""gelu"""
           lowerCamelCase__ 					=			0.1
           lowerCamelCase__ 					=			0.1
           lowerCamelCase__ 					=			5_1_2
           lowerCamelCase__ 					=			1_6
           lowerCamelCase__ 					=			2
           lowerCamelCase__ 					=			0.02
           lowerCamelCase__ 					=			3
           lowerCamelCase__ 					=			4
           lowerCamelCase__ 					=			None
           lowerCamelCase__ 					=			self.vocab_size - 1
           lowerCamelCase__ 					=			self.vocab_size - 1
           lowerCamelCase__ 					=			self.vocab_size - 1




     def 	__snake_case				(      self					:				str     )			->    List[Any]:
           lowerCamelCase__ 					=			ids_tensor([self.batch_size, self.seq_length]		,					self.vocab_size     )

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

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

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

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

           lowerCamelCase__ 					=			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=snake_case_		,					)

           lowerCamelCase__ 					=			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 	__snake_case				(      self					:				Optional[int]		,					A__					:				List[Any]		,					A__					:				Union[str, Any]		,					A__					:				List[str]		,					A__					:				List[str]		,					A__					:				int		,					*A__					:				List[Any]     )			->    Union[str, Any]:
           lowerCamelCase__ 					=			TFGPTJModel(config=snake_case_     )
           lowerCamelCase__ 					=			{
               """input_ids""": input_ids,
               """attention_mask""": input_mask,
               """token_type_ids""": token_type_ids,
           }
           lowerCamelCase__ 					=			model(snake_case_     )

           lowerCamelCase__ 					=			[input_ids, None, input_mask]  # None is the input for 'past'
           lowerCamelCase__ 					=			model(snake_case_     )

           lowerCamelCase__ 					=			model(snake_case_     )

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




     def 	__snake_case				(      self					:				Any		,					A__					:				str		,					A__					:				List[Any]		,					A__					:				Optional[Any]		,					A__					:				int		,					A__					:				Optional[Any]		,					*A__					:				Dict     )			->    int:
           lowerCamelCase__ 					=			TFGPTJModel(config=snake_case_     )

           # first forward pass
           lowerCamelCase__ 					=			model(snake_case_		,					token_type_ids=snake_case_		,					use_cache=snake_case_     )
           lowerCamelCase__ 					=			model(snake_case_		,					token_type_ids=snake_case_     )
           lowerCamelCase__ 					=			model(snake_case_		,					token_type_ids=snake_case_		,					use_cache=snake_case_     )

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

           lowerCamelCase__ 					=			outputs.to_tuple()

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

           # append to next input_ids and token_type_ids
           lowerCamelCase__ 					=			tf.concat([input_ids, next_tokens]		,					axis=-1     )
           lowerCamelCase__ 					=			tf.concat([token_type_ids, next_token_types]		,					axis=-1     )

           lowerCamelCase__ 					=			model(snake_case_		,					token_type_ids=snake_case_     )["""last_hidden_state"""]
           lowerCamelCase__ 					=			model(snake_case_		,					token_type_ids=snake_case_		,					past_key_values=snake_case_     )[
               """last_hidden_state"""
           ]

           # select random slice
           lowerCamelCase__ 					=			int(ids_tensor((1,)		,					shape_list(snake_case_     )[-1]     )     )
           lowerCamelCase__ 					=			output_from_no_past[:, -1, random_slice_idx]
           lowerCamelCase__ 					=			output_from_past[:, 0, random_slice_idx]

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




     def 	__snake_case				(      self					:				Any		,					A__					:				int		,					A__					:				str		,					A__					:				str		,					A__					:				int		,					A__					:				str		,					*A__					:				List[str]     )			->    str:
           lowerCamelCase__ 					=			TFGPTJModel(config=snake_case_     )

           # create attention mask
           lowerCamelCase__ 					=			self.seq_length // 2
           lowerCamelCase__ 					=			tf.ones((self.batch_size, half_seq_length)		,					dtype=tf.intaa     )
           lowerCamelCase__ 					=			tf.zeros((self.batch_size, self.seq_length - half_seq_length)		,					dtype=tf.intaa     )
           lowerCamelCase__ 					=			tf.concat([attn_mask_begin, attn_mask_end]		,					axis=1     )

           # first forward pass
           lowerCamelCase__ 					=			model(snake_case_		,					attention_mask=snake_case_     ).to_tuple()

           # create hypothetical next token and extent to next_input_ids
           lowerCamelCase__ 					=			ids_tensor((self.batch_size, 1)		,					config.vocab_size     )

           # change a random masked slice from input_ids
           lowerCamelCase__ 					=			ids_tensor((1,)		,					snake_case_     ).numpy() + 1
           lowerCamelCase__ 					=			ids_tensor((self.batch_size, self.seq_length)		,					config.vocab_size     )
           lowerCamelCase__ 					=			tf.range(self.seq_length     ) == (self.seq_length - random_seq_idx_to_change)
           lowerCamelCase__ 					=			tf.transpose(
               tf.broadcast_to(tf.expand_dims(snake_case_		,					-1     )		,					(self.seq_length, self.batch_size)     )     )
           lowerCamelCase__ 					=			tf.where(snake_case_		,					snake_case_		,					snake_case_     )

           # append to next input_ids and attn_mask
           lowerCamelCase__ 					=			tf.concat([input_ids, next_tokens]		,					axis=-1     )
           lowerCamelCase__ 					=			tf.concat([attn_mask, tf.ones((shape_list(snake_case_     )[0], 1)		,					dtype=tf.intaa     )]		,					axis=1     )

           # get two different outputs
           lowerCamelCase__ 					=			model(snake_case_		,					attention_mask=snake_case_     )["""last_hidden_state"""]
           lowerCamelCase__ 					=			model(snake_case_		,					past_key_values=snake_case_		,					attention_mask=snake_case_     )[
               """last_hidden_state"""
           ]

           # select random slice
           lowerCamelCase__ 					=			int(ids_tensor((1,)		,					shape_list(snake_case_     )[-1]     )     )
           lowerCamelCase__ 					=			output_from_no_past[:, -1, random_slice_idx]
           lowerCamelCase__ 					=			output_from_past[:, 0, random_slice_idx]

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




     def 	__snake_case				(      self					:				Dict		,					A__					:				Optional[Any]		,					A__					:				Optional[int]		,					A__					:				Optional[int]		,					A__					:				Any		,					A__					:				int		,					*A__					:				Tuple     )			->    Any:
           lowerCamelCase__ 					=			TFGPTJModel(config=snake_case_     )

           lowerCamelCase__ 					=			input_ids[:1, :]
           lowerCamelCase__ 					=			input_mask[:1, :]
           lowerCamelCase__ 					=			token_type_ids[:1, :]
           lowerCamelCase__ 					=			1

           # first forward pass
           lowerCamelCase__ 					=			model(snake_case_		,					attention_mask=snake_case_		,					token_type_ids=snake_case_		,					use_cache=snake_case_     )

           lowerCamelCase__ 					=			outputs.to_tuple()

           # create hypothetical next token and extent to next_input_ids
           lowerCamelCase__ 					=			ids_tensor((self.batch_size, 3)		,					config.vocab_size     )
           lowerCamelCase__ 					=			ids_tensor((self.batch_size, 3)		,					2     )
           lowerCamelCase__ 					=			ids_tensor((self.batch_size, 3)		,					self.type_vocab_size     )

           # append to next input_ids and token_type_ids
           lowerCamelCase__ 					=			tf.concat([input_ids, next_tokens]		,					axis=-1     )
           lowerCamelCase__ 					=			tf.concat([input_mask, next_attn_mask]		,					axis=-1     )
           lowerCamelCase__ 					=			tf.concat([token_type_ids, next_token_types]		,					axis=-1     )

           lowerCamelCase__ 					=			model(
               snake_case_		,					token_type_ids=snake_case_		,					attention_mask=snake_case_     )["""last_hidden_state"""]
           lowerCamelCase__ 					=			model(
               snake_case_		,					token_type_ids=snake_case_		,					attention_mask=snake_case_		,					past_key_values=snake_case_		,					)["""last_hidden_state"""]
           self.parent.assertTrue(output_from_past.shape[1] == next_tokens.shape[1]     )

           # select random slice
           lowerCamelCase__ 					=			int(ids_tensor((1,)		,					shape_list(snake_case_     )[-1]     )     )
           lowerCamelCase__ 					=			output_from_no_past[:, -3:, random_slice_idx]
           lowerCamelCase__ 					=			output_from_past[:, :, random_slice_idx]

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




     def 	__snake_case				(      self					:				Dict		,					A__					:				Dict		,					A__					:				Any		,					A__					:				Tuple		,					A__					:				Optional[Any]		,					A__					:				List[Any]		,					*A__					:				Optional[Any]     )			->    Any:
           lowerCamelCase__ 					=			TFGPTJForCausalLM(config=snake_case_     )
           lowerCamelCase__ 					=			{
               """input_ids""": input_ids,
               """attention_mask""": input_mask,
               """token_type_ids""": token_type_ids,
           }
           lowerCamelCase__ 					=			model(snake_case_     )
           self.parent.assertEqual(result.logits.shape		,					(self.batch_size, self.seq_length, self.vocab_size)     )




     def 	__snake_case				(      self					:				Optional[int]     )			->    Dict:
           lowerCamelCase__ 					=			self.prepare_config_and_inputs()

           (
               lowerCamelCase__   
           ) 					=			config_and_inputs

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






@require_tf
class      A     (						_a     ,  _a     ,  _a     ,  unittest.TestCase						):
     """simple docstring"""

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

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




     def 	__snake_case				(      self					:				Optional[int]		,					A__					:				Optional[int]		,					A__					:				Union[str, Any]		,					A__					:				Any		,					A__					:				Dict		,					A__					:				List[Any]     )			->    Any:
           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 	__snake_case				(      self					:				Tuple     )			->    Optional[Any]:
           lowerCamelCase__ 					=			TFGPTJModelTester(self     )
           lowerCamelCase__ 					=			ConfigTester(self		,					config_class=snake_case_		,					n_embd=3_7     )




     def 	__snake_case				(      self					:				List[str]     )			->    str:
           self.config_tester.run_common_tests()




     def 	__snake_case				(      self					:				Tuple     )			->    List[Any]:
           lowerCamelCase__ 					=			self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_gptj_model(*snake_case_     )




     def 	__snake_case				(      self					:				Optional[Any]     )			->    Tuple:
           lowerCamelCase__ 					=			self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_gptj_model_past(*snake_case_     )




     def 	__snake_case				(      self					:				List[Any]     )			->    str:
           lowerCamelCase__ 					=			self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_gptj_model_attention_mask_past(*snake_case_     )




     def 	__snake_case				(      self					:				List[Any]     )			->    Tuple:
           lowerCamelCase__ 					=			self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_gptj_model_past_large_inputs(*snake_case_     )




     def 	__snake_case				(      self					:				Optional[int]     )			->    Optional[int]:
           lowerCamelCase__ 					=			self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_gptj_lm_head_model(*snake_case_     )




     @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 	__snake_case				(      self					:				Dict     )			->    Union[str, Any]:
           lowerCamelCase__ 					=			TFGPTJModel.from_pretrained('EleutherAI/gpt-j-6B'		,					from_pt=snake_case_     )
           self.assertIsNotNone(snake_case_     )




     @unittest.skip(reason='Currently, model embeddings are going to undergo a major refactor.'     )
     def 	__snake_case				(      self					:				Tuple     )			->    Optional[int]:
           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      A     (						unittest.TestCase						):
     """simple docstring"""

     def 	__snake_case				(      self					:				int     )			->    str:
           lowerCamelCase__ 					=			TFGPTJForCausalLM.from_pretrained('EleutherAI/gpt-j-6B'		,					from_pt=snake_case_     )
           lowerCamelCase__ 					=			tf.convert_to_tensor([[4_6_4, 3_2_9_0]]		,					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
           lowerCamelCase__ 					=			[4_6_4, 3_2_9_0, 3_1_8, 2_5_7, 5_8_2, 3_3_8, 1_2_6_6, 1_5_4_5, 1_3, 6_3_2, 3_1_8, 2_5_7, 9_1_1_2, 1_5_1_8_5, 1_1, 2_9_0, 3_4_0, 3_1_8, 2_5_7, 1_5_4_5]
           # fmt: on
           lowerCamelCase__ 					=			model.generate(snake_case_		,					do_sample=snake_case_     )
           self.assertListEqual(output_ids[0].numpy().tolist()		,					snake_case_     )




     def 	__snake_case				(      self					:				Union[str, Any]     )			->    Optional[Any]:
           lowerCamelCase__ 					=			AutoTokenizer.from_pretrained('EleutherAI/gpt-j-6B'		,					revision='float16'     )
           lowerCamelCase__ 					=			TFGPTJForCausalLM.from_pretrained('EleutherAI/gpt-j-6B'		,					revision='float16'		,					from_pt=snake_case_     )

           lowerCamelCase__ 					=			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'     ):
                 lowerCamelCase__ 					=			model.generate(**snake_case_		,					do_sample=snake_case_		,					seed=[4_2, 0]     )
           lowerCamelCase__ 					=			tokenizer.decode(output_ids[0]		,					skip_special_tokens=snake_case_     )

           lowerCamelCase__ 					=			"""Today is a nice day and I’m going to go for a walk. I’"""
           self.assertEqual(snake_case_		,					snake_case_     )




     def 	__snake_case				(      self					:				Dict     )			->    Any:
           lowerCamelCase__ 					=			TFGPTJForCausalLM.from_pretrained('EleutherAI/gpt-j-6B'		,					revision='float16'		,					from_pt=snake_case_     )
           lowerCamelCase__ 					=			AutoTokenizer.from_pretrained('EleutherAI/gpt-j-6B'		,					revision='float16'     )

           lowerCamelCase__ 					=			"""left"""

           # Define PAD Token = EOS Token = 50256
           lowerCamelCase__ 					=			tokenizer.eos_token
           lowerCamelCase__ 					=			model.config.eos_token_id

           # use different length sentences to test batching
           lowerCamelCase__ 					=			[
               """Hello, my dog is a little""",
               """Today, I""",
           ]
           lowerCamelCase__ 					=			[
               """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 	__snake_case				(      self					:				List[Any]     )			->    List[Any]:
           # Confirms that we get the expected results with left-padded beam search
           lowerCamelCase__ 					=			self._get_beam_search_test_objects()

           lowerCamelCase__ 					=			tokenizer(snake_case_		,					return_tensors='tf'		,					padding=snake_case_     )
           lowerCamelCase__ 					=			model.generate(**snake_case_		,					do_sample=snake_case_		,					num_beams=2     )
           lowerCamelCase__ 					=			tokenizer.batch_decode(snake_case_		,					skip_special_tokens=snake_case_     )
           self.assertListEqual(snake_case_		,					snake_case_     )




     def 	__snake_case				(      self					:				Optional[int]     )			->    Any:
           # Confirms that left-padding is working properly
           lowerCamelCase__ 					=			self._get_beam_search_test_objects()

           lowerCamelCase__ 					=			tokenizer(snake_case_		,					return_tensors='tf'		,					padding=snake_case_     )
           lowerCamelCase__ 					=			tokenizer(sentences[0]		,					return_tensors='tf'     )
           lowerCamelCase__ 					=			model.generate(**snake_case_		,					do_sample=snake_case_		,					num_beams=2     )
           lowerCamelCase__ 					=			(
               shape_list(inputs_non_padded['input_ids']     )[-1]
               - tf.reduce_sum(tf.cast(inputs['attention_mask'][-1]		,					tf.intaa     )     ).numpy()
           )
           lowerCamelCase__ 					=			tokenizer(sentences[1]		,					return_tensors='tf'     )
           lowerCamelCase__ 					=			model.generate(
               **snake_case_		,					do_sample=snake_case_		,					num_beams=2		,					max_length=model.config.max_length - num_paddings     )
           lowerCamelCase__ 					=			tokenizer.decode(output_non_padded[0]		,					skip_special_tokens=snake_case_     )
           lowerCamelCase__ 					=			tokenizer.decode(output_padded[0]		,					skip_special_tokens=snake_case_     )
           self.assertListEqual(snake_case_		,					[non_padded_sentence, padded_sentence]     )




     def 	__snake_case				(      self					:				Any     )			->    List[Any]:
           # Confirms that XLA is working properly
           lowerCamelCase__ 					=			self._get_beam_search_test_objects()

           lowerCamelCase__ 					=			tokenizer(snake_case_		,					return_tensors='tf'		,					padding=snake_case_     )
           lowerCamelCase__ 					=			tf.function(model.generate		,					jit_compile=snake_case_     )
           lowerCamelCase__ 					=			xla_generate(**snake_case_		,					do_sample=snake_case_		,					num_beams=2     )
           lowerCamelCase__ 					=			tokenizer.batch_decode(snake_case_		,					skip_special_tokens=snake_case_     )
           self.assertListEqual(snake_case_		,					snake_case_     )






from __future__ import annotations

def      lowerCAmelCase_     (					lowerCamelCase:    list[int | float]  )   ->     int | float:


      '''simple docstring'''




      if len(lowerCamelCase  ) == 0:
            raise ValueError('find_max() arg is an empty sequence'  )
      lowerCamelCase__ 					=			nums[0]
      for x in nums:
            if x > max_num:
                  lowerCamelCase__ 					=			x
      return max_num


if __name__ == "__main__":
      import doctest

      doctest.testmod(verbose=True)