






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




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

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




							def 						__snake_case					(   self     :   Tuple	):
											lowerCAmelCase__      			=						ids_tensor([self.batch_size, self.seq_length - 1] ,							self.vocab_size	)
											lowerCAmelCase__      			=						tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size	) ,							1	)
											lowerCAmelCase__      			=						tf.concat([input_ids, eos_tensor] ,							axis=1	)

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

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




							def 						__snake_case					(   self     :   List[Any] ,							SCREAMING_SNAKE_CASE_     :   Any ,							SCREAMING_SNAKE_CASE_     :   str	):
											lowerCAmelCase__      			=						TFLEDModel(config=SCREAMING_SNAKE_CASE_	).get_decoder()
											lowerCAmelCase__      			=						inputs_dict['''input_ids''']

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

											# first forward pass
											lowerCAmelCase__      			=						model(SCREAMING_SNAKE_CASE_ ,							attention_mask=SCREAMING_SNAKE_CASE_ ,							use_cache=SCREAMING_SNAKE_CASE_	)

											lowerCAmelCase__			,		lowerCAmelCase__      			=						outputs.to_tuple()

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

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

											lowerCAmelCase__      			=						model(SCREAMING_SNAKE_CASE_ ,							attention_mask=SCREAMING_SNAKE_CASE_	)[0]
											lowerCAmelCase__      			=						model(SCREAMING_SNAKE_CASE_ ,							attention_mask=SCREAMING_SNAKE_CASE_ ,							past_key_values=SCREAMING_SNAKE_CASE_	)[0]

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

											# select random slice
											lowerCAmelCase__      			=						int(ids_tensor((1,) ,							output_from_past.shape[-1]	)	)
											lowerCAmelCase__      			=						output_from_no_past[:, -3:, random_slice_idx]
											lowerCAmelCase__      			=						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_		(lowercase__						:       Any     ,					lowercase__						:       Optional[int]     ,					lowercase__						:       Any     ,					lowercase__						:       Tuple=None     ,					lowercase__						:       Optional[Any]=None     ,					lowercase__						:       Optional[Any]=None     ,					lowercase__						:       Tuple=None     ,					)		->    Dict:


				'''simple docstring'''




				if attention_mask is None:
								lowerCAmelCase__      			=						tf.cast(tf.math.not_equal(lowercase__     ,					config.pad_token_id       )     ,					tf.inta       )
				if decoder_attention_mask is None:
								lowerCAmelCase__      			=						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:
								lowerCAmelCase__      			=						tf.ones((config.encoder_layers, config.encoder_attention_heads)       )
				if decoder_head_mask is None:
								lowerCAmelCase__      			=						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_       (				snake_case__     ,    snake_case__     ,    unittest.TestCase					):
							UpperCamelCase_		:Tuple		      =    (TFLEDForConditionalGeneration, TFLEDModel) if is_tf_available() else ()
							UpperCamelCase_		:str		      =    (TFLEDForConditionalGeneration,) if is_tf_available() else ()
							UpperCamelCase_		:Optional[int]		      =    (
							    {
							        'conversational': TFLEDForConditionalGeneration,
							        'feature-extraction': TFLEDModel,
							        'summarization': TFLEDForConditionalGeneration,
							        'text2text-generation': TFLEDForConditionalGeneration,
							        'translation': TFLEDForConditionalGeneration,
							    }
							    if is_tf_available()
							    else {}
							)
							UpperCamelCase_		:str		      =    True
							UpperCamelCase_		:Union[str, Any]		      =    False
							UpperCamelCase_		:int		      =    False
							UpperCamelCase_		:List[str]		      =    False




							def 						__snake_case					(   self     :   str	):
											lowerCAmelCase__      			=						TFLEDModelTester(self	)
											lowerCAmelCase__      			=						ConfigTester(self ,							config_class=SCREAMING_SNAKE_CASE_	)




							def 						__snake_case					(   self     :   Union[str, Any]	):
											self.config_tester.run_common_tests()




							def 						__snake_case					(   self     :   Optional[Any]	):
											lowerCAmelCase__      			=						self.model_tester.prepare_config_and_inputs_for_common()
											self.model_tester.check_decoder_model_past_large_inputs(*SCREAMING_SNAKE_CASE_	)




							def 						__snake_case					(   self     :   Any	):
											lowerCAmelCase__			,		lowerCAmelCase__      			=						self.model_tester.prepare_config_and_inputs_for_common()
											lowerCAmelCase__      			=						tf.zeros_like(inputs_dict['''attention_mask''']	)
											lowerCAmelCase__      			=						2
											lowerCAmelCase__      			=						tf.where(
											    tf.range(self.model_tester.seq_length	)[None, :] < num_global_attn_indices ,							1 ,							inputs_dict['''global_attention_mask'''] ,							)

											lowerCAmelCase__      			=						True
											lowerCAmelCase__      			=						self.model_tester.seq_length
											lowerCAmelCase__      			=						self.model_tester.encoder_seq_length

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

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

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

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

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




							@unittest.skip('''LED keeps using potentially symbolic tensors in conditionals and breaks tracing.'''	)
							def 						__snake_case					(   self     :   List[str]	):
											pass




							def 						__snake_case					(   self     :   Tuple	):
											# TODO: Head-masking not yet implement
											pass





