





'''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 				__snake_case     :
	__lowerCAmelCase       						=  LEDConfig
	__lowerCAmelCase       						=  {}
	__lowerCAmelCase       						=  '''gelu'''





	def __init__(		self      ,       UpperCamelCase_      ,       UpperCamelCase_=13      ,       UpperCamelCase_=7      ,       UpperCamelCase_=True      ,       UpperCamelCase_=False      ,       UpperCamelCase_=99      ,       UpperCamelCase_=32      ,       UpperCamelCase_=2      ,       UpperCamelCase_=4      ,       UpperCamelCase_=37      ,       UpperCamelCase_=0.1      ,       UpperCamelCase_=0.1      ,       UpperCamelCase_=20      ,       UpperCamelCase_=2      ,       UpperCamelCase_=1      ,       UpperCamelCase_=0      ,       UpperCamelCase_=4      ,       )      ->		Any:
		snake_case__	       =   parent
		snake_case__	       =   batch_size
		snake_case__	       =   seq_length
		snake_case__	       =   is_training
		snake_case__	       =   use_labels
		snake_case__	       =   vocab_size
		snake_case__	       =   hidden_size
		snake_case__	       =   num_hidden_layers
		snake_case__	       =   num_attention_heads
		snake_case__	       =   intermediate_size

		snake_case__	       =   hidden_dropout_prob
		snake_case__	       =   attention_probs_dropout_prob
		snake_case__	       =   max_position_embeddings
		snake_case__	       =   eos_token_id
		snake_case__	       =   pad_token_id
		snake_case__	       =   bos_token_id
		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
		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
		snake_case__	       =   (
		    self.seq_length + (self.attention_window - self.seq_length % self.attention_window) % self.attention_window
		)





	def   _snake_case   (		self		)      ->		Any:
		snake_case__	       =   ids_tensor([self.batch_size, self.seq_length - 1]      ,       self.vocab_size		)
		snake_case__	       =   tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size		)      ,       1		)
		snake_case__	       =   tf.concat([input_ids, eos_tensor]      ,       axis=1		)

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

		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      ,       )
		snake_case__	       =   prepare_led_inputs_dict(UpperCamelCase_      ,       UpperCamelCase_      ,       UpperCamelCase_		)
		snake_case__	       =   tf.concat(
		    [tf.zeros_like(UpperCamelCase_		)[:, :-1], tf.ones_like(UpperCamelCase_		)[:, -1:]]      ,       axis=-1      ,       )
		snake_case__	       =   global_attention_mask
		return config, inputs_dict





	def   _snake_case   (		self      ,       UpperCamelCase_      ,       UpperCamelCase_		)      ->		List[Any]:
		snake_case__	       =   TFLEDModel(config=UpperCamelCase_		).get_decoder()
		snake_case__	       =   inputs_dict['input_ids']

		snake_case__	       =   input_ids[:1, :]
		snake_case__	       =   inputs_dict['attention_mask'][:1, :]
		snake_case__	       =   1

		# first forward pass
		snake_case__	       =   model(UpperCamelCase_      ,       attention_mask=UpperCamelCase_      ,       use_cache=UpperCamelCase_		)

		snake_case__       ,  snake_case__	       =   outputs.to_tuple()

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

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

		snake_case__	       =   model(UpperCamelCase_      ,       attention_mask=UpperCamelCase_		)[0]
		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
		snake_case__	       =   int(ids_tensor((1,)      ,       output_from_past.shape[-1]		)		)
		snake_case__	       =   output_from_no_past[:, -3:, random_slice_idx]
		snake_case__	       =   output_from_past[:, :, random_slice_idx]

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





def 						__lowerCamelCase			(	UpperCAmelCase_					,				UpperCAmelCase_					,				UpperCAmelCase_					,				UpperCAmelCase_=None					,				UpperCAmelCase_=None					,				UpperCAmelCase_=None					,				UpperCAmelCase_=None					,				)   ->Optional[int]:
	if attention_mask is None:
		snake_case__	       =   tf.cast(tf.math.not_equal(UpperCAmelCase_					,				config.pad_token_id							)					,				tf.inta							)
	if decoder_attention_mask is None:
		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:
		snake_case__	       =   tf.ones((config.encoder_layers, config.encoder_attention_heads)							)
	if decoder_head_mask is None:
		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 				__snake_case     (					__magic_name__       ,       __magic_name__       ,       unittest.TestCase							):
	__lowerCAmelCase       						=  (TFLEDForConditionalGeneration, TFLEDModel) if is_tf_available() else ()
	__lowerCAmelCase       						=  (TFLEDForConditionalGeneration,) if is_tf_available() else ()
	__lowerCAmelCase       						=  (
	    {
	        '''conversational''': TFLEDForConditionalGeneration,
	        '''feature-extraction''': TFLEDModel,
	        '''summarization''': TFLEDForConditionalGeneration,
	        '''text2text-generation''': TFLEDForConditionalGeneration,
	        '''translation''': TFLEDForConditionalGeneration,
	    }
	    if is_tf_available()
	    else {}
	)
	__lowerCAmelCase       						=  True
	__lowerCAmelCase       						=  False
	__lowerCAmelCase       						=  False
	__lowerCAmelCase       						=  False





	def   _snake_case   (		self		)      ->		Union[str, Any]:
		snake_case__	       =   TFLEDModelTester(self		)
		snake_case__	       =   ConfigTester(self      ,       config_class=UpperCamelCase_		)





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





	def   _snake_case   (		self		)      ->		str:
		snake_case__	       =   self.model_tester.prepare_config_and_inputs_for_common()
		self.model_tester.check_decoder_model_past_large_inputs(*UpperCamelCase_		)





	def   _snake_case   (		self		)      ->		str:
		snake_case__       ,  snake_case__	       =   self.model_tester.prepare_config_and_inputs_for_common()
		snake_case__	       =   tf.zeros_like(inputs_dict['attention_mask']		)
		snake_case__	       =   2
		snake_case__	       =   tf.where(
		    tf.range(self.model_tester.seq_length		)[None, :] < num_global_attn_indices      ,       1      ,       inputs_dict['global_attention_mask']      ,       )

		snake_case__	       =   True
		snake_case__	       =   self.model_tester.seq_length
		snake_case__	       =   self.model_tester.encoder_seq_length

		def check_decoder_attentions_output(UpperCamelCase_		):
			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_		):
			snake_case__	       =   [t.numpy() for t in outputs.encoder_attentions]
			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:
			snake_case__	       =   True
			snake_case__	       =   False
			snake_case__	       =   False
			snake_case__	       =   model_class(UpperCamelCase_		)
			snake_case__	       =   model(self._prepare_for_class(UpperCamelCase_      ,       UpperCamelCase_		)		)
			snake_case__	       =   len(UpperCamelCase_		)
			self.assertEqual(config.output_hidden_states      ,       UpperCamelCase_		)
			check_encoder_attentions_output(UpperCamelCase_		)

			if self.is_encoder_decoder:
				snake_case__	       =   model_class(UpperCamelCase_		)
				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"]
			snake_case__	       =   True
			snake_case__	       =   model_class(UpperCamelCase_		)
			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
			snake_case__	       =   True
			snake_case__	       =   True
			snake_case__	       =   model_class(UpperCamelCase_		)
			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   _snake_case   (		self		)      ->		Tuple:
		pass





	def   _snake_case   (		self		)      ->		Optional[Any]:
		# TODO: Head-masking not yet implement
		pass





def 						__lowerCamelCase			(	UpperCAmelCase_							)   ->Union[str, Any]:
	return tf.constant(UpperCAmelCase_					,				dtype=tf.intaa							)


a__					:			Optional[Any]          =     1e-4







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





	def   _snake_case   (		self		)      ->		Any:
		snake_case__	       =   TFLEDForConditionalGeneration.from_pretrained('allenai/led-base-16384'		).led

		# change to intended input here
		snake_case__	       =   _long_tensor([512 * [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69]]		)
		snake_case__	       =   _long_tensor([128 * [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69]]		)
		snake_case__	       =   prepare_led_inputs_dict(model.config      ,       UpperCamelCase_      ,       UpperCamelCase_		)
		snake_case__	       =   model(**UpperCamelCase_		)[0]
		snake_case__	       =   (1, 1024, 768)
		self.assertEqual(output.shape      ,       UpperCamelCase_		)
		# change to expected output here
		snake_case__	       =   tf.convert_to_tensor(
		    [[2.3_0_5_0, 2.8_2_7_9, 0.6_5_3_1], [-1.8_4_5_7, -0.1_4_5_5, -3.5_6_6_1], [-1.0_1_8_6, 0.4_5_8_6, -2.2_0_4_3]]      ,       )
		tf.debugging.assert_near(output[:, :3, :3]      ,       UpperCamelCase_      ,       atol=1E-3		)





	def   _snake_case   (		self		)      ->		List[str]:
		snake_case__	       =   TFLEDForConditionalGeneration.from_pretrained('allenai/led-base-16384'		)

		# change to intended input here
		snake_case__	       =   _long_tensor([512 * [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69]]		)
		snake_case__	       =   _long_tensor([128 * [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69]]		)
		snake_case__	       =   prepare_led_inputs_dict(model.config      ,       UpperCamelCase_      ,       UpperCamelCase_		)
		snake_case__	       =   model(**UpperCamelCase_		)[0]
		snake_case__	       =   (1, 1024, model.config.vocab_size)
		self.assertEqual(output.shape      ,       UpperCamelCase_		)
		# change to expected output here
		snake_case__	       =   tf.convert_to_tensor(
		    [[3_3.6_5_0_7, 6.4_5_7_2, 1_6.8_0_8_9], [5.8_7_3_9, -2.4_2_3_8, 1_1.2_9_0_2], [-3.2_1_3_9, -4.3_1_4_9, 4.2_7_8_3]]      ,       )
		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 				__snake_case     (					__magic_name__							):





	@staticmethod
	@abstractmethod
	def   _snake_case   (		UpperCamelCase_		)      ->		Dict:
		raise NotImplementedError()





	@abstractmethod
	def   _snake_case   (		self		)      ->		List[Any]:
		raise NotImplementedError()



