




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






			'''simple docstring'''





			UpperCamelCase__								=LEDConfig
			UpperCamelCase__								={}
			UpperCamelCase__								='''gelu'''

			def __init__( self			:   Tuple       ,						lowerCamelCase_			:   Dict       ,						lowerCamelCase_			:   Dict=13       ,						lowerCamelCase_			:   str=7       ,						lowerCamelCase_			:   Union[str, Any]=True       ,						lowerCamelCase_			:   List[str]=False       ,						lowerCamelCase_			:   Dict=99       ,						lowerCamelCase_			:   Optional[Any]=32       ,						lowerCamelCase_			:   Union[str, Any]=2       ,						lowerCamelCase_			:   Dict=4       ,						lowerCamelCase_			:   List[Any]=37       ,						lowerCamelCase_			:   Union[str, Any]=0.1       ,						lowerCamelCase_			:   List[Any]=0.1       ,						lowerCamelCase_			:   List[Any]=20       ,						lowerCamelCase_			:   Tuple=2       ,						lowerCamelCase_			:   Union[str, Any]=1       ,						lowerCamelCase_			:   Dict=0       ,						lowerCamelCase_			:   Optional[int]=4       ,						)						->   Any:
								__magic_name__							:	Dict    							=							parent
								__magic_name__							:	Optional[int]    							=							batch_size
								__magic_name__							:	Dict    							=							seq_length
								__magic_name__							:	List[str]    							=							is_training
								__magic_name__							:	Optional[int]    							=							use_labels
								__magic_name__							:	Union[str, Any]    							=							vocab_size
								__magic_name__							:	Any    							=							hidden_size
								__magic_name__							:	List[Any]    							=							num_hidden_layers
								__magic_name__							:	List[Any]    							=							num_attention_heads
								__magic_name__							:	Optional[int]    							=							intermediate_size

								__magic_name__							:	List[Any]    							=							hidden_dropout_prob
								__magic_name__							:	Tuple    							=							attention_probs_dropout_prob
								__magic_name__							:	Optional[int]    							=							max_position_embeddings
								__magic_name__							:	Union[str, Any]    							=							eos_token_id
								__magic_name__							:	str    							=							pad_token_id
								__magic_name__							:	Tuple    							=							bos_token_id
								__magic_name__							:	str    							=							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
								__magic_name__							:	Dict    							=							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
								__magic_name__							:	Optional[Any]    							=							(
								    self.seq_length + (self.attention_window - self.seq_length % self.attention_window) % self.attention_window
								)

			def 							UpperCAmelCase__       ( self			:   Any  )						->   Optional[int]:
								__magic_name__							:	Union[str, Any]    							=							ids_tensor([self.batch_size, self.seq_length - 1]       ,						self.vocab_size  )
								__magic_name__							:	Optional[int]    							=							tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size  )       ,						1  )
								__magic_name__							:	Union[str, Any]    							=							tf.concat([input_ids, eos_tensor]       ,						axis=1  )

								__magic_name__							:	Optional[int]    							=							ids_tensor([self.batch_size, self.seq_length]       ,						self.vocab_size  )

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


			def 							UpperCAmelCase__       ( self			:   str       ,						lowerCamelCase_			:   str       ,						lowerCamelCase_			:   str  )						->   Tuple:
								__magic_name__							:	Tuple    							=							TFLEDModel(config=snake_case__  ).get_decoder()
								__magic_name__							:	str    							=							inputs_dict['''input_ids''']

								__magic_name__							:	Optional[Any]    							=							input_ids[:1, :]
								__magic_name__							:	Dict    							=							inputs_dict['''attention_mask'''][:1, :]
								__magic_name__							:	Dict    							=							1

								# first forward pass
								__magic_name__							:	Tuple    							=							model(snake_case__       ,						attention_mask=snake_case__       ,						use_cache=snake_case__  )

								__magic_name__							,    __magic_name__							:	Dict    							=							outputs.to_tuple()

								# create hypothetical next token and extent to next_input_ids
								__magic_name__							:	Union[str, Any]    							=							ids_tensor((self.batch_size, 3)       ,						config.vocab_size  )
								__magic_name__							:	Optional[int]    							=							tf.cast(ids_tensor((self.batch_size, 3)       ,						2  )       ,						tf.inta  )

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

								__magic_name__							:	int    							=							model(snake_case__       ,						attention_mask=snake_case__  )[0]
								__magic_name__							:	Optional[Any]    							=							model(snake_case__       ,						attention_mask=snake_case__       ,						past_key_values=snake_case__  )[0]

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

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

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







