




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 	UpperCamelCase__			:

    '''simple docstring'''



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







    def 					_lowercase  (					self  )  ->   int:
           lowerCamelCase     :   Tuple    				=      ids_tensor([self.batch_size, self.seq_length]  ,						self.vocab_size  )

           lowerCamelCase     :   Any    				=      None
           if self.use_input_mask:
                  lowerCamelCase     :   Union[str, Any]    				=      random_attention_mask([self.batch_size, self.seq_length]  )

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

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

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

           lowerCamelCase     :   List[str]    				=      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__  ,						)

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

           lowerCamelCase     :   Tuple    				=      [input_ids, None, input_mask]  # None is the input for 'past'
           lowerCamelCase     :   int    				=      model(UpperCamelCase__  )

           lowerCamelCase     :   List[str]    				=      model(UpperCamelCase__  )

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







    def 					_lowercase  (					self  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						*UpperCamelCase__  )  ->   int:
           lowerCamelCase     :   List[str]    				=      TFGPTJModel(config=UpperCamelCase__  )

           # first forward pass
           lowerCamelCase     :   List[str]    				=      model(UpperCamelCase__  ,						token_type_ids=UpperCamelCase__  ,						use_cache=UpperCamelCase__  )
           lowerCamelCase     :   Optional[int]    				=      model(UpperCamelCase__  ,						token_type_ids=UpperCamelCase__  )
           lowerCamelCase     :   Optional[Any]    				=      model(UpperCamelCase__  ,						token_type_ids=UpperCamelCase__  ,						use_cache=UpperCamelCase__  )

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

           lowerCamelCase				,   lowerCamelCase     :   str    				=      outputs.to_tuple()

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

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

           lowerCamelCase     :   str    				=      model(UpperCamelCase__  ,						token_type_ids=UpperCamelCase__  )["last_hidden_state"]
           lowerCamelCase     :   Optional[int]    				=      model(UpperCamelCase__  ,						token_type_ids=UpperCamelCase__  ,						past_key_values=UpperCamelCase__  )[
               "last_hidden_state"
           ]

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

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







    def 					_lowercase  (					self  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						*UpperCamelCase__  )  ->   Optional[int]:
           lowerCamelCase     :   Union[str, Any]    				=      TFGPTJModel(config=UpperCamelCase__  )

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

           # first forward pass
           lowerCamelCase				,   lowerCamelCase     :   Union[str, Any]    				=      model(UpperCamelCase__  ,						attention_mask=UpperCamelCase__  ).to_tuple()

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

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

           # append to next input_ids and attn_mask
           lowerCamelCase     :   List[str]    				=      tf.concat([input_ids, next_tokens]  ,						axis=-1  )
           lowerCamelCase     :   Union[str, Any]    				=      tf.concat([attn_mask, tf.ones((shape_list(UpperCamelCase__  )[0], 1)  ,						dtype=tf.intaa  )]  ,						axis=1  )

           # get two different outputs
           lowerCamelCase     :   List[str]    				=      model(UpperCamelCase__  ,						attention_mask=UpperCamelCase__  )["last_hidden_state"]
           lowerCamelCase     :   Any    				=      model(UpperCamelCase__  ,						past_key_values=UpperCamelCase__  ,						attention_mask=UpperCamelCase__  )[
               "last_hidden_state"
           ]

           # select random slice
           lowerCamelCase     :   Optional[Any]    				=      int(ids_tensor((1,)  ,						shape_list(UpperCamelCase__  )[-1]  )  )
           lowerCamelCase     :   Union[str, Any]    				=      output_from_no_past[:, -1, random_slice_idx]
           lowerCamelCase     :   int    				=      output_from_past[:, 0, random_slice_idx]

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







    def 					_lowercase  (					self  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						*UpperCamelCase__  )  ->   Optional[int]:
           lowerCamelCase     :   int    				=      TFGPTJModel(config=UpperCamelCase__  )

           lowerCamelCase     :   Union[str, Any]    				=      input_ids[:1, :]
           lowerCamelCase     :   List[str]    				=      input_mask[:1, :]
           lowerCamelCase     :   str    				=      token_type_ids[:1, :]
           lowerCamelCase     :   Any    				=      1

           # first forward pass
           lowerCamelCase     :   str    				=      model(UpperCamelCase__  ,						attention_mask=UpperCamelCase__  ,						token_type_ids=UpperCamelCase__  ,						use_cache=UpperCamelCase__  )

           lowerCamelCase				,   lowerCamelCase     :   str    				=      outputs.to_tuple()

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

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

           lowerCamelCase     :   Dict    				=      model(
               UpperCamelCase__  ,						token_type_ids=UpperCamelCase__  ,						attention_mask=UpperCamelCase__  )["last_hidden_state"]
           lowerCamelCase     :   List[str]    				=      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
           lowerCamelCase     :   Optional[Any]    				=      int(ids_tensor((1,)  ,						shape_list(UpperCamelCase__  )[-1]  )  )
           lowerCamelCase     :   Optional[Any]    				=      output_from_no_past[:, -3:, random_slice_idx]
           lowerCamelCase     :   int    				=      output_from_past[:, :, random_slice_idx]

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







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







    def 					_lowercase  (					self  )  ->   Tuple:
           lowerCamelCase     :   Any    				=      self.prepare_config_and_inputs()

           (
               (
               lowerCamelCase
           )				,   (
               lowerCamelCase
           )				,   (
               lowerCamelCase
           )				,   (
               lowerCamelCase
           )				,   (
               lowerCamelCase
           )				,   (
               lowerCamelCase
           )				,   (
               lowerCamelCase
           )				,   (
               lowerCamelCase
           )				,   (
               lowerCamelCase
           )				,   
           )     :   Tuple    				=      config_and_inputs

           lowerCamelCase     :   Optional[int]    				=      {
               "input_ids": input_ids,
               "token_type_ids": token_type_ids,
               "attention_mask": input_mask,
           }
           return config, inputs_dict




@require_tf
class 	UpperCamelCase__			(lowerCAmelCase__					,       lowerCAmelCase__					,       lowerCAmelCase__					,       unittest.TestCase  ):

    '''simple docstring'''



    lowerCamelCase_   :			Optional[int]			  =   (
        (TFGPTJForCausalLM, TFGPTJForSequenceClassification, TFGPTJForQuestionAnswering, TFGPTJModel)
        if is_tf_available()
        else ()
    )

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







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







    def 					_lowercase  (					self  )  ->   List[Any]:
           self.config_tester.run_common_tests()







    def 					_lowercase  (					self  )  ->   int:
           lowerCamelCase     :   Any    				=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_gptj_model(*UpperCamelCase__  )







    def 					_lowercase  (					self  )  ->   Any:
           lowerCamelCase     :   str    				=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_gptj_model_past(*UpperCamelCase__  )







    def 					_lowercase  (					self  )  ->   Union[str, Any]:
           lowerCamelCase     :   Tuple    				=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_gptj_model_attention_mask_past(*UpperCamelCase__  )







    def 					_lowercase  (					self  )  ->   int:
           lowerCamelCase     :   List[str]    				=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_gptj_model_past_large_inputs(*UpperCamelCase__  )







    def 					_lowercase  (					self  )  ->   int:
           lowerCamelCase     :   Dict    				=      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 					_lowercase  (					self  )  ->   Any:
           lowerCamelCase     :   str    				=      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 					_lowercase  (					self  )  ->   Tuple:
           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 	UpperCamelCase__			(unittest.TestCase  ):

    '''simple docstring'''



    def 					_lowercase  (					self  )  ->   Any:
           lowerCamelCase     :   Dict    				=      TFGPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B"  ,						from_pt=UpperCamelCase__  )
           lowerCamelCase     :   List[Any]    				=      tf.convert_to_tensor([[464, 3290]]  ,						dtype=tf.intaa  )  # The dog
           # fmt: off
           # The dog is a man's best friend. It is a loyal companion, and it is a friend
           lowerCamelCase     :   Optional[int]    				=      [464, 3290, 318, 257, 582, 338, 1266, 1545, 13, 632, 318, 257, 9112, 1_5185, 11, 290, 340, 318, 257, 1545]
           # fmt: on
           lowerCamelCase     :   Any    				=      model.generate(UpperCamelCase__  ,						do_sample=UpperCamelCase__  )
           self.assertListEqual(output_ids[0].numpy().tolist()  ,						UpperCamelCase__  )







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

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

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







    def 					_lowercase  (					self  )  ->   str:
           lowerCamelCase     :   List[str]    				=      TFGPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B"  ,						revision="float16"  ,						from_pt=UpperCamelCase__  )
           lowerCamelCase     :   int    				=      AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B"  ,						revision="float16"  )

           lowerCamelCase     :   Optional[Any]    				=      "left"

           # Define PAD Token = EOS Token = 50256
           lowerCamelCase     :   Optional[Any]    				=      tokenizer.eos_token
           lowerCamelCase     :   Union[str, Any]    				=      model.config.eos_token_id

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

           lowerCamelCase     :   Dict    				=      tokenizer(UpperCamelCase__  ,						return_tensors="tf"  ,						padding=UpperCamelCase__  )
           lowerCamelCase     :   str    				=      model.generate(**UpperCamelCase__  ,						do_sample=UpperCamelCase__  ,						num_beams=2  )
           lowerCamelCase     :   Optional[int]    				=      tokenizer.batch_decode(UpperCamelCase__  ,						skip_special_tokens=UpperCamelCase__  )
           self.assertListEqual(UpperCamelCase__  ,						UpperCamelCase__  )







    def 					_lowercase  (					self  )  ->   Dict:
           # Confirms that left-padding is working properly
           lowerCamelCase				,   lowerCamelCase				,   lowerCamelCase				,   lowerCamelCase     :   List[Any]    				=      self._get_beam_search_test_objects()

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







    def 					_lowercase  (					self  )  ->   Any:
           # Confirms that XLA is working properly
           lowerCamelCase				,   lowerCamelCase				,   lowerCamelCase				,   lowerCamelCase     :   Any    				=      self._get_beam_search_test_objects()

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












import os
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device


if is_torch_available():
       import torch

       from transformers import GPTaLMHeadModel




@require_torch
@require_sentencepiece
@require_tokenizers
class 	UpperCamelCase__			(unittest.TestCase  ):

    '''simple docstring'''



    @slow
    @unittest.skip("Model is not available."  )
    def 					_lowercase  (					self  )  ->   List[Any]:
           lowerCamelCase     :   List[Any]    				=      "nvidia/megatron-gpt2-345m/"
           if "MYDIR" in os.environ:
                  lowerCamelCase     :   Union[str, Any]    				=      os.path.join(os.environ["MYDIR"]  ,						UpperCamelCase__  )
           lowerCamelCase     :   Any    				=      GPTaLMHeadModel.from_pretrained(UpperCamelCase__  )
           model.to(UpperCamelCase__  )
           model.half()

           lowerCamelCase     :   Dict    				=      torch.tensor(
               [[101, 7110, 1005, 1056, 2023, 1_1333, 1_7413, 1029, 102]]  ,						device=UpperCamelCase__  ,						dtype=torch.long  ,						)

           with torch.no_grad():
                  lowerCamelCase     :   Any    				=      model(UpperCamelCase__  ).logits

           lowerCamelCase     :   Union[str, Any]    				=      torch.Size((1, 9, 5_0257)  )
           self.assertEqual(output.shape  ,						UpperCamelCase__  )

           lowerCamelCase     :   List[str]    				=      torch.tensor(
               [
                   4.9414,
                   -0.2920,
                   -1.2148,
                   -4.0273,
                   -0.5161,
                   -5.2109,
                   -1.2412,
                   -1.8301,
                   -1.7734,
                   -4.7148,
                   -0.2317,
                   -1.0811,
                   -2.1777,
                   0.4141,
                   -3.7969,
                   -4.0586,
                   -2.5332,
                   -3.3809,
                   4.3867,
               ]  ,						device=UpperCamelCase__  ,						dtype=torch.half  ,						)

           for i in range(19  ):
                  lowerCamelCase				,   lowerCamelCase     :   Optional[int]    				=      8 * i // 17, 2792 * i  # along the diagonal
                  lowerCamelCase				,   lowerCamelCase     :   Tuple    				=      output[0, r, c], expected_diag[i]
                  lowerCamelCase     :   List[str]    				=      F'''row={r} col={c} computed={computed} expected={expected}'''
                  self.assertAlmostEqual(UpperCamelCase__  ,						UpperCamelCase__  ,						delta=1e-4  ,						msg=UpperCamelCase__  )






