


'''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     __SCREAMING_SNAKE_CASE       :






     def __init__(      self  ,				__UpperCamelCase		)							->     Optional[int]:
          _a            =		parent
          _a            =		13
          _a            =		7
          _a            =		True
          _a            =		True
          _a            =		True
          _a            =		True
          _a            =		True
          _a            =		99
          _a            =		32
          _a            =		4
          _a            =		2
          _a            =		4
          _a            =		37
          _a            =		"gelu"
          _a            =		0.1
          _a            =		0.1
          _a            =		512
          _a            =		16
          _a            =		2
          _a            =		0.02
          _a            =		3
          _a            =		4
          _a            =		None
          _a            =		self.vocab_size - 1
          _a            =		self.vocab_size - 1
          _a            =		self.vocab_size - 1






     def 			a_       (      self		)							->     Optional[int]:
          _a            =		ids_tensor([self.batch_size, self.seq_length]  ,				self.vocab_size		)

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

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

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

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

          _a            =		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=__UpperCamelCase  ,				)

          _a            =		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  ,				__UpperCamelCase  ,				__UpperCamelCase  ,				__UpperCamelCase  ,				__UpperCamelCase  ,				__UpperCamelCase  ,				*__UpperCamelCase		)							->     Optional[Any]:
          _a            =		TFGPTJModel(config=__UpperCamelCase		)
          _a            =		{
              "input_ids": input_ids,
              "attention_mask": input_mask,
              "token_type_ids": token_type_ids,
          }
          _a            =		model(__UpperCamelCase		)

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

          _a            =		model(__UpperCamelCase		)

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






     def 			a_       (      self  ,				__UpperCamelCase  ,				__UpperCamelCase  ,				__UpperCamelCase  ,				__UpperCamelCase  ,				__UpperCamelCase  ,				*__UpperCamelCase		)							->     List[str]:
          _a            =		TFGPTJModel(config=__UpperCamelCase		)

          # first forward pass
          _a            =		model(__UpperCamelCase  ,				token_type_ids=__UpperCamelCase  ,				use_cache=__UpperCamelCase		)
          _a            =		model(__UpperCamelCase  ,				token_type_ids=__UpperCamelCase		)
          _a            =		model(__UpperCamelCase  ,				token_type_ids=__UpperCamelCase  ,				use_cache=__UpperCamelCase		)

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

          _a      ,      _a            =		outputs.to_tuple()

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

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

          _a            =		model(__UpperCamelCase  ,				token_type_ids=__UpperCamelCase		)["last_hidden_state"]
          _a            =		model(__UpperCamelCase  ,				token_type_ids=__UpperCamelCase  ,				past_key_values=__UpperCamelCase		)[
              "last_hidden_state"
          ]

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

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






     def 			a_       (      self  ,				__UpperCamelCase  ,				__UpperCamelCase  ,				__UpperCamelCase  ,				__UpperCamelCase  ,				__UpperCamelCase  ,				*__UpperCamelCase		)							->     Tuple:
          _a            =		TFGPTJModel(config=__UpperCamelCase		)

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

          # first forward pass
          _a      ,      _a            =		model(__UpperCamelCase  ,				attention_mask=__UpperCamelCase		).to_tuple()

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

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

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

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

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

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






     def 			a_       (      self  ,				__UpperCamelCase  ,				__UpperCamelCase  ,				__UpperCamelCase  ,				__UpperCamelCase  ,				__UpperCamelCase  ,				*__UpperCamelCase		)							->     Any:
          _a            =		TFGPTJModel(config=__UpperCamelCase		)

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

          # first forward pass
          _a            =		model(__UpperCamelCase  ,				attention_mask=__UpperCamelCase  ,				token_type_ids=__UpperCamelCase  ,				use_cache=__UpperCamelCase		)

          _a      ,      _a            =		outputs.to_tuple()

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

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

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

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

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






     def 			a_       (      self  ,				__UpperCamelCase  ,				__UpperCamelCase  ,				__UpperCamelCase  ,				__UpperCamelCase  ,				__UpperCamelCase  ,				*__UpperCamelCase		)							->     Optional[int]:
          _a            =		TFGPTJForCausalLM(config=__UpperCamelCase		)
          _a            =		{
              "input_ids": input_ids,
              "attention_mask": input_mask,
              "token_type_ids": token_type_ids,
          }
          _a            =		model(__UpperCamelCase		)
          self.parent.assertEqual(result.logits.shape  ,				(self.batch_size, self.seq_length, self.vocab_size)		)






     def 			a_       (      self		)							->     Dict:
          _a            =		self.prepare_config_and_inputs()

          (
              (
              _a
          )      ,      (
              _a
          )      ,      (
              _a
          )      ,      (
              _a
          )      ,      (
              _a
          )      ,      (
              _a
          )      ,      (
              _a
          )      ,      (
              _a
          )      ,      (
              _a
          )      ,      
          )            =		config_and_inputs

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







@require_tf
class     __SCREAMING_SNAKE_CASE       (   lowerCamelCase__			,      lowerCamelCase__			,      lowerCamelCase__			,      unittest.TestCase			):
     UpperCAmelCase			       =					(
         (TFGPTJForCausalLM, TFGPTJForSequenceClassification, TFGPTJForQuestionAnswering, TFGPTJModel)
         if is_tf_available()
         else ()
     )

     UpperCAmelCase			       =					(TFGPTJForCausalLM,) if is_tf_available() else ()
     UpperCAmelCase			       =					(
         {
             '''feature-extraction''': TFGPTJModel,
             '''question-answering''': TFGPTJForQuestionAnswering,
             '''text-classification''': TFGPTJForSequenceClassification,
             '''text-generation''': TFGPTJForCausalLM,
             '''zero-shot''': TFGPTJForSequenceClassification,
         }
         if is_tf_available()
         else {}
     )
     UpperCAmelCase			       =					False
     UpperCAmelCase			       =					False
     UpperCAmelCase			       =					False
     UpperCAmelCase			       =					False






     def 			a_       (      self  ,				__UpperCamelCase  ,				__UpperCamelCase  ,				__UpperCamelCase  ,				__UpperCamelCase  ,				__UpperCamelCase		)							->     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 			a_       (      self		)							->     Dict:
          _a            =		TFGPTJModelTester(self		)
          _a            =		ConfigTester(self  ,				config_class=__UpperCamelCase  ,				n_embd=37		)






     def 			a_       (      self		)							->     Any:
          self.config_tester.run_common_tests()






     def 			a_       (      self		)							->     List[Any]:
          _a            =		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_gptj_model(*__UpperCamelCase		)






     def 			a_       (      self		)							->     List[str]:
          _a            =		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_gptj_model_past(*__UpperCamelCase		)






     def 			a_       (      self		)							->     int:
          _a            =		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_gptj_model_attention_mask_past(*__UpperCamelCase		)






     def 			a_       (      self		)							->     Optional[Any]:
          _a            =		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_gptj_model_past_large_inputs(*__UpperCamelCase		)






     def 			a_       (      self		)							->     str:
          _a            =		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_gptj_lm_head_model(*__UpperCamelCase		)






     @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		)							->     Union[str, Any]:
          _a            =		TFGPTJModel.from_pretrained("EleutherAI/gpt-j-6B"  ,				from_pt=__UpperCamelCase		)
          self.assertIsNotNone(__UpperCamelCase		)






     @unittest.skip(reason="Currently, model embeddings are going to undergo a major refactor."		)
     def 			a_       (      self		)							->     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     __SCREAMING_SNAKE_CASE       (   unittest.TestCase			):






     def 			a_       (      self		)							->     Optional[Any]:
          _a            =		TFGPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B"  ,				from_pt=__UpperCamelCase		)
          _a            =		tf.convert_to_tensor([[464, 3_290]]  ,				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
          _a            =		[464, 3_290, 318, 257, 582, 338, 1_266, 1_545, 13, 632, 318, 257, 9_112, 15_185, 11, 290, 340, 318, 257, 1_545]
          # fmt: on
          _a            =		model.generate(__UpperCamelCase  ,				do_sample=__UpperCamelCase		)
          self.assertListEqual(output_ids[0].numpy().tolist()  ,				__UpperCamelCase		)






     def 			a_       (      self		)							->     Any:
          _a            =		AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B"  ,				revision="float16"		)
          _a            =		TFGPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B"  ,				revision="float16"  ,				from_pt=__UpperCamelCase		)

          _a            =		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"		):
               _a            =		model.generate(**__UpperCamelCase  ,				do_sample=__UpperCamelCase  ,				seed=[42, 0]		)
          _a            =		tokenizer.decode(output_ids[0]  ,				skip_special_tokens=__UpperCamelCase		)

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






     def 			a_       (      self		)							->     str:
          _a            =		TFGPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B"  ,				revision="float16"  ,				from_pt=__UpperCamelCase		)
          _a            =		AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B"  ,				revision="float16"		)

          _a            =		"left"

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

          # use different length sentences to test batching
          _a            =		[
              "Hello, my dog is a little",
              "Today, I",
          ]
          _a            =		[
              "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
          _a      ,      _a      ,      _a      ,      _a            =		self._get_beam_search_test_objects()

          _a            =		tokenizer(__UpperCamelCase  ,				return_tensors="tf"  ,				padding=__UpperCamelCase		)
          _a            =		model.generate(**__UpperCamelCase  ,				do_sample=__UpperCamelCase  ,				num_beams=2		)
          _a            =		tokenizer.batch_decode(__UpperCamelCase  ,				skip_special_tokens=__UpperCamelCase		)
          self.assertListEqual(__UpperCamelCase  ,				__UpperCamelCase		)






     def 			a_       (      self		)							->     List[str]:
          # Confirms that left-padding is working properly
          _a      ,      _a      ,      _a      ,      _a            =		self._get_beam_search_test_objects()

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






     def 			a_       (      self		)							->     Tuple:
          # Confirms that XLA is working properly
          _a      ,      _a      ,      _a      ,      _a            =		self._get_beam_search_test_objects()

          _a            =		tokenizer(__UpperCamelCase  ,				return_tensors="tf"  ,				padding=__UpperCamelCase		)
          _a            =		tf.function(model.generate  ,				jit_compile=__UpperCamelCase		)
          _a            =		xla_generate(**__UpperCamelCase  ,				do_sample=__UpperCamelCase  ,				num_beams=2		)
          _a            =		tokenizer.batch_decode(__UpperCamelCase  ,				skip_special_tokens=__UpperCamelCase		)
          self.assertListEqual(__UpperCamelCase  ,				__UpperCamelCase		)





'''simple docstring'''

def __UpperCamelCase (      __lowerCamelCase		:			int					)   ->  int:







     '''simple docstring'''






     if not isinstance(__lowerCamelCase   ,	__lowerCamelCase					) or number < 0:
          raise ValueError("Input must be a non-negative integer"					)

     _a            =		number & ~(number - 1)
     _a            =		0
     while intermediate:
          intermediate >>= 1
          index += 1
     return index - 1


if __name__ == "__main__":
   import doctest

   doctest.testmod(verbose=True)