def 		lowercase__   (		__A:	Optional[int]							,__A:	str							,__A:	Optional[Any]							,__A:	Union[str, Any]=None							,__A:	Optional[Any]=None							,__A:	Union[str, Any]=None							,__A:	Optional[Any]=None							,):





					'''simple docstring'''






					if attention_mask is None:
										__magic_name__							:	Any    							=							tf.cast(tf.math.not_equal(lowerCAmelCase__							,config.pad_token_id				)							,tf.inta				)
					if decoder_attention_mask is None:
										__magic_name__							:	Optional[Any]    							=							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:
										__magic_name__							:	int    							=							tf.ones((config.encoder_layers, config.encoder_attention_heads)				)
					if decoder_head_mask is None:
										__magic_name__							:	Any    							=							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    lowerCamelCase    (		__a				,__a				,unittest.TestCase						):






			'''simple docstring'''





			UpperCamelCase__								=(TFLEDForConditionalGeneration, TFLEDModel) if is_tf_available() else ()
			UpperCamelCase__								=(TFLEDForConditionalGeneration,) if is_tf_available() else ()
			UpperCamelCase__								=(
			    {
			        '''conversational''': TFLEDForConditionalGeneration,
			        '''feature-extraction''': TFLEDModel,
			        '''summarization''': TFLEDForConditionalGeneration,
			        '''text2text-generation''': TFLEDForConditionalGeneration,
			        '''translation''': TFLEDForConditionalGeneration,
			    }
			    if is_tf_available()
			    else {}
			)
			UpperCamelCase__								=True
			UpperCamelCase__								=False
			UpperCamelCase__								=False
			UpperCamelCase__								=False

			def 							UpperCAmelCase__       ( self			:   Union[str, Any]  )						->   Tuple:
								__magic_name__							:	Any    							=							TFLEDModelTester(self  )
								__magic_name__							:	Tuple    							=							ConfigTester(self       ,						config_class=snake_case__  )

			def 							UpperCAmelCase__       ( self			:   Tuple  )						->   List[str]:
								self.config_tester.run_common_tests()

			def 							UpperCAmelCase__       ( self			:   Any  )						->   Dict:
								__magic_name__							:	Optional[Any]    							=							self.model_tester.prepare_config_and_inputs_for_common()
								self.model_tester.check_decoder_model_past_large_inputs(*snake_case__  )

			def 							UpperCAmelCase__       ( self			:   Any  )						->   Optional[int]:
								__magic_name__							,    __magic_name__							:	str    							=							self.model_tester.prepare_config_and_inputs_for_common()
								__magic_name__							:	str    							=							tf.zeros_like(inputs_dict['''attention_mask''']  )
								__magic_name__							:	Optional[Any]    							=							2
								__magic_name__							:	List[str]    							=							tf.where(
								    tf.range(self.model_tester.seq_length  )[None, :] < num_global_attn_indices       ,						1       ,						inputs_dict['''global_attention_mask''']       ,						)

								__magic_name__							:	List[Any]    							=							True
								__magic_name__							:	int    							=							self.model_tester.seq_length
								__magic_name__							:	Optional[int]    							=							self.model_tester.encoder_seq_length

								def check_decoder_attentions_output(lowerCamelCase_			:   Dict  ):
													__magic_name__							:	Tuple    							=							outputs.decoder_attentions
													self.assertEqual(len(snake_case__  )       ,						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(lowerCamelCase_			:   Dict  ):
													__magic_name__							:	Tuple    							=							[t.numpy() for t in outputs.encoder_attentions]
													__magic_name__							:	List[Any]    							=							[t.numpy() for t in outputs.encoder_global_attentions]
													self.assertEqual(len(snake_case__  )       ,						self.model_tester.num_hidden_layers  )
													self.assertEqual(len(snake_case__  )       ,						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:
													__magic_name__							:	Union[str, Any]    							=							True
													__magic_name__							:	Optional[int]    							=							False
													__magic_name__							:	Tuple    							=							False
													__magic_name__							:	List[str]    							=							model_class(snake_case__  )
													__magic_name__							:	List[Any]    							=							model(self._prepare_for_class(snake_case__       ,						snake_case__  )  )
													__magic_name__							:	Any    							=							len(snake_case__  )
													self.assertEqual(config.output_hidden_states       ,						snake_case__  )
													check_encoder_attentions_output(snake_case__  )

													if self.is_encoder_decoder:
																		__magic_name__							:	Optional[Any]    							=							model_class(snake_case__  )
																		__magic_name__							:	List[str]    							=							model(self._prepare_for_class(snake_case__       ,						snake_case__  )  )
																		self.assertEqual(config.output_hidden_states       ,						snake_case__  )
																		check_decoder_attentions_output(snake_case__  )

													# Check that output attentions can also be changed via the config
													del inputs_dict["output_attentions"]
													__magic_name__							:	str    							=							True
													__magic_name__							:	int    							=							model_class(snake_case__  )
													__magic_name__							:	Union[str, Any]    							=							model(self._prepare_for_class(snake_case__       ,						snake_case__  )  )
													self.assertEqual(config.output_hidden_states       ,						snake_case__  )
													check_encoder_attentions_output(snake_case__  )

													# Check attention is always last and order is fine
													__magic_name__							:	Tuple    							=							True
													__magic_name__							:	Tuple    							=							True
													__magic_name__							:	str    							=							model_class(snake_case__  )
													__magic_name__							:	Optional[Any]    							=							model(self._prepare_for_class(snake_case__       ,						snake_case__  )  )

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

			@unittest.skip('''LED keeps using potentially symbolic tensors in conditionals and breaks tracing.'''  )
			def 							UpperCAmelCase__       ( self			:   List[str]  )						->   Union[str, Any]:
								pass


			def 							UpperCAmelCase__       ( self			:   Dict  )						->   List[Any]:
								# TODO: Head-masking not yet implement
								pass







def 		lowercase__   (		__A:	str				):





					'''simple docstring'''






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


__lowerCamelCase   :						Union[str, Any]				=			1E-4


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






			'''simple docstring'''





			def 							UpperCAmelCase__       ( self			:   Optional[int]  )						->   Tuple:
								__magic_name__							:	Optional[Any]    							=							TFLEDForConditionalGeneration.from_pretrained('''allenai/led-base-16384'''  ).led

								# change to intended input here
								__magic_name__							:	Any    							=							_long_tensor([512 * [0, 31414, 232, 328, 740, 1140, 12695, 69]]  )
								__magic_name__							:	str    							=							_long_tensor([128 * [0, 31414, 232, 328, 740, 1140, 12695, 69]]  )
								__magic_name__							:	List[str]    							=							prepare_led_inputs_dict(model.config       ,						snake_case__       ,						snake_case__  )
								__magic_name__							:	Optional[Any]    							=							model(**snake_case__  )[0]
								__magic_name__							:	List[Any]    							=							(1, 1024, 768)
								self.assertEqual(output.shape       ,						snake_case__  )
								# change to expected output here
								__magic_name__							:	str    							=							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]       ,						snake_case__       ,						atol=1E-3  )


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

								# change to intended input here
								__magic_name__							:	str    							=							_long_tensor([512 * [0, 31414, 232, 328, 740, 1140, 12695, 69]]  )
								__magic_name__							:	Optional[Any]    							=							_long_tensor([128 * [0, 31414, 232, 328, 740, 1140, 12695, 69]]  )
								__magic_name__							:	Tuple    							=							prepare_led_inputs_dict(model.config       ,						snake_case__       ,						snake_case__  )
								__magic_name__							:	List[str]    							=							model(**snake_case__  )[0]
								__magic_name__							:	Union[str, Any]    							=							(1, 1024, model.config.vocab_size)
								self.assertEqual(output.shape       ,						snake_case__  )
								# change to expected output here
								__magic_name__							:	Union[str, Any]    							=							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]       ,						snake_case__       ,						atol=1E-3       ,						rtol=1E-3  )









