


"""simple docstring"""







from __future__ import annotations

import unittest

from transformers import LEDConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
		import tensorflow as tf

		from transformers import TFLEDForConditionalGeneration, TFLEDModel




@require_tf
class  UpperCamelCase_	:

	"""simple docstring"""





	snake_case__				:  int    					=			LEDConfig
	snake_case__				:  Tuple    					=			{}
	snake_case__				:  List[Any]    					=			"gelu"




	def __init__( self			:				Dict			,			UpperCAmelCase__			:				Dict			,			UpperCAmelCase__			:				List[str]=1_3			,			UpperCAmelCase__			:				Any=7			,			UpperCAmelCase__			:				Optional[int]=True			,			UpperCAmelCase__			:				str=False			,			UpperCAmelCase__			:				Optional[int]=9_9			,			UpperCAmelCase__			:				int=3_2			,			UpperCAmelCase__			:				Union[str, Any]=2			,			UpperCAmelCase__			:				Dict=4			,			UpperCAmelCase__			:				int=3_7			,			UpperCAmelCase__			:				Optional[int]=0.1			,			UpperCAmelCase__			:				Optional[int]=0.1			,			UpperCAmelCase__			:				List[str]=2_0			,			UpperCAmelCase__			:				int=2			,			UpperCAmelCase__			:				Dict=1			,			UpperCAmelCase__			:				Union[str, Any]=0			,			UpperCAmelCase__			:				List[str]=4			,			)       ->	Dict:
						__SCREAMING_SNAKE_CASE						   =				parent
						__SCREAMING_SNAKE_CASE						   =				batch_size
						__SCREAMING_SNAKE_CASE						   =				seq_length
						__SCREAMING_SNAKE_CASE						   =				is_training
						__SCREAMING_SNAKE_CASE						   =				use_labels
						__SCREAMING_SNAKE_CASE						   =				vocab_size
						__SCREAMING_SNAKE_CASE						   =				hidden_size
						__SCREAMING_SNAKE_CASE						   =				num_hidden_layers
						__SCREAMING_SNAKE_CASE						   =				num_attention_heads
						__SCREAMING_SNAKE_CASE						   =				intermediate_size

						__SCREAMING_SNAKE_CASE						   =				hidden_dropout_prob
						__SCREAMING_SNAKE_CASE						   =				attention_probs_dropout_prob
						__SCREAMING_SNAKE_CASE						   =				max_position_embeddings
						__SCREAMING_SNAKE_CASE						   =				eos_token_id
						__SCREAMING_SNAKE_CASE						   =				pad_token_id
						__SCREAMING_SNAKE_CASE						   =				bos_token_id
						__SCREAMING_SNAKE_CASE						   =				attention_window

						# `ModelTesterMixin.test_attention_outputs` is expecting attention tensors to be of size
						# [num_attention_heads, encoder_seq_length, encoder_key_length], but TFLongformerSelfAttention
						# returns attention of shape [num_attention_heads, encoder_seq_length, self.attention_window + 1]
						# because its local attention only attends to `self.attention_window` and one before and one after
						__SCREAMING_SNAKE_CASE						   =				self.attention_window + 2

						# because of padding `encoder_seq_length`, is different from `seq_length`. Relevant for
						# the `test_attention_outputs` and `test_hidden_states_output` tests
						__SCREAMING_SNAKE_CASE						   =				(
						    self.seq_length + (self.attention_window - self.seq_length % self.attention_window) % self.attention_window
						)




	def   UpperCAmelCase_      ( self			:				Tuple   )       ->	List[str]:
						__SCREAMING_SNAKE_CASE						   =				ids_tensor([self.batch_size, self.seq_length - 1]			,			self.vocab_size   )
						__SCREAMING_SNAKE_CASE						   =				tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size   )			,			1   )
						__SCREAMING_SNAKE_CASE						   =				tf.concat([input_ids, eos_tensor]			,			axis=1   )

						__SCREAMING_SNAKE_CASE						   =				ids_tensor([self.batch_size, self.seq_length]			,			self.vocab_size   )

						__SCREAMING_SNAKE_CASE						   =				self.config_cls(
						    vocab_size=self.vocab_size			,			d_model=self.hidden_size			,			encoder_layers=self.num_hidden_layers			,			decoder_layers=self.num_hidden_layers			,			encoder_attention_heads=self.num_attention_heads			,			decoder_attention_heads=self.num_attention_heads			,			encoder_ffn_dim=self.intermediate_size			,			decoder_ffn_dim=self.intermediate_size			,			dropout=self.hidden_dropout_prob			,			attention_dropout=self.attention_probs_dropout_prob			,			max_position_embeddings=self.max_position_embeddings			,			eos_token_ids=[2]			,			bos_token_id=self.bos_token_id			,			pad_token_id=self.pad_token_id			,			decoder_start_token_id=self.pad_token_id			,			attention_window=self.attention_window			,			**self.config_updates			,			)
						__SCREAMING_SNAKE_CASE						   =				prepare_led_inputs_dict(UpperCAmelCase__			,			UpperCAmelCase__			,			UpperCAmelCase__   )
						__SCREAMING_SNAKE_CASE						   =				tf.concat(
						    [tf.zeros_like(UpperCAmelCase__   )[:, :-1], tf.ones_like(UpperCAmelCase__   )[:, -1:]]			,			axis=-1			,			)
						__SCREAMING_SNAKE_CASE						   =				global_attention_mask
						return config, inputs_dict




	def   UpperCAmelCase_      ( self			:				Optional[Any]			,			UpperCAmelCase__			:				Tuple			,			UpperCAmelCase__			:				Union[str, Any]   )       ->	int:
						__SCREAMING_SNAKE_CASE						   =				TFLEDModel(config=UpperCAmelCase__   ).get_decoder()
						__SCREAMING_SNAKE_CASE						   =				inputs_dict["input_ids"]

						__SCREAMING_SNAKE_CASE						   =				input_ids[:1, :]
						__SCREAMING_SNAKE_CASE						   =				inputs_dict["attention_mask"][:1, :]
						__SCREAMING_SNAKE_CASE						   =				1

						# first forward pass
						__SCREAMING_SNAKE_CASE						   =				model(UpperCAmelCase__			,			attention_mask=UpperCAmelCase__			,			use_cache=UpperCAmelCase__   )

						__SCREAMING_SNAKE_CASE   ,  __SCREAMING_SNAKE_CASE						   =				outputs.to_tuple()

						# create hypothetical next token and extent to next_input_ids
						__SCREAMING_SNAKE_CASE						   =				ids_tensor((self.batch_size, 3)			,			config.vocab_size   )
						__SCREAMING_SNAKE_CASE						   =				tf.cast(ids_tensor((self.batch_size, 3)			,			2   )			,			tf.inta   )

						# append to next input_ids and
						__SCREAMING_SNAKE_CASE						   =				tf.concat([input_ids, next_tokens]			,			axis=-1   )
						__SCREAMING_SNAKE_CASE						   =				tf.concat([attention_mask, next_attn_mask]			,			axis=-1   )

						__SCREAMING_SNAKE_CASE						   =				model(UpperCAmelCase__			,			attention_mask=UpperCAmelCase__   )[0]
						__SCREAMING_SNAKE_CASE						   =				model(UpperCAmelCase__			,			attention_mask=UpperCAmelCase__			,			past_key_values=UpperCAmelCase__   )[0]

						self.parent.assertEqual(next_tokens.shape[1]			,			output_from_past.shape[1]   )

						# select random slice
						__SCREAMING_SNAKE_CASE						   =				int(ids_tensor((1,)			,			output_from_past.shape[-1]   )   )
						__SCREAMING_SNAKE_CASE						   =				output_from_no_past[:, -3:, random_slice_idx]
						__SCREAMING_SNAKE_CASE						   =				output_from_past[:, :, random_slice_idx]

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




