

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 	__magic_name__    :
			SCREAMING_SNAKE_CASE							   =						LEDConfig
			SCREAMING_SNAKE_CASE							   =						{}
			SCREAMING_SNAKE_CASE							   =						'gelu'





			def __init__(		self					,  __snake_case					,  __snake_case=13					,  __snake_case=7					,  __snake_case=True					,  __snake_case=False					,  __snake_case=99					,  __snake_case=32					,  __snake_case=2					,  __snake_case=4					,  __snake_case=37					,  __snake_case=0.1					,  __snake_case=0.1					,  __snake_case=20					,  __snake_case=2					,  __snake_case=1					,  __snake_case=0					,  __snake_case=4					,  )   -> Optional[int]:




						'''simple docstring'''
						__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 		__magic_name__    (		self					)   -> str:




						'''simple docstring'''
						__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(__snake_case					,  __snake_case					,  __snake_case					)
						__a        =tf.concat(
						    [tf.zeros_like(__snake_case					)[:, :-1], tf.ones_like(__snake_case					)[:, -1:]]					,  axis=-1					,  )
						__a        =global_attention_mask
						return config, inputs_dict





			def 		__magic_name__    (		self					,  __snake_case					,  __snake_case					)   -> int:




						'''simple docstring'''
						__a        =TFLEDModel(config=__snake_case					).get_decoder()
						__a        =inputs_dict['input_ids']

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

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

						__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(__snake_case					,  attention_mask=__snake_case					)[0]
						__a        =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
						__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(__snake_case					,  __snake_case					,  rtol=1e-3					)






def        UpperCamelCase_(							_snake_case   :  Dict       ,     _snake_case   :  Any       ,     _snake_case   :  Union[str, Any]       ,     _snake_case   :  List[str]=None       ,     _snake_case   :  List[str]=None       ,     _snake_case   :  Optional[Any]=None       ,     _snake_case   :  Any=None       ,     ):




			"""simple docstring"""






			if attention_mask is None:
						__a        =tf.cast(tf.math.not_equal(_snake_case       ,     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 	__magic_name__    (			lowerCAmelCase_  ,    lowerCAmelCase_  ,    unittest.TestCase      ):
			SCREAMING_SNAKE_CASE							   =						(TFLEDForConditionalGeneration, TFLEDModel) if is_tf_available() else ()
			SCREAMING_SNAKE_CASE							   =						(TFLEDForConditionalGeneration,) if is_tf_available() else ()
			SCREAMING_SNAKE_CASE							   =						(
			    {
			        'conversational': TFLEDForConditionalGeneration,
			        'feature-extraction': TFLEDModel,
			        'summarization': TFLEDForConditionalGeneration,
			        'text2text-generation': TFLEDForConditionalGeneration,
			        'translation': TFLEDForConditionalGeneration,
			    }
			    if is_tf_available()
			    else {}
			)
			SCREAMING_SNAKE_CASE							   =						True
			SCREAMING_SNAKE_CASE							   =						False
			SCREAMING_SNAKE_CASE							   =						False
			SCREAMING_SNAKE_CASE							   =						False





			def 		__magic_name__    (		self					)   -> Optional[Any]:




						'''simple docstring'''
						__a        =TFLEDModelTester(self					)
						__a        =ConfigTester(self					,  config_class=__snake_case					)





			def 		__magic_name__    (		self					)   -> Optional[int]:




						'''simple docstring'''
						self.config_tester.run_common_tests()





			def 		__magic_name__    (		self					)   -> Tuple:




						'''simple docstring'''
						__a        =self.model_tester.prepare_config_and_inputs_for_common()
						self.model_tester.check_decoder_model_past_large_inputs(*__snake_case					)





			def 		__magic_name__    (		self					)   -> Any:




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

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




						'''simple docstring'''
						pass





			def 		__magic_name__    (		self					)   -> List[Any]:




						'''simple docstring'''
						# TODO: Head-masking not yet implement
						pass






def        UpperCamelCase_(							_snake_case   :  Dict	):




			"""simple docstring"""






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


_lowerCAmelCase					:			Optional[int]  				=			1E-4






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





			def 		__magic_name__    (		self					)   -> Optional[Any]:




						'''simple docstring'''
						__a        =TFLEDForConditionalGeneration.from_pretrained('allenai/led-base-16384'					).led

						# change to intended input here
						__a        =_long_tensor([512 * [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69]]					)
						__a        =_long_tensor([128 * [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69]]					)
						__a        =prepare_led_inputs_dict(model.config					,  __snake_case					,  __snake_case					)
						__a        =model(**__snake_case					)[0]
						__a        =(1, 1024, 768)
						self.assertEqual(output.shape					,  __snake_case					)
						# change to expected output here
						__a        =tf.convert_to_tensor(
						    [[2.3050, 2.8279, 0.6531], [-1.8457, -0.1455, -3.5661], [-1.0186, 0.4586, -2.2043]]					,  )
						tf.debugging.assert_near(output[:, :3, :3]					,  __snake_case					,  atol=1e-3					)





			def 		__magic_name__    (		self					)   -> Tuple:




						'''simple docstring'''
						__a        =TFLEDForConditionalGeneration.from_pretrained('allenai/led-base-16384'					)

						# change to intended input here
						__a        =_long_tensor([512 * [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69]]					)
						__a        =_long_tensor([128 * [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69]]					)
						__a        =prepare_led_inputs_dict(model.config					,  __snake_case					,  __snake_case					)
						__a        =model(**__snake_case					)[0]
						__a        =(1, 1024, model.config.vocab_size)
						self.assertEqual(output.shape					,  __snake_case					)
						# change to expected output here
						__a        =tf.convert_to_tensor(
						    [[33.6507, 6.4572, 16.8089], [5.8739, -2.4238, 11.2902], [-3.2139, -4.3149, 4.2783]]					,  )
						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_tokenizers_available,
    is_torch_available,
)


_lowerCAmelCase					:			Tuple  				=			{
    "configuration_llama": ["LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP", "LlamaConfig"],
}

try:
						if not is_sentencepiece_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_lowerCAmelCase					:			Optional[int]  				=			["LlamaTokenizer"]

try:
						if not is_tokenizers_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_lowerCAmelCase					:			Optional[int]  				=			["LlamaTokenizerFast"]

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_lowerCAmelCase					:			List[Any]  				=			[
						    "LlamaForCausalLM",
						    "LlamaModel",
						    "LlamaPreTrainedModel",
						    "LlamaForSequenceClassification",
						]


if TYPE_CHECKING:
						from .configuration_llama import LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP, LlamaConfig

						try:
												if not is_sentencepiece_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .tokenization_llama import LlamaTokenizer

						try:
												if not is_tokenizers_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .tokenization_llama_fast import LlamaTokenizerFast

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_llama import LlamaForCausalLM, LlamaForSequenceClassification, LlamaModel, LlamaPreTrainedModel


else:
						import sys

						_lowerCAmelCase					:			int  				=			_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)






