
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__    :							str									=		LEDConfig
					UpperCAmelCase__    :							Any									=		{}
					UpperCAmelCase__    :							Any									=		"""gelu"""





					def __init__(self				:   Tuple    ,   _A				:   Union[str, Any]    ,   _A				:   List[str]=1_3    ,   _A				:   Union[str, Any]=7    ,   _A				:   Optional[int]=True    ,   _A				:   Dict=False    ,   _A				:   Dict=9_9    ,   _A				:   List[str]=3_2    ,   _A				:   List[Any]=2    ,   _A				:   Any=4    ,   _A				:   Any=3_7    ,   _A				:   int=0.1    ,   _A				:   Tuple=0.1    ,   _A				:   Optional[int]=2_0    ,   _A				:   str=2    ,   _A				:   Union[str, Any]=1    ,   _A				:   Dict=0    ,   _A				:   Optional[int]=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 		UpperCAmelCase(self				:   Any       )						->						List[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(__a    ,   __a    ,   __a       )
									snake_case								= tf.concat(
									    [tf.zeros_like(__a       )[:, :-1], tf.ones_like(__a       )[:, -1:]]    ,   axis=-1    ,   )
									snake_case								= global_attention_mask
									return config, inputs_dict





					def 		UpperCAmelCase(self				:   List[str]    ,   _A				:   str    ,   _A				:   Optional[int]       )						->						List[Any]:
									snake_case								= TFLEDModel(config=__a       ).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(__a    ,   attention_mask=__a    ,   use_cache=__a       )

									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(__a    ,   attention_mask=__a       )[0]
									snake_case								= model(__a    ,   attention_mask=__a    ,   past_key_values=__a       )[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(__a    ,   __a    ,   rtol=1E-3       )



def 						lowercase_      (       A__			,		A__			,		A__			,		A__=None			,		A__=None			,		A__=None			,		A__=None			,		)							->							str:
				"""simple docstring"""

				if attention_mask is None:
								snake_case								= tf.cast(tf.math.not_equal(snake_case_			,		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 					lowerCamelCase		(  UpperCamelCase__							,			UpperCamelCase__							,			unittest.TestCase  ):
					UpperCAmelCase__    :							List[Any]									=		(TFLEDForConditionalGeneration, TFLEDModel) if is_tf_available() else ()
					UpperCAmelCase__    :							Any									=		(TFLEDForConditionalGeneration,) if is_tf_available() else ()
					UpperCAmelCase__    :							Dict									=		(
					    {
					        """conversational""": TFLEDForConditionalGeneration,
					        """feature-extraction""": TFLEDModel,
					        """summarization""": TFLEDForConditionalGeneration,
					        """text2text-generation""": TFLEDForConditionalGeneration,
					        """translation""": TFLEDForConditionalGeneration,
					    }
					    if is_tf_available()
					    else {}
					)
					UpperCAmelCase__    :							Any									=		True
					UpperCAmelCase__    :							str									=		False
					UpperCAmelCase__    :							Optional[Any]									=		False
					UpperCAmelCase__    :							int									=		False





					def 		UpperCAmelCase(self				:   Union[str, Any]       )						->						Tuple:
									snake_case								= TFLEDModelTester(self       )
									snake_case								= ConfigTester(self    ,   config_class=__a       )





					def 		UpperCAmelCase(self				:   str       )						->						Any:
									self.config_tester.run_common_tests()





					def 		UpperCAmelCase(self				:   Union[str, Any]       )						->						Union[str, Any]:
									snake_case								= self.model_tester.prepare_config_and_inputs_for_common()
									self.model_tester.check_decoder_model_past_large_inputs(*__a       )





					def 		UpperCAmelCase(self				:   Optional[int]       )						->						Any:
									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(_A				:   Optional[int]       ):
													snake_case								= outputs.decoder_attentions
													self.assertEqual(len(__a       )    ,   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(_A				:   Dict       ):
													snake_case								= [t.numpy() for t in outputs.encoder_attentions]
													snake_case								= [t.numpy() for t in outputs.encoder_global_attentions]
													self.assertEqual(len(__a       )    ,   self.model_tester.num_hidden_layers       )
													self.assertEqual(len(__a       )    ,   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(__a       )
													snake_case								= model(self._prepare_for_class(__a    ,   __a       )       )
													snake_case								= len(__a       )
													self.assertEqual(config.output_hidden_states    ,   __a       )
													check_encoder_attentions_output(__a       )

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

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

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

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





					@unittest.skip("LED keeps using potentially symbolic tensors in conditionals and breaks tracing."       )
					def 		UpperCAmelCase(self				:   str       )						->						Tuple:
									pass





					def 		UpperCAmelCase(self				:   Union[str, Any]       )						->						List[str]:
									# TODO: Head-masking not yet implement
									pass



def 						lowercase_      (       A__		)							->							str:
				"""simple docstring"""

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


_A   =			1e-4







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





					def 		UpperCAmelCase(self				:   str       )						->						Optional[Any]:
									snake_case								= TFLEDForConditionalGeneration.from_pretrained("allenai/led-base-16384"       ).led

									# change to intended input here
									snake_case								= _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]]       )
									snake_case								= _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]]       )
									snake_case								= prepare_led_inputs_dict(model.config    ,   __a    ,   __a       )
									snake_case								= model(**__a       )[0]
									snake_case								= (1, 1_0_2_4, 7_6_8)
									self.assertEqual(output.shape    ,   __a       )
									# change to expected output here
									snake_case								= tf.convert_to_tensor(
									    [[2.30_50, 2.82_79, 0.65_31], [-1.84_57, -0.14_55, -3.56_61], [-1.01_86, 0.45_86, -2.20_43]]    ,   )
									tf.debugging.assert_near(output[:, :3, :3]    ,   __a    ,   atol=1E-3       )





					def 		UpperCAmelCase(self				:   Any       )						->						int:
									snake_case								= TFLEDForConditionalGeneration.from_pretrained("allenai/led-base-16384"       )

									# change to intended input here
									snake_case								= _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]]       )
									snake_case								= _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]]       )
									snake_case								= prepare_led_inputs_dict(model.config    ,   __a    ,   __a       )
									snake_case								= model(**__a       )[0]
									snake_case								= (1, 1_0_2_4, model.config.vocab_size)
									self.assertEqual(output.shape    ,   __a       )
									# change to expected output here
									snake_case								= tf.convert_to_tensor(
									    [[33.65_07, 6.45_72, 16.80_89], [5.87_39, -2.42_38, 11.29_02], [-3.21_39, -4.31_49, 4.27_83]]    ,   )
									tf.debugging.assert_near(output[:, :3, :3]    ,   __a    ,   atol=1E-3    ,   rtol=1E-3       )