def    UpperCAmelCase__       (lowerCAmelCase_  ,						lowerCAmelCase_  ,						lowerCAmelCase_  ,						lowerCAmelCase_=None  ,						lowerCAmelCase_=None  ,						lowerCAmelCase_=None  ,						lowerCAmelCase_=None  ,						):




					'''simple docstring'''


					if attention_mask is None:
										__SCREAMING_SNAKE_CASE						   =				tf.cast(tf.math.not_equal(lowerCAmelCase_  ,						config.pad_token_id   )  ,						tf.inta   )
					if decoder_attention_mask is None:
										__SCREAMING_SNAKE_CASE						   =				tf.concat(
										    [
										        tf.ones(decoder_input_ids[:, :1].shape  ,						dtype=tf.inta   ),
										        tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:]  ,						config.pad_token_id   )  ,						tf.inta   ),
										    ]  ,						axis=-1  ,						)
					if head_mask is None:
										__SCREAMING_SNAKE_CASE						   =				tf.ones((config.encoder_layers, config.encoder_attention_heads)   )
					if decoder_head_mask is None:
										__SCREAMING_SNAKE_CASE						   =				tf.ones((config.decoder_layers, config.decoder_attention_heads)   )
					return {
					    "input_ids": input_ids,
					    "attention_mask": attention_mask,
					    "decoder_input_ids": decoder_input_ids,
					    "decoder_attention_mask": decoder_attention_mask,
					    "head_mask": head_mask,
					    "decoder_head_mask": decoder_head_mask,
					}