from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_speech_available,
    is_tf_available,
    is_torch_available,
)


__lowerCamelCase   :						List[Any]				=			{
    '''configuration_speech_to_text''': ['''SPEECH_TO_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''Speech2TextConfig'''],
    '''processing_speech_to_text''': ['''Speech2TextProcessor'''],
}

try:
	if not is_sentencepiece_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__lowerCamelCase   :						Any				=			['''Speech2TextTokenizer''']

try:
	if not is_speech_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__lowerCamelCase   :						Optional[Any]				=			['''Speech2TextFeatureExtractor''']

try:
	if not is_tf_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__lowerCamelCase   :						Any				=			[
	    '''TF_SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST''',
	    '''TFSpeech2TextForConditionalGeneration''',
	    '''TFSpeech2TextModel''',
	    '''TFSpeech2TextPreTrainedModel''',
	]

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__lowerCamelCase   :						Any				=			[
	    '''SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST''',
	    '''Speech2TextForConditionalGeneration''',
	    '''Speech2TextModel''',
	    '''Speech2TextPreTrainedModel''',
	]


if TYPE_CHECKING:
	from .configuration_speech_to_text import SPEECH_TO_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP, SpeechaTextConfig
	from .processing_speech_to_text import SpeechaTextProcessor

	try:
		if not is_sentencepiece_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .tokenization_speech_to_text import SpeechaTextTokenizer

	try:
		if not is_speech_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .feature_extraction_speech_to_text import SpeechaTextFeatureExtractor

	try:
		if not is_tf_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_tf_speech_to_text import (
		    TF_SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFSpeechaTextForConditionalGeneration,
		    TFSpeechaTextModel,
		    TFSpeechaTextPreTrainedModel,
		)

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_speech_to_text import (
		    SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST,
		    SpeechaTextForConditionalGeneration,
		    SpeechaTextModel,
		    SpeechaTextPreTrainedModel,
		)

else:
	import sys

	__lowerCamelCase   :						List[str]				=			_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)



