






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






						def 					_lowerCAmelCase				(		self )       -> Union[str, Any]:
								snake_case_ :							int         =				ids_tensor([self.batch_size, self.seq_length]						, self.vocab_size )

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

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

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

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

								snake_case_ :							Optional[int]         =				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=_SCREAMING_SNAKE_CASE						, )

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

								snake_case_ :							int         =				[input_ids, None, input_mask]  # None is the input for 'past'
								snake_case_ :							List[str]         =				model(_SCREAMING_SNAKE_CASE )

								snake_case_ :							Dict         =				model(_SCREAMING_SNAKE_CASE )

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






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, *_SCREAMING_SNAKE_CASE )       -> Union[str, Any]:
								snake_case_ :							List[str]         =				TFGPTJModel(config=_SCREAMING_SNAKE_CASE )

								# first forward pass
								snake_case_ :							List[Any]         =				model(_SCREAMING_SNAKE_CASE						, token_type_ids=_SCREAMING_SNAKE_CASE						, use_cache=_SCREAMING_SNAKE_CASE )
								snake_case_ :							Any         =				model(_SCREAMING_SNAKE_CASE						, token_type_ids=_SCREAMING_SNAKE_CASE )
								snake_case_ :							Union[str, Any]         =				model(_SCREAMING_SNAKE_CASE						, token_type_ids=_SCREAMING_SNAKE_CASE						, use_cache=_SCREAMING_SNAKE_CASE )

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

								snake_case_       ,				snake_case_ :							str         =				outputs.to_tuple()

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

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

								snake_case_ :							List[str]         =				model(_SCREAMING_SNAKE_CASE						, token_type_ids=_SCREAMING_SNAKE_CASE )["last_hidden_state"]
								snake_case_ :							Any         =				model(_SCREAMING_SNAKE_CASE						, token_type_ids=_SCREAMING_SNAKE_CASE						, past_key_values=_SCREAMING_SNAKE_CASE )[
								    "last_hidden_state"
								]

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

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






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

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

								# first forward pass
								snake_case_       ,				snake_case_ :							int         =				model(_SCREAMING_SNAKE_CASE						, attention_mask=_SCREAMING_SNAKE_CASE ).to_tuple()

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

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

								# append to next input_ids and attn_mask
								snake_case_ :							List[str]         =				tf.concat([input_ids, next_tokens]						, axis=-1 )
								snake_case_ :							Optional[int]         =				tf.concat([attn_mask, tf.ones((shape_list(_SCREAMING_SNAKE_CASE )[0], 1)						, dtype=tf.intaa )]						, axis=1 )

								# get two different outputs
								snake_case_ :							List[Any]         =				model(_SCREAMING_SNAKE_CASE						, attention_mask=_SCREAMING_SNAKE_CASE )["last_hidden_state"]
								snake_case_ :							int         =				model(_SCREAMING_SNAKE_CASE						, past_key_values=_SCREAMING_SNAKE_CASE						, attention_mask=_SCREAMING_SNAKE_CASE )[
								    "last_hidden_state"
								]

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

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






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, *_SCREAMING_SNAKE_CASE )       -> Optional[int]:
								snake_case_ :							List[Any]         =				TFGPTJModel(config=_SCREAMING_SNAKE_CASE )

								snake_case_ :							Dict         =				input_ids[:1, :]
								snake_case_ :							int         =				input_mask[:1, :]
								snake_case_ :							Optional[Any]         =				token_type_ids[:1, :]
								snake_case_ :							List[Any]         =				1

								# first forward pass
								snake_case_ :							List[str]         =				model(_SCREAMING_SNAKE_CASE						, attention_mask=_SCREAMING_SNAKE_CASE						, token_type_ids=_SCREAMING_SNAKE_CASE						, use_cache=_SCREAMING_SNAKE_CASE )

								snake_case_       ,				snake_case_ :							Any         =				outputs.to_tuple()

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

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

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

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

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






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






						def 					_lowerCAmelCase				(		self )       -> Dict:
								snake_case_ :							Union[str, Any]         =				self.prepare_config_and_inputs()

								(
								    (
								    snake_case_
								)       ,				(
								    snake_case_
								)       ,				(
								    snake_case_
								)       ,				(
								    snake_case_
								)       ,				(
								    snake_case_
								)       ,				(
								    snake_case_
								)       ,				(
								    snake_case_
								)       ,				(
								    snake_case_
								)       ,				(
								    snake_case_
								)       ,				
								) :							Optional[Any]         =				config_and_inputs

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




