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 						a__							:
			A__    :       Any        =      LEDConfig
			A__    :       Optional[Any]        =      {}
			A__    :       Union[str, Any]        =      'gelu'



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

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



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

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

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







			def __SCREAMING_SNAKE_CASE							(						self				,						UpperCAmelCase				,						UpperCAmelCase   )			->   Tuple:
								__a		 =       TFLEDModel(config=UpperCAmelCase   ).get_decoder()
								__a		 =       inputs_dict['input_ids']

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

								# first forward pass
								__a		 =       model(UpperCAmelCase				,						attention_mask=UpperCAmelCase				,						use_cache=UpperCAmelCase   )

								__a ,				__a		 =       outputs.to_tuple()

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

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

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

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






def 							lowerCAmelCase(     __lowerCamelCase				,    __lowerCamelCase				,    __lowerCamelCase				,    __lowerCamelCase=None				,    __lowerCamelCase=None				,    __lowerCamelCase=None				,    __lowerCamelCase=None				,    ):
					if attention_mask is None:
										__a		 =       tf.cast(tf.math.not_equal(__lowerCamelCase				,    config.pad_token_id						)				,    tf.inta						)
					if decoder_attention_mask is None:
										__a		 =       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:
										__a		 =       tf.ones((config.encoder_layers, config.encoder_attention_heads)						)
					if decoder_head_mask is None:
										__a		 =       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 						a__							(       __snake_case				,						__snake_case				,						unittest.TestCase		):
			A__    :       Dict        =      (TFLEDForConditionalGeneration, TFLEDModel) if is_tf_available() else ()
			A__    :       Tuple        =      (TFLEDForConditionalGeneration,) if is_tf_available() else ()
			A__    :       List[str]        =      (
			    {
			        'conversational': TFLEDForConditionalGeneration,
			        'feature-extraction': TFLEDModel,
			        'summarization': TFLEDForConditionalGeneration,
			        'text2text-generation': TFLEDForConditionalGeneration,
			        'translation': TFLEDForConditionalGeneration,
			    }
			    if is_tf_available()
			    else {}
			)
			A__    :       List[str]        =      True
			A__    :       Union[str, Any]        =      False
			A__    :       Any        =      False
			A__    :       Any        =      False



			def __SCREAMING_SNAKE_CASE							(						self   )			->   Tuple:
								__a		 =       TFLEDModelTester(self   )
								__a		 =       ConfigTester(self				,						config_class=UpperCAmelCase   )



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



			def __SCREAMING_SNAKE_CASE							(						self   )			->   List[str]:
								__a		 =       self.model_tester.prepare_config_and_inputs_for_common()
								self.model_tester.check_decoder_model_past_large_inputs(*UpperCAmelCase   )



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

								__a		 =       True
								__a		 =       self.model_tester.seq_length
								__a		 =       self.model_tester.encoder_seq_length

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

													if self.is_encoder_decoder:
																		__a		 =       model_class(UpperCAmelCase   )
																		__a		 =       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"]
													__a		 =       True
													__a		 =       model_class(UpperCAmelCase   )
													__a		 =       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
													__a		 =       True
													__a		 =       True
													__a		 =       model_class(UpperCAmelCase   )
													__a		 =       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 __SCREAMING_SNAKE_CASE							(						self   )			->   str:
								pass







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






def 							lowerCAmelCase(     __lowerCamelCase						):
					return tf.constant(__lowerCamelCase				,    dtype=tf.intaa						)


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



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



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

								# change to intended input here
								__a		 =       _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]]   )
								__a		 =       _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]]   )
								__a		 =       prepare_led_inputs_dict(model.config				,						UpperCAmelCase				,						UpperCAmelCase   )
								__a		 =       model(**UpperCAmelCase   )[0]
								__a		 =       (1, 1_0_2_4, 7_6_8)
								self.assertEqual(output.shape				,						UpperCAmelCase   )
								# change to expected output here
								__a		 =       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 __SCREAMING_SNAKE_CASE							(						self   )			->   Union[str, Any]:
								__a		 =       TFLEDForConditionalGeneration.from_pretrained('allenai/led-base-16384'   )

								# change to intended input here
								__a		 =       _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]]   )
								__a		 =       _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]]   )
								__a		 =       prepare_led_inputs_dict(model.config				,						UpperCAmelCase				,						UpperCAmelCase   )
								__a		 =       model(**UpperCAmelCase   )[0]
								__a		 =       (1, 1_0_2_4, model.config.vocab_size)
								self.assertEqual(output.shape				,						UpperCAmelCase   )
								# change to expected output here
								__a		 =       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   )





import warnings
from contextlib import contextmanager

from ...processing_utils import ProcessorMixin
from .feature_extraction_wavaveca import WavaVecaFeatureExtractor
from .tokenization_wavaveca import WavaVecaCTCTokenizer