def 	lowerCAmelCase_		(lowercase__						:       Dict       )		->    Any:


				'''simple docstring'''




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


_UpperCAmelCase			:		int   				=    1E-4




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




							def 						__snake_case					(   self     :   Dict	):
											lowerCAmelCase__      			=						TFLEDForConditionalGeneration.from_pretrained('''allenai/led-base-16384'''	).led

											# change to intended input here
											lowerCAmelCase__      			=						_long_tensor([512 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]	)
											lowerCAmelCase__      			=						_long_tensor([128 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]	)
											lowerCAmelCase__      			=						prepare_led_inputs_dict(model.config ,							SCREAMING_SNAKE_CASE_ ,							SCREAMING_SNAKE_CASE_	)
											lowerCAmelCase__      			=						model(**SCREAMING_SNAKE_CASE_	)[0]
											lowerCAmelCase__      			=						(1, 1_024, 768)
											self.assertEqual(output.shape ,							SCREAMING_SNAKE_CASE_	)
											# change to expected output here
											lowerCAmelCase__      			=						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] ,							SCREAMING_SNAKE_CASE_ ,							atol=1e-3	)




							def 						__snake_case					(   self     :   Dict	):
											lowerCAmelCase__      			=						TFLEDForConditionalGeneration.from_pretrained('''allenai/led-base-16384'''	)

											# change to intended input here
											lowerCAmelCase__      			=						_long_tensor([512 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]	)
											lowerCAmelCase__      			=						_long_tensor([128 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]	)
											lowerCAmelCase__      			=						prepare_led_inputs_dict(model.config ,							SCREAMING_SNAKE_CASE_ ,							SCREAMING_SNAKE_CASE_	)
											lowerCAmelCase__      			=						model(**SCREAMING_SNAKE_CASE_	)[0]
											lowerCAmelCase__      			=						(1, 1_024, model.config.vocab_size)
											self.assertEqual(output.shape ,							SCREAMING_SNAKE_CASE_	)
											# change to expected output here
											lowerCAmelCase__      			=						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] ,							SCREAMING_SNAKE_CASE_ ,							atol=1e-3 ,							rtol=1e-3	)













import pytest

import datasets


# Import fixture modules as plugins
_UpperCAmelCase			:		Tuple   				=    ["tests.fixtures.files", "tests.fixtures.hub", "tests.fixtures.fsspec"]





def 	lowerCAmelCase_		(lowercase__						:       Optional[int]     ,					lowercase__						:       int       )		->    Any:


				'''simple docstring'''




				for item in items:
								if any(marker in item.keywords for marker in ['''integration''', '''unit''']       ):
												continue
								item.add_marker(pytest.mark.unit       )





def 	lowerCAmelCase_		(lowercase__						:       str       )		->    List[str]:


				'''simple docstring'''




				config.addinivalue_line('''markers'''     ,					'''torchaudio_latest: mark test to run with torchaudio>=0.12'''       )





@pytest.fixture(autouse=lowercase__       )
def 	lowerCAmelCase_		(lowercase__						:       int     ,					lowercase__						:       Dict       )		->    Optional[Any]:


				'''simple docstring'''




				lowerCAmelCase__      			=						tmp_path_factory.getbasetemp() / '''cache'''
				lowerCAmelCase__      			=						test_hf_cache_home / '''datasets'''
				lowerCAmelCase__      			=						test_hf_cache_home / '''metrics'''
				lowerCAmelCase__      			=						test_hf_cache_home / '''modules'''
				monkeypatch.setattr('''datasets.config.HF_DATASETS_CACHE'''     ,					str(lowercase__       )       )
				monkeypatch.setattr('''datasets.config.HF_METRICS_CACHE'''     ,					str(lowercase__       )       )
				monkeypatch.setattr('''datasets.config.HF_MODULES_CACHE'''     ,					str(lowercase__       )       )
				lowerCAmelCase__      			=						test_hf_datasets_cache / '''downloads'''
				monkeypatch.setattr('''datasets.config.DOWNLOADED_DATASETS_PATH'''     ,					str(lowercase__       )       )
				lowerCAmelCase__      			=						test_hf_datasets_cache / '''downloads''' / '''extracted'''
				monkeypatch.setattr('''datasets.config.EXTRACTED_DATASETS_PATH'''     ,					str(lowercase__       )       )





@pytest.fixture(autouse=lowercase__     ,					scope='''session'''       )
def 	lowerCAmelCase_		()		->    int:


				'''simple docstring'''




				datasets.disable_progress_bar()





@pytest.fixture(autouse=lowercase__       )
def 	lowerCAmelCase_		(lowercase__						:       Dict       )		->    List[str]:


				'''simple docstring'''




				monkeypatch.setattr('''datasets.config.HF_UPDATE_DOWNLOAD_COUNTS'''     ,					lowercase__       )





@pytest.fixture
def 	lowerCAmelCase_		(lowercase__						:       List[str]       )		->    List[Any]:


				'''simple docstring'''




				monkeypatch.setattr('''sqlalchemy.util.deprecations.SILENCE_UBER_WARNING'''     ,					lowercase__       )