@require_tf
class 						UpperCAmelCase_  (		SCREAMING_SNAKE_CASE__    ,     SCREAMING_SNAKE_CASE__    ,     SCREAMING_SNAKE_CASE__    ,     unittest.TestCase    ):


						'''simple docstring'''




						A   :      Union[str, Any]						 =    (
						    (TFGPTJForCausalLM, TFGPTJForSequenceClassification, TFGPTJForQuestionAnswering, TFGPTJModel)
						    if is_tf_available()
						    else ()
						)

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






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )       -> 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 					_lowerCAmelCase				(		self )       -> Any:
								snake_case_ :							Dict         =				TFGPTJModelTester(self )
								snake_case_ :							Dict         =				ConfigTester(self						, config_class=_SCREAMING_SNAKE_CASE						, n_embd=37 )






						def 					_lowerCAmelCase				(		self )       -> int:
								self.config_tester.run_common_tests()






						def 					_lowerCAmelCase				(		self )       -> str:
								snake_case_ :							str         =				self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_gptj_model(*_SCREAMING_SNAKE_CASE )






						def 					_lowerCAmelCase				(		self )       -> Dict:
								snake_case_ :							Dict         =				self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_gptj_model_past(*_SCREAMING_SNAKE_CASE )






						def 					_lowerCAmelCase				(		self )       -> Tuple:
								snake_case_ :							Any         =				self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_gptj_model_attention_mask_past(*_SCREAMING_SNAKE_CASE )






						def 					_lowerCAmelCase				(		self )       -> Optional[Any]:
								snake_case_ :							Tuple         =				self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_gptj_model_past_large_inputs(*_SCREAMING_SNAKE_CASE )






						def 					_lowerCAmelCase				(		self )       -> Optional[Any]:
								snake_case_ :							Tuple         =				self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_gptj_lm_head_model(*_SCREAMING_SNAKE_CASE )






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






						@unittest.skip(reason="Currently, model embeddings are going to undergo a major refactor." )
						def 					_lowerCAmelCase				(		self )       -> List[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 					_lowerCAmelCase				(		self )       -> List[str]:
								snake_case_ :							Any         =				TFGPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B"						, from_pt=_SCREAMING_SNAKE_CASE )
								snake_case_ :							Dict         =				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_ :							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
								snake_case_ :							List[str]         =				model.generate(_SCREAMING_SNAKE_CASE						, do_sample=_SCREAMING_SNAKE_CASE )
								self.assertListEqual(output_ids[0].numpy().tolist()						, _SCREAMING_SNAKE_CASE )






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

								snake_case_ :							Any         =				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_ :							int         =				model.generate(**_SCREAMING_SNAKE_CASE						, do_sample=_SCREAMING_SNAKE_CASE						, seed=[42, 0] )
								snake_case_ :							List[str]         =				tokenizer.decode(output_ids[0]						, skip_special_tokens=_SCREAMING_SNAKE_CASE )

								snake_case_ :							List[str]         =				"Today is a nice day and I’m going to go for a walk. I’"
								self.assertEqual(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )






						def 					_lowerCAmelCase				(		self )       -> Dict:
								snake_case_ :							Optional[Any]         =				TFGPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B"						, revision="float16"						, from_pt=_SCREAMING_SNAKE_CASE )
								snake_case_ :							Dict         =				AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B"						, revision="float16" )

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

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

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

								snake_case_ :							Dict         =				tokenizer(_SCREAMING_SNAKE_CASE						, return_tensors="tf"						, padding=_SCREAMING_SNAKE_CASE )
								snake_case_ :							Tuple         =				model.generate(**_SCREAMING_SNAKE_CASE						, do_sample=_SCREAMING_SNAKE_CASE						, num_beams=2 )
								snake_case_ :							str         =				tokenizer.batch_decode(_SCREAMING_SNAKE_CASE						, skip_special_tokens=_SCREAMING_SNAKE_CASE )
								self.assertListEqual(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )






						def 					_lowerCAmelCase				(		self )       -> Optional[int]:
								# Confirms that left-padding is working properly
								snake_case_       ,				snake_case_       ,				snake_case_       ,				snake_case_ :							List[Any]         =				self._get_beam_search_test_objects()

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






						def 					_lowerCAmelCase				(		self )       -> str:
								# Confirms that XLA is working properly
								snake_case_       ,				snake_case_       ,				snake_case_       ,				snake_case_ :							int         =				self._get_beam_search_test_objects()

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








from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowercase     :							Any	     =			logging.get_logger(__name__)

lowercase     :							int	     =			{
    # pcqm4mv1 now deprecated
    '''graphormer-base''': '''https://huggingface.co/clefourrier/graphormer-base-pcqm4mv2/resolve/main/config.json''',
    # See all Graphormer models at https://huggingface.co/models?filter=graphormer
}




class 						UpperCAmelCase_  (		SCREAMING_SNAKE_CASE__    ):


						'''simple docstring'''




						A   :      Optional[int]						 =    'graphormer'
						A   :      List[str]						 =    ['past_key_values']






						def __init__(		self						, _SCREAMING_SNAKE_CASE = 1						, _SCREAMING_SNAKE_CASE = 512 * 9						, _SCREAMING_SNAKE_CASE = 512 * 3						, _SCREAMING_SNAKE_CASE = 512						, _SCREAMING_SNAKE_CASE = 512						, _SCREAMING_SNAKE_CASE = 512						, _SCREAMING_SNAKE_CASE = 128						, _SCREAMING_SNAKE_CASE = 5						, _SCREAMING_SNAKE_CASE = 1024						, _SCREAMING_SNAKE_CASE = "multi_hop"						, _SCREAMING_SNAKE_CASE = 512						, _SCREAMING_SNAKE_CASE = False						, _SCREAMING_SNAKE_CASE = 12						, _SCREAMING_SNAKE_CASE = 768						, _SCREAMING_SNAKE_CASE = 768						, _SCREAMING_SNAKE_CASE = 32						, _SCREAMING_SNAKE_CASE = 0.1						, _SCREAMING_SNAKE_CASE = 0.1						, _SCREAMING_SNAKE_CASE = 0.0						, _SCREAMING_SNAKE_CASE = False						, _SCREAMING_SNAKE_CASE = False						, _SCREAMING_SNAKE_CASE = False						, _SCREAMING_SNAKE_CASE = "gelu"						, _SCREAMING_SNAKE_CASE = None						, _SCREAMING_SNAKE_CASE = False						, _SCREAMING_SNAKE_CASE = 0						, _SCREAMING_SNAKE_CASE = False						, _SCREAMING_SNAKE_CASE = 0.0						, _SCREAMING_SNAKE_CASE = 8						, _SCREAMING_SNAKE_CASE = None						, _SCREAMING_SNAKE_CASE = None						, _SCREAMING_SNAKE_CASE = True						, _SCREAMING_SNAKE_CASE = True						, _SCREAMING_SNAKE_CASE=0						, _SCREAMING_SNAKE_CASE=1						, _SCREAMING_SNAKE_CASE=2						, **_SCREAMING_SNAKE_CASE						, )       -> List[str]:
								snake_case_ :							str         =				num_classes
								snake_case_ :							Union[str, Any]         =				num_atoms
								snake_case_ :							Dict         =				num_in_degree
								snake_case_ :							Union[str, Any]         =				num_out_degree
								snake_case_ :							List[Any]         =				num_edges
								snake_case_ :							Union[str, Any]         =				num_spatial
								snake_case_ :							int         =				num_edge_dis
								snake_case_ :							List[Any]         =				edge_type
								snake_case_ :							Any         =				multi_hop_max_dist
								snake_case_ :							Tuple         =				spatial_pos_max
								snake_case_ :							Union[str, Any]         =				max_nodes
								snake_case_ :							str         =				num_hidden_layers
								snake_case_ :							str         =				embedding_dim
								snake_case_ :							List[str]         =				embedding_dim
								snake_case_ :							Dict         =				ffn_embedding_dim
								snake_case_ :							Any         =				num_attention_heads
								snake_case_ :							List[str]         =				dropout
								snake_case_ :							List[str]         =				attention_dropout
								snake_case_ :							Optional[Any]         =				layerdrop
								snake_case_ :							Optional[int]         =				encoder_normalize_before
								snake_case_ :							Union[str, Any]         =				pre_layernorm
								snake_case_ :							int         =				apply_graphormer_init
								snake_case_ :							List[Any]         =				activation_fn
								snake_case_ :							str         =				embed_scale
								snake_case_ :							int         =				freeze_embeddings
								snake_case_ :							Optional[Any]         =				num_trans_layers_to_freeze
								snake_case_ :							int         =				share_input_output_embed
								snake_case_ :							int         =				traceable
								snake_case_ :							Tuple         =				q_noise
								snake_case_ :							Tuple         =				qn_block_size

								# These parameters are here for future extensions
								# atm, the model only supports self attention
								snake_case_ :							List[Any]         =				kdim
								snake_case_ :							Tuple         =				vdim
								snake_case_ :							Any         =				self_attention
								snake_case_ :							Any         =				bias

								super().__init__(
								    pad_token_id=_SCREAMING_SNAKE_CASE						, bos_token_id=_SCREAMING_SNAKE_CASE						, eos_token_id=_SCREAMING_SNAKE_CASE						, **_SCREAMING_SNAKE_CASE						, )