class 						a__							(       __snake_case		):
			A__    :       Any        =      'Wav2Vec2FeatureExtractor'
			A__    :       str        =      'AutoTokenizer'



			def __init__(						self				,						UpperCAmelCase				,						UpperCAmelCase   )			->   Union[str, Any]:
								super().__init__(UpperCAmelCase				,						UpperCAmelCase   )
								__a		 =       self.feature_extractor
								__a		 =       False



			@classmethod
			def __SCREAMING_SNAKE_CASE							(						cls				,						UpperCAmelCase				,						**UpperCAmelCase   )			->   Dict:
								try:
													return super().from_pretrained(UpperCAmelCase				,						**UpperCAmelCase   )
								except OSError:
													warnings.warn(
													    f'''Loading a tokenizer inside {cls.__name__} from a config that does not'''
													    ' include a `tokenizer_class` attribute is deprecated and will be '
													    'removed in v5. Please add `\'tokenizer_class\': \'Wav2Vec2CTCTokenizer\'`'
													    ' attribute to either your `config.json` or `tokenizer_config.json` '
													    'file to suppress this warning: '				,						UpperCAmelCase				,						)

													__a		 =       WavaVecaFeatureExtractor.from_pretrained(UpperCAmelCase				,						**UpperCAmelCase   )
													__a		 =       WavaVecaCTCTokenizer.from_pretrained(UpperCAmelCase				,						**UpperCAmelCase   )

													return cls(feature_extractor=UpperCAmelCase				,						tokenizer=UpperCAmelCase   )



			def __call__(						self				,						*UpperCAmelCase				,						**UpperCAmelCase   )			->   List[Any]:
								# For backward compatibility
								if self._in_target_context_manager:
													return self.current_processor(*UpperCAmelCase				,						**UpperCAmelCase   )

								if "raw_speech" in kwargs:
													warnings.warn('Using `raw_speech` as a keyword argument is deprecated. Use `audio` instead.'   )
													__a		 =       kwargs.pop('raw_speech'   )
								else:
													__a		 =       kwargs.pop('audio'				,						UpperCAmelCase   )
								__a		 =       kwargs.pop('sampling_rate'				,						UpperCAmelCase   )
								__a		 =       kwargs.pop('text'				,						UpperCAmelCase   )
								if len(UpperCAmelCase   ) > 0:
													__a		 =       args[0]
													__a		 =       args[1:]

								if audio is None and text is None:
													raise ValueError('You need to specify either an `audio` or `text` input to process.'   )

								if audio is not None:
													__a		 =       self.feature_extractor(UpperCAmelCase				,						*UpperCAmelCase				,						sampling_rate=UpperCAmelCase				,						**UpperCAmelCase   )
								if text is not None:
													__a		 =       self.tokenizer(UpperCAmelCase				,						**UpperCAmelCase   )

								if text is None:
													return inputs
								elif audio is None:
													return encodings
								else:
													__a		 =       encodings['input_ids']
													return inputs



			def __SCREAMING_SNAKE_CASE							(						self				,						*UpperCAmelCase				,						**UpperCAmelCase   )			->   Tuple:
								# For backward compatibility
								if self._in_target_context_manager:
													return self.current_processor.pad(*UpperCAmelCase				,						**UpperCAmelCase   )

								__a		 =       kwargs.pop('input_features'				,						UpperCAmelCase   )
								__a		 =       kwargs.pop('labels'				,						UpperCAmelCase   )
								if len(UpperCAmelCase   ) > 0:
													__a		 =       args[0]
													__a		 =       args[1:]

								if input_features is not None:
													__a		 =       self.feature_extractor.pad(UpperCAmelCase				,						*UpperCAmelCase				,						**UpperCAmelCase   )
								if labels is not None:
													__a		 =       self.tokenizer.pad(UpperCAmelCase				,						**UpperCAmelCase   )

								if labels is None:
													return input_features
								elif input_features is None:
													return labels
								else:
													__a		 =       labels['input_ids']
													return input_features



			def __SCREAMING_SNAKE_CASE							(						self				,						*UpperCAmelCase				,						**UpperCAmelCase   )			->   List[str]:
								return self.tokenizer.batch_decode(*UpperCAmelCase				,						**UpperCAmelCase   )



			def __SCREAMING_SNAKE_CASE							(						self				,						*UpperCAmelCase				,						**UpperCAmelCase   )			->   Union[str, Any]:
								return self.tokenizer.decode(*UpperCAmelCase				,						**UpperCAmelCase   )







			@contextmanager
			def __SCREAMING_SNAKE_CASE							(						self   )			->   List[Any]:
								warnings.warn(
								    '`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your '
								    'labels by using the argument `text` of the regular `__call__` method (either in the same call as '
								    'your audio inputs, or in a separate call.'   )
								__a		 =       True
								__a		 =       self.tokenizer
								yield
								__a		 =       self.feature_extractor
								__a		 =       False