@require_tf
class  UpperCamelCase_	( UpperCamelCase			,  UpperCamelCase			,  unittest.TestCase):

	"""simple docstring"""





	snake_case__				:  str    					=			(TFLEDForConditionalGeneration, TFLEDModel) if is_tf_available() else ()
	snake_case__				:  List[str]    					=			(TFLEDForConditionalGeneration,) if is_tf_available() else ()
	snake_case__				:  Union[str, Any]    					=			(
	    {
	        "conversational": TFLEDForConditionalGeneration,
	        "feature-extraction": TFLEDModel,
	        "summarization": TFLEDForConditionalGeneration,
	        "text2text-generation": TFLEDForConditionalGeneration,
	        "translation": TFLEDForConditionalGeneration,
	    }
	    if is_tf_available()
	    else {}
	)
	snake_case__				:  str    					=			True
	snake_case__				:  Tuple    					=			False
	snake_case__				:  List[str]    					=			False
	snake_case__				:  int    					=			False




	def   UpperCAmelCase_      ( self			:				List[str]   )       ->	Union[str, Any]:
						__SCREAMING_SNAKE_CASE						   =				TFLEDModelTester(self   )
						__SCREAMING_SNAKE_CASE						   =				ConfigTester(self			,			config_class=UpperCAmelCase__   )




	def   UpperCAmelCase_      ( self			:				Optional[Any]   )       ->	Tuple:
						self.config_tester.run_common_tests()




	def   UpperCAmelCase_      ( self			:				Optional[int]   )       ->	Optional[Any]:
						__SCREAMING_SNAKE_CASE						   =				self.model_tester.prepare_config_and_inputs_for_common()
						self.model_tester.check_decoder_model_past_large_inputs(*UpperCAmelCase__   )




	def   UpperCAmelCase_      ( self			:				Any   )       ->	Optional[int]:
						__SCREAMING_SNAKE_CASE   ,  __SCREAMING_SNAKE_CASE						   =				self.model_tester.prepare_config_and_inputs_for_common()
						__SCREAMING_SNAKE_CASE						   =				tf.zeros_like(inputs_dict["attention_mask"]   )
						__SCREAMING_SNAKE_CASE						   =				2
						__SCREAMING_SNAKE_CASE						   =				tf.where(
						    tf.range(self.model_tester.seq_length   )[None, :] < num_global_attn_indices			,			1			,			inputs_dict["global_attention_mask"]			,			)

						__SCREAMING_SNAKE_CASE						   =				True
						__SCREAMING_SNAKE_CASE						   =				self.model_tester.seq_length
						__SCREAMING_SNAKE_CASE						   =				self.model_tester.encoder_seq_length

						def check_decoder_attentions_output(UpperCAmelCase__			:				Optional[int]   ):
											__SCREAMING_SNAKE_CASE						   =				outputs.decoder_attentions
											self.assertEqual(len(UpperCAmelCase__   )			,			self.model_tester.num_hidden_layers   )
											self.assertListEqual(
											    list(decoder_attentions[0].shape[-3:]   )			,			[self.model_tester.num_attention_heads, seq_length, seq_length]			,			)

						def check_encoder_attentions_output(UpperCAmelCase__			:				Optional[Any]   ):
											__SCREAMING_SNAKE_CASE						   =				[t.numpy() for t in outputs.encoder_attentions]
											__SCREAMING_SNAKE_CASE						   =				[t.numpy() for t in outputs.encoder_global_attentions]
											self.assertEqual(len(UpperCAmelCase__   )			,			self.model_tester.num_hidden_layers   )
											self.assertEqual(len(UpperCAmelCase__   )			,			self.model_tester.num_hidden_layers   )
											self.assertListEqual(
											    list(attentions[0].shape[-3:]   )			,			[self.model_tester.num_attention_heads, seq_length, seq_length]			,			)
											self.assertListEqual(
											    list(global_attentions[0].shape[-3:]   )			,			[self.model_tester.num_attention_heads, encoder_seq_length, num_global_attn_indices]			,			)

						for model_class in self.all_model_classes:
											__SCREAMING_SNAKE_CASE						   =				True
											__SCREAMING_SNAKE_CASE						   =				False
											__SCREAMING_SNAKE_CASE						   =				False
											__SCREAMING_SNAKE_CASE						   =				model_class(UpperCAmelCase__   )
											__SCREAMING_SNAKE_CASE						   =				model(self._prepare_for_class(UpperCAmelCase__			,			UpperCAmelCase__   )   )
											__SCREAMING_SNAKE_CASE						   =				len(UpperCAmelCase__   )
											self.assertEqual(config.output_hidden_states			,			UpperCAmelCase__   )
											check_encoder_attentions_output(UpperCAmelCase__   )

											if self.is_encoder_decoder:
																__SCREAMING_SNAKE_CASE						   =				model_class(UpperCAmelCase__   )
																__SCREAMING_SNAKE_CASE						   =				model(self._prepare_for_class(UpperCAmelCase__			,			UpperCAmelCase__   )   )
																self.assertEqual(config.output_hidden_states			,			UpperCAmelCase__   )
																check_decoder_attentions_output(UpperCAmelCase__   )

											# Check that output attentions can also be changed via the config
											del inputs_dict["output_attentions"]
											__SCREAMING_SNAKE_CASE						   =				True
											__SCREAMING_SNAKE_CASE						   =				model_class(UpperCAmelCase__   )
											__SCREAMING_SNAKE_CASE						   =				model(self._prepare_for_class(UpperCAmelCase__			,			UpperCAmelCase__   )   )
											self.assertEqual(config.output_hidden_states			,			UpperCAmelCase__   )
											check_encoder_attentions_output(UpperCAmelCase__   )

											# Check attention is always last and order is fine
											__SCREAMING_SNAKE_CASE						   =				True
											__SCREAMING_SNAKE_CASE						   =				True
											__SCREAMING_SNAKE_CASE						   =				model_class(UpperCAmelCase__   )
											__SCREAMING_SNAKE_CASE						   =				model(self._prepare_for_class(UpperCAmelCase__			,			UpperCAmelCase__   )   )

											self.assertEqual(out_len + (2 if self.is_encoder_decoder else 1)			,			len(UpperCAmelCase__   )   )
											self.assertEqual(model.config.output_hidden_states			,			UpperCAmelCase__   )
											check_encoder_attentions_output(UpperCAmelCase__   )




	@unittest.skip("LED keeps using potentially symbolic tensors in conditionals and breaks tracing."   )
	def   UpperCAmelCase_      ( self			:				List[Any]   )       ->	Optional[Any]:
						pass




	def   UpperCAmelCase_      ( self			:				List[str]   )       ->	str:
						# TODO: Head-masking not yet implement
						pass




def    UpperCAmelCase__       (lowerCAmelCase_   ):




					'''simple docstring'''


					return tf.constant(lowerCAmelCase_  ,						dtype=tf.intaa   )


a__				:							str						=     1E-4




@slow
@require_tf
class  UpperCamelCase_	( unittest.TestCase):

	"""simple docstring"""





	def   UpperCAmelCase_      ( self			:				int   )       ->	str:
						__SCREAMING_SNAKE_CASE						   =				TFLEDForConditionalGeneration.from_pretrained("allenai/led-base-16384"   ).led

						# change to intended input here
						__SCREAMING_SNAKE_CASE						   =				_long_tensor([5_1_2 * [0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9]]   )
						__SCREAMING_SNAKE_CASE						   =				_long_tensor([1_2_8 * [0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9]]   )
						__SCREAMING_SNAKE_CASE						   =				prepare_led_inputs_dict(model.config			,			UpperCAmelCase__			,			UpperCAmelCase__   )
						__SCREAMING_SNAKE_CASE						   =				model(**UpperCAmelCase__   )[0]
						__SCREAMING_SNAKE_CASE						   =				(1, 1_0_2_4, 7_6_8)
						self.assertEqual(output.shape			,			UpperCAmelCase__   )
						# change to expected output here
						__SCREAMING_SNAKE_CASE						   =				tf.convert_to_tensor(
						    [[2.3_050, 2.8_279, 0.6_531], [-1.8_457, -0.1_455, -3.5_661], [-1.0_186, 0.4_586, -2.2_043]]			,			)
						tf.debugging.assert_near(output[:, :3, :3]			,			UpperCAmelCase__			,			atol=1E-3   )




	def   UpperCAmelCase_      ( self			:				Any   )       ->	Optional[int]:
						__SCREAMING_SNAKE_CASE						   =				TFLEDForConditionalGeneration.from_pretrained("allenai/led-base-16384"   )

						# change to intended input here
						__SCREAMING_SNAKE_CASE						   =				_long_tensor([5_1_2 * [0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9]]   )
						__SCREAMING_SNAKE_CASE						   =				_long_tensor([1_2_8 * [0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9]]   )
						__SCREAMING_SNAKE_CASE						   =				prepare_led_inputs_dict(model.config			,			UpperCAmelCase__			,			UpperCAmelCase__   )
						__SCREAMING_SNAKE_CASE						   =				model(**UpperCAmelCase__   )[0]
						__SCREAMING_SNAKE_CASE						   =				(1, 1_0_2_4, model.config.vocab_size)
						self.assertEqual(output.shape			,			UpperCAmelCase__   )
						# change to expected output here
						__SCREAMING_SNAKE_CASE						   =				tf.convert_to_tensor(
						    [[33.6_507, 6.4_572, 16.8_089], [5.8_739, -2.4_238, 11.2_902], [-3.2_139, -4.3_149, 4.2_783]]			,			)
						tf.debugging.assert_near(output[:, :3, :3]			,			UpperCAmelCase__			,			atol=1E-3			,			rtol=1E-3   )




"""simple docstring"""







from abc import ABC, abstractmethod
from argparse import ArgumentParser




class  UpperCamelCase_	( UpperCamelCase):

	"""simple docstring"""





	@staticmethod
	@abstractmethod
	def   UpperCAmelCase_      ( UpperCAmelCase__			:				ArgumentParser   )       ->	int:
						raise NotImplementedError()




	@abstractmethod
	def   UpperCAmelCase_      ( self			:				int   )       ->	Optional[int]:
						raise NotImplementedError()
