





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  ,    UpperCAmelCase     )-> List[Any]:
     __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				 =       5_12
     __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 		SCREAMING_SNAKE_CASE__	( self     )-> Tuple:
     __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 		SCREAMING_SNAKE_CASE__	( self  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    *UpperCAmelCase     )-> Dict:
     __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 		SCREAMING_SNAKE_CASE__	( self  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    *UpperCAmelCase     )-> Dict:
     __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				 =       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 		SCREAMING_SNAKE_CASE__	( self  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    *UpperCAmelCase     )-> str:
     __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				 =       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 		SCREAMING_SNAKE_CASE__	( self  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    *UpperCAmelCase     )-> str:
     __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				 =       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 		SCREAMING_SNAKE_CASE__	( self  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    *UpperCAmelCase     )-> List[str]:
     __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 		SCREAMING_SNAKE_CASE__	( self     )-> Optional[Any]:
     __A				 =       self.prepare_config_and_inputs()

     (
         __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 						_lowerCAmelCase(						lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_		,   unittest.TestCase):

  """simple docstring"""







  lowerCamelCase__      					=   (
      (TFGPTJForCausalLM, TFGPTJForSequenceClassification, TFGPTJForQuestionAnswering, TFGPTJModel)
      if is_tf_available()
      else ()
  )

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

  def 		SCREAMING_SNAKE_CASE__	( self  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase     )-> Tuple:







     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 		SCREAMING_SNAKE_CASE__	( self     )-> List[str]:
     __A				 =       TFGPTJModelTester(self     )
     __A				 =       ConfigTester(self  ,    config_class=UpperCAmelCase  ,    n_embd=37     )

  def 		SCREAMING_SNAKE_CASE__	( self     )-> str:
     self.config_tester.run_common_tests()

  def 		SCREAMING_SNAKE_CASE__	( self     )-> Dict:
     __A				 =       self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_gptj_model(*UpperCAmelCase     )

  def 		SCREAMING_SNAKE_CASE__	( self     )-> str:
     __A				 =       self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_gptj_model_past(*UpperCAmelCase     )

  def 		SCREAMING_SNAKE_CASE__	( self     )-> List[str]:
     __A				 =       self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_gptj_model_attention_mask_past(*UpperCAmelCase     )

  def 		SCREAMING_SNAKE_CASE__	( self     )-> Tuple:
     __A				 =       self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_gptj_model_past_large_inputs(*UpperCAmelCase     )

  def 		SCREAMING_SNAKE_CASE__	( self     )-> List[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 		SCREAMING_SNAKE_CASE__	( self     )-> Optional[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 		SCREAMING_SNAKE_CASE__	( self     )-> str:
     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 		SCREAMING_SNAKE_CASE__	( self     )-> Tuple:
     __A				 =       TFGPTJForCausalLM.from_pretrained('''EleutherAI/gpt-j-6B'''  ,    from_pt=UpperCAmelCase     )
     __A				 =       tf.convert_to_tensor([[4_64, 32_90]]  ,    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				 =       [4_64, 32_90, 3_18, 2_57, 5_82, 3_38, 12_66, 15_45, 13, 6_32, 3_18, 2_57, 91_12, 1_51_85, 11, 2_90, 3_40, 3_18, 2_57, 15_45]
     # fmt: on
     __A				 =       model.generate(UpperCAmelCase  ,    do_sample=UpperCAmelCase     )
     self.assertListEqual(output_ids[0].numpy().tolist()  ,    UpperCAmelCase     )

  def 		SCREAMING_SNAKE_CASE__	( self     )-> Dict:
     __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 		SCREAMING_SNAKE_CASE__	( self     )-> List[Any]:
     __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 		SCREAMING_SNAKE_CASE__	( self     )-> str:
     __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 		SCREAMING_SNAKE_CASE__	( self     )-> Dict:
     __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 		SCREAMING_SNAKE_CASE__	( self     )-> Any:
     __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     )








from __future__ import annotations
def 				__UpperCamelCase  (		snake_case    ,	snake_case     )  ->						list:


   '''simple docstring'''







   __A				 =       [None] * (len(snake_case     ) + len(snake_case     ))

   __A				 =       0  # pointer to current index for left Half
   __A				 =       0  # pointer to current index for the right Half
   __A				 =       0  # pointer to current index for the sorted array Half

   while pointera < len(snake_case     ) and pointera < len(snake_case     ):
      if left_half[pointera] < right_half[pointera]:
         __A				 =       left_half[pointera]
         pointera += 1
         index += 1
      else:
         __A				 =       right_half[pointera]
         pointera += 1
         index += 1
   while pointera < len(snake_case     ):
      __A				 =       left_half[pointera]
      pointera += 1
      index += 1

   while pointera < len(snake_case     ):
      __A				 =       right_half[pointera]
      pointera += 1
      index += 1

   return sorted_array
def 				__UpperCamelCase  (		snake_case     )  ->						list:


   '''simple docstring'''







   if len(snake_case     ) <= 1:
      return array
   # the actual formula to calculate the middle element = left + (right - left) // 2
   # this avoids integer overflow in case of large N
   __A				 =       0 + (len(snake_case     ) - 0) // 2

   # Split the array into halves till the array length becomes equal to One
   # merge the arrays of single length returned by mergeSort function and
   # pass them into the merge arrays function which merges the array
   __A				 =       array[:middle]
   __A				 =       array[middle:]

   return merge(merge_sort(snake_case     )    ,	merge_sort(snake_case     )     )


if __name__ == "__main__":
    import doctest

    doctest.testmod()