from . import (
    albert,
    align,
    altclip,
    audio_spectrogram_transformer,
    auto,
    autoformer,
    bark,
    bart,
    barthez,
    bartpho,
    beit,
    bert,
    bert_generation,
    bert_japanese,
    bertweet,
    big_bird,
    bigbird_pegasus,
    biogpt,
    bit,
    blenderbot,
    blenderbot_small,
    blip,
    blip_a,
    bloom,
    bridgetower,
    byta,
    camembert,
    canine,
    chinese_clip,
    clap,
    clip,
    clipseg,
    codegen,
    conditional_detr,
    convbert,
    convnext,
    convnextva,
    cpm,
    cpmant,
    ctrl,
    cvt,
    dataavec,
    deberta,
    deberta_va,
    decision_transformer,
    deformable_detr,
    deit,
    deprecated,
    deta,
    detr,
    dialogpt,
    dinat,
    distilbert,
    dit,
    donut,
    dpr,
    dpt,
    efficientformer,
    efficientnet,
    electra,
    encodec,
    encoder_decoder,
    ernie,
    ernie_m,
    esm,
    falcon,
    flaubert,
    flava,
    fnet,
    focalnet,
    fsmt,
    funnel,
    git,
    glpn,
    gpta,
    gpt_bigcode,
    gpt_neo,
    gpt_neox,
    gpt_neox_japanese,
    gpt_swa,
    gptj,
    gptsan_japanese,
    graphormer,
    groupvit,
    herbert,
    hubert,
    ibert,
    imagegpt,
    informer,
    instructblip,
    jukebox,
    layoutlm,
    layoutlmva,
    layoutlmva,
    layoutxlm,
    led,
    levit,
    lilt,
    llama,
    longformer,
    longta,
    luke,
    lxmert,
    mam_aaa,
    marian,
    markuplm,
    maskaformer,
    maskformer,
    mbart,
    mbartaa,
    mega,
    megatron_bert,
    megatron_gpta,
    mgp_str,
    mluke,
    mobilebert,
    mobilenet_va,
    mobilenet_va,
    mobilevit,
    mobilevitva,
    mpnet,
    mra,
    mta,
    musicgen,
    mvp,
    nat,
    nezha,
    nllb,
    nllb_moe,
    nystromformer,
    oneformer,
    open_llama,
    openai,
    opt,
    owlvit,
    pegasus,
    pegasus_x,
    perceiver,
    phobert,
    pixastruct,
    plbart,
    poolformer,
    prophetnet,
    qdqbert,
    rag,
    realm,
    reformer,
    regnet,
    rembert,
    resnet,
    roberta,
    roberta_prelayernorm,
    roc_bert,
    roformer,
    rwkv,
    sam,
    segformer,
    sew,
    sew_d,
    speech_encoder_decoder,
    speech_to_text,
    speech_to_text_a,
    speechta,
    splinter,
    squeezebert,
    swiftformer,
    swin,
    swinasr,
    swinva,
    switch_transformers,
    ta,
    table_transformer,
    tapas,
    time_series_transformer,
    timesformer,
    timm_backbone,
    transfo_xl,
    trocr,
    tvlt,
    umta,
    unispeech,
    unispeech_sat,
    upernet,
    videomae,
    vilt,
    vision_encoder_decoder,
    vision_text_dual_encoder,
    visual_bert,
    vit,
    vit_hybrid,
    vit_mae,
    vit_msn,
    vivit,
    wavaveca,
    wavaveca_conformer,
    wavaveca_phoneme,
    wavaveca_with_lm,
    wavlm,
    whisper,
    x_clip,
    xglm,
    xlm,
    xlm_prophetnet,
    xlm_roberta,
    xlm_roberta_xl,
    xlnet,
    xmod,
    yolos,
    yoso,
)

