

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							:







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







      def       lowerCAmelCase_       (	self  )		->			Any:
        snake_case_					= ids_tensor([self.batch_size, self.seq_length]				,    self.vocab_size  )

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

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

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

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

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

        snake_case_					= 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       lowerCAmelCase_       (	self				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase				,    *lowerCamelCase  )		->			Union[str, Any]:
        snake_case_					= TFGPTJModel(config=lowerCamelCase  )
        snake_case_					= {
            """input_ids""": input_ids,
            """attention_mask""": input_mask,
            """token_type_ids""": token_type_ids,
        }
        snake_case_					= model(lowerCamelCase  )

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

        snake_case_					= model(lowerCamelCase  )

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







      def       lowerCAmelCase_       (	self				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase				,    *lowerCamelCase  )		->			Optional[Any]:
        snake_case_					= TFGPTJModel(config=lowerCamelCase  )

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

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

        snake_case_		,							snake_case_					= outputs.to_tuple()

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

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

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

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

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







      def       lowerCAmelCase_       (	self				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase				,    *lowerCamelCase  )		->			Union[str, Any]:
        snake_case_					= TFGPTJModel(config=lowerCamelCase  )

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

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

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

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

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

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

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

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







      def       lowerCAmelCase_       (	self				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase				,    *lowerCamelCase  )		->			Optional[Any]:
        snake_case_					= TFGPTJModel(config=lowerCamelCase  )

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

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

        snake_case_		,							snake_case_					= outputs.to_tuple()

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

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

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

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

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







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







      def       lowerCAmelCase_       (	self  )		->			str:
        snake_case_					= self.prepare_config_and_inputs()

        (
            (
            snake_case_
        )		,							(
            snake_case_
        )		,							(
            snake_case_
        )		,							(
            snake_case_
        )		,							(
            snake_case_
        )		,							(
            snake_case_
        )		,							(
            snake_case_
        )		,							(
            snake_case_
        )		,							(
            snake_case_
        )		,							
        )					= config_and_inputs

        snake_case_					= {
            """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      ):
      lowerCamelCase_      : str  			=   (
          (TFGPTJForCausalLM, TFGPTJForSequenceClassification, TFGPTJForQuestionAnswering, TFGPTJModel)
          if is_tf_available()
          else ()
      )

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







      def       lowerCAmelCase_       (	self				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase				,    lowerCamelCase  )		->			List[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       lowerCAmelCase_       (	self  )		->			Union[str, Any]:
        snake_case_					= TFGPTJModelTester(self  )
        snake_case_					= ConfigTester(self				,    config_class=lowerCamelCase				,    n_embd=37  )







      def       lowerCAmelCase_       (	self  )		->			Optional[int]:
        self.config_tester.run_common_tests()







      def       lowerCAmelCase_       (	self  )		->			List[str]:
        snake_case_					= self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_gptj_model(*lowerCamelCase  )







      def       lowerCAmelCase_       (	self  )		->			Any:
        snake_case_					= self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_gptj_model_past(*lowerCamelCase  )







      def       lowerCAmelCase_       (	self  )		->			int:
        snake_case_					= self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_gptj_model_attention_mask_past(*lowerCamelCase  )







      def       lowerCAmelCase_       (	self  )		->			Optional[int]:
        snake_case_					= self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_gptj_model_past_large_inputs(*lowerCamelCase  )







      def       lowerCAmelCase_       (	self  )		->			Optional[Any]:
        snake_case_					= self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_gptj_lm_head_model(*lowerCamelCase  )







      @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       lowerCAmelCase_       (	self  )		->			Optional[Any]:
        snake_case_					= TFGPTJModel.from_pretrained("""EleutherAI/gpt-j-6B"""				,    from_pt=lowerCamelCase  )
        self.assertIsNotNone(lowerCamelCase  )







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







      def       lowerCAmelCase_       (	self  )		->			Optional[int]:
        snake_case_					= TFGPTJForCausalLM.from_pretrained("""EleutherAI/gpt-j-6B"""				,    from_pt=lowerCamelCase  )
        snake_case_					= 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
        snake_case_					= [464, 3290, 318, 257, 582, 338, 1266, 1545, 13, 632, 318, 257, 9112, 15185, 11, 290, 340, 318, 257, 1545]
        # fmt: on
        snake_case_					= model.generate(lowerCamelCase				,    do_sample=lowerCamelCase  )
        self.assertListEqual(output_ids[0].numpy().tolist()				,    lowerCamelCase  )







      def       lowerCAmelCase_       (	self  )		->			List[Any]:
        snake_case_					= AutoTokenizer.from_pretrained("""EleutherAI/gpt-j-6B"""				,    revision="""float16"""  )
        snake_case_					= TFGPTJForCausalLM.from_pretrained("""EleutherAI/gpt-j-6B"""				,    revision="""float16"""				,    from_pt=lowerCamelCase  )

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

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







      def       lowerCAmelCase_       (	self  )		->			Optional[int]:
        snake_case_					= TFGPTJForCausalLM.from_pretrained("""EleutherAI/gpt-j-6B"""				,    revision="""float16"""				,    from_pt=lowerCamelCase  )
        snake_case_					= AutoTokenizer.from_pretrained("""EleutherAI/gpt-j-6B"""				,    revision="""float16"""  )

        snake_case_					= """left"""

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

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

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







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

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







      def       lowerCAmelCase_       (	self  )		->			Optional[int]:
        # Confirms that XLA is working properly
        snake_case_		,							snake_case_		,							snake_case_		,							snake_case_					= self._get_beam_search_test_objects()

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


def   UpperCamelCase(   lowercase_				)      ->							float:





  '''simple docstring'''





  return -num if num < 0 else num



def   UpperCamelCase(   lowercase_				)      ->							int:





  '''simple docstring'''





  if len(lowercase_				) == 0:
    raise ValueError("""abs_min() arg is an empty sequence"""				)
  snake_case_					= x[0]
  for i in x:
    if abs_val(lowercase_				) < abs_val(lowercase_				):
      snake_case_					= i
  return j



def   UpperCamelCase(   lowercase_				)      ->							int:





  '''simple docstring'''





  if len(lowercase_				) == 0:
    raise ValueError("""abs_max() arg is an empty sequence"""				)
  snake_case_					= x[0]
  for i in x:
    if abs(lowercase_				) > abs(lowercase_				):
      snake_case_					= i
  return j



def   UpperCamelCase(   lowercase_				)      ->							int:





  '''simple docstring'''





  if len(lowercase_				) == 0:
    raise ValueError("""abs_max_sort() arg is an empty sequence"""				)
  return sorted(lowercase_     ,   key=lowercase_				)[-1]



def   UpperCamelCase(   )      ->							List[Any]:





  '''simple docstring'''





  assert abs_val(0				) == 0
  assert abs_val(34				) == 34
  assert abs_val(-100000000000				) == 100000000000

  snake_case_					= [-3, -1, 2, -11]
  assert abs_max(lowercase_				) == -11
  assert abs_max_sort(lowercase_				) == -11
  assert abs_min(lowercase_				) == -1


if __name__ == "__main__":
    import doctest

    doctest.testmod()

    test_abs_val()
    print(abs_val(-34))  # --> 34