




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__  )  ->   str:
											lowerCamelCase     :   Optional[int]    				=      parent
											lowerCamelCase     :   Dict    				=      13
											lowerCamelCase     :   List[str]    				=      7
											lowerCamelCase     :   str    				=      True
											lowerCamelCase     :   int    				=      True
											lowerCamelCase     :   Optional[int]    				=      True
											lowerCamelCase     :   Dict    				=      True
											lowerCamelCase     :   List[str]    				=      True
											lowerCamelCase     :   Dict    				=      99
											lowerCamelCase     :   Union[str, Any]    				=      32
											lowerCamelCase     :   Optional[int]    				=      4
											lowerCamelCase     :   str    				=      2
											lowerCamelCase     :   Any    				=      4
											lowerCamelCase     :   List[str]    				=      37
											lowerCamelCase     :   int    				=      "gelu"
											lowerCamelCase     :   str    				=      0.1
											lowerCamelCase     :   int    				=      0.1
											lowerCamelCase     :   Optional[int]    				=      512
											lowerCamelCase     :   List[Any]    				=      16
											lowerCamelCase     :   Any    				=      2
											lowerCamelCase     :   str    				=      0.02
											lowerCamelCase     :   List[Any]    				=      3
											lowerCamelCase     :   List[str]    				=      4
											lowerCamelCase     :   str    				=      None
											lowerCamelCase     :   Union[str, Any]    				=      self.vocab_size - 1
											lowerCamelCase     :   Union[str, Any]    				=      self.vocab_size - 1
											lowerCamelCase     :   Union[str, Any]    				=      self.vocab_size - 1







				def 					_lowercase  (					self  )  ->   Optional[Any]:
											lowerCamelCase     :   List[str]    				=      ids_tensor([self.batch_size, self.seq_length]  ,						self.vocab_size  )

											lowerCamelCase     :   Dict    				=      None
											if self.use_input_mask:
																		lowerCamelCase     :   int    				=      random_attention_mask([self.batch_size, self.seq_length]  )

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

											lowerCamelCase     :   Optional[Any]    				=      None
											if self.use_mc_token_ids:
																		lowerCamelCase     :   Union[str, Any]    				=      ids_tensor([self.batch_size, self.num_choices]  ,						self.seq_length  )

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

											lowerCamelCase     :   Optional[Any]    				=      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=a_  ,						)

											lowerCamelCase     :   Any    				=      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     :   Optional[Any]    				=      TFGPTJModel(config=a_  )
											lowerCamelCase     :   List[Any]    				=      {
											    "input_ids": input_ids,
											    "attention_mask": input_mask,
											    "token_type_ids": token_type_ids,
											}
											lowerCamelCase     :   Optional[int]    				=      model(a_  )

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

											lowerCamelCase     :   int    				=      model(a_  )

											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__  )  ->   List[str]:
											lowerCamelCase     :   Dict    				=      TFGPTJModel(config=a_  )

											# first forward pass
											lowerCamelCase     :   Tuple    				=      model(a_  ,						token_type_ids=a_  ,						use_cache=a_  )
											lowerCamelCase     :   Optional[Any]    				=      model(a_  ,						token_type_ids=a_  )
											lowerCamelCase     :   Union[str, Any]    				=      model(a_  ,						token_type_ids=a_  ,						use_cache=a_  )

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

											lowerCamelCase				,   lowerCamelCase     :   int    				=      outputs.to_tuple()

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

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

											lowerCamelCase     :   Union[str, Any]    				=      model(a_  ,						token_type_ids=a_  )["last_hidden_state"]
											lowerCamelCase     :   List[Any]    				=      model(a_  ,						token_type_ids=a_  ,						past_key_values=a_  )[
											    "last_hidden_state"
											]

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

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







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

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

											# first forward pass
											lowerCamelCase				,   lowerCamelCase     :   Optional[Any]    				=      model(a_  ,						attention_mask=a_  ).to_tuple()

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

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

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

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

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

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







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

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

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

											lowerCamelCase				,   lowerCamelCase     :   Tuple    				=      outputs.to_tuple()

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

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

											lowerCamelCase     :   str    				=      model(
											    a_  ,						token_type_ids=a_  ,						attention_mask=a_  )["last_hidden_state"]
											lowerCamelCase     :   List[Any]    				=      model(
											    a_  ,						token_type_ids=a_  ,						attention_mask=a_  ,						past_key_values=a_  ,						)["last_hidden_state"]
											self.parent.assertTrue(output_from_past.shape[1] == next_tokens.shape[1]  )

											# select random slice
											lowerCamelCase     :   Tuple    				=      int(ids_tensor((1,)  ,						shape_list(a_  )[-1]  )  )
											lowerCamelCase     :   Optional[Any]    				=      output_from_no_past[:, -3:, random_slice_idx]
											lowerCamelCase     :   List[str]    				=      output_from_past[:, :, random_slice_idx]

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







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







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

											(
											    (
											    lowerCamelCase      
											)				,   (
											    lowerCamelCase      
											)				,   (
											    lowerCamelCase      
											)				,   (
											    lowerCamelCase      
											)				,   (
											    lowerCamelCase      
											)				,   (
											    lowerCamelCase      
											)				,   (
											    lowerCamelCase      
											)				,   (
											    lowerCamelCase      
											)				,   (
											    lowerCamelCase      
											)				,         
											)     :   Any    				=      config_and_inputs

											lowerCamelCase     :   Dict    				=      {
											    "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_   :			List[Any]			  =   (
				    (TFGPTJForCausalLM, TFGPTJForSequenceClassification, TFGPTJForQuestionAnswering, TFGPTJModel)
				    if is_tf_available()
				    else ()
				)

				lowerCamelCase_   :			int			  =   (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_   :			Dict			  =   False
				lowerCamelCase_   :			Optional[int]			  =   False
				lowerCamelCase_   :			List[Any]			  =   False
				lowerCamelCase_   :			List[str]			  =   False







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







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







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







				def 					_lowercase  (					self  )  ->   List[Any]:
											lowerCamelCase     :   Dict    				=      self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_gptj_model_past(*a_  )







				def 					_lowercase  (					self  )  ->   Optional[Any]:
											lowerCamelCase     :   Any    				=      self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_gptj_model_attention_mask_past(*a_  )







				def 					_lowercase  (					self  )  ->   List[str]:
											lowerCamelCase     :   Union[str, Any]    				=      self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_gptj_model_past_large_inputs(*a_  )







				def 					_lowercase  (					self  )  ->   int:
											lowerCamelCase     :   Optional[Any]    				=      self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_gptj_lm_head_model(*a_  )







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







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

				'''simple docstring'''



				def 					_lowercase  (					self  )  ->   Tuple:
											lowerCamelCase     :   Union[str, Any]    				=      TFGPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B"  ,						from_pt=a_  )
											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     :   Dict    				=      [464, 3290, 318, 257, 582, 338, 1266, 1545, 13, 632, 318, 257, 9112, 1_5185, 11, 290, 340, 318, 257, 1545]
											# fmt: on
											lowerCamelCase     :   List[str]    				=      model.generate(a_  ,						do_sample=a_  )
											self.assertListEqual(output_ids[0].numpy().tolist()  ,						a_  )







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

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

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







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

											lowerCamelCase     :   Union[str, Any]    				=      "left"

											# Define PAD Token = EOS Token = 50256
											lowerCamelCase     :   Any    				=      tokenizer.eos_token
											lowerCamelCase     :   Tuple    				=      model.config.eos_token_id

											# use different length sentences to test batching
											lowerCamelCase     :   List[str]    				=      [
											    "Hello, my dog is a little",
											    "Today, I",
											]
											lowerCamelCase     :   List[str]    				=      [
											    "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  )  ->   Optional[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     :   Union[str, Any]    				=      tokenizer(a_  ,						return_tensors="tf"  ,						padding=a_  )
											lowerCamelCase     :   List[str]    				=      model.generate(**a_  ,						do_sample=a_  ,						num_beams=2  )
											lowerCamelCase     :   Union[str, Any]    				=      tokenizer.batch_decode(a_  ,						skip_special_tokens=a_  )
											self.assertListEqual(a_  ,						a_  )







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

											lowerCamelCase     :   Optional[int]    				=      tokenizer(a_  ,						return_tensors="tf"  ,						padding=a_  )
											lowerCamelCase     :   Union[str, Any]    				=      tokenizer(sentences[0]  ,						return_tensors="tf"  )
											lowerCamelCase     :   List[Any]    				=      model.generate(**a_  ,						do_sample=a_  ,						num_beams=2  )
											lowerCamelCase     :   Optional[Any]    				=      (
											    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     :   Tuple    				=      model.generate(
											    **a_  ,						do_sample=a_  ,						num_beams=2  ,						max_length=model.config.max_length - num_paddings  )
											lowerCamelCase     :   List[Any]    				=      tokenizer.decode(output_non_padded[0]  ,						skip_special_tokens=a_  )
											lowerCamelCase     :   str    				=      tokenizer.decode(output_padded[0]  ,						skip_special_tokens=a_  )
											self.assertListEqual(a_  ,						[non_padded_sentence, padded_sentence]  )







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

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












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__  )






