




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

    '''simple docstring'''






    UpperCAmelCase__													=   LEDConfig
    UpperCAmelCase__													=   {}
    UpperCAmelCase__													=   '''gelu'''


    def __init__(    self   :			Dict      ,				UpperCAmelCase__   :			Tuple      ,				UpperCAmelCase__   :			Tuple=13      ,				UpperCAmelCase__   :			Union[str, Any]=7      ,				UpperCAmelCase__   :			List[Any]=True      ,				UpperCAmelCase__   :			Tuple=False      ,				UpperCAmelCase__   :			int=99      ,				UpperCAmelCase__   :			int=32      ,				UpperCAmelCase__   :			Optional[int]=2      ,				UpperCAmelCase__   :			Optional[int]=4      ,				UpperCAmelCase__   :			Tuple=37      ,				UpperCAmelCase__   :			Optional[int]=0.1      ,				UpperCAmelCase__   :			Optional[int]=0.1      ,				UpperCAmelCase__   :			Union[str, Any]=20      ,				UpperCAmelCase__   :			List[Any]=2      ,				UpperCAmelCase__   :			Tuple=1      ,				UpperCAmelCase__   :			Union[str, Any]=0      ,				UpperCAmelCase__   :			Any=4      ,				)		->Dict:





           '''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  SCREAMING_SNAKE_CASE							(    self   :			List[str])		->Union[str, Any]:





           '''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(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   :			Optional[int]      ,				UpperCAmelCase__   :			List[Any]      ,				UpperCAmelCase__   :			List[str])		->Tuple:





           '''simple docstring'''

           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 SCREAMING_SNAKE_CASE       ( lowercase_				,       lowercase_				,       lowercase_				,       lowercase_=None				,       lowercase_=None				,       lowercase_=None				,       lowercase_=None				,       )     -> int:



       """simple docstring"""

       if attention_mask is None:
              A__   				=					tf.cast(tf.math.not_equal(lowercase_				,       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        UpperCamelCase_					(				UpperCAmelCase__	,      UpperCAmelCase__	,      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  SCREAMING_SNAKE_CASE							(    self   :			List[Any])		->List[str]:





           '''simple docstring'''

           A__   				=					TFLEDModelTester(self)
           A__   				=					ConfigTester(self      ,				config_class=UpperCAmelCase__)


    def  SCREAMING_SNAKE_CASE							(    self   :			List[str])		->Dict:





           '''simple docstring'''

           self.config_tester.run_common_tests()


    def  SCREAMING_SNAKE_CASE							(    self   :			Tuple)		->List[str]:





           '''simple docstring'''

           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   :			List[str])		->Optional[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(UpperCAmelCase__   :			Dict):
                  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__   :			Union[str, Any]):
                  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)		->Dict:





           '''simple docstring'''

           pass





    def  SCREAMING_SNAKE_CASE							(    self   :			Optional[int])		->str:





           '''simple docstring'''

           pass







def SCREAMING_SNAKE_CASE       ( lowercase_		)     -> List[Any]:



       """simple docstring"""

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


_lowerCamelCase			:	Optional[Any]							 =     1E-4


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

    '''simple docstring'''






    def  SCREAMING_SNAKE_CASE							(    self   :			Any)		->Union[str, Any]:





           '''simple docstring'''

           A__   				=					TFLEDForConditionalGeneration.from_pretrained('''allenai/led-base-16384''').led

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





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





           '''simple docstring'''

           A__   				=					TFLEDForConditionalGeneration.from_pretrained('''allenai/led-base-16384''')

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








from argparse import ArgumentParser

from . import BaseTransformersCLICommand







def SCREAMING_SNAKE_CASE       ( lowercase_		)     -> str:



       """simple docstring"""

       return DownloadCommand(args.model				,       args.cache_dir				,       args.force				,       args.trust_remote_code		)


class        UpperCamelCase_					(				UpperCAmelCase__  ):

    '''simple docstring'''






    @staticmethod
    def  SCREAMING_SNAKE_CASE							(    UpperCAmelCase__   :			ArgumentParser)		->Any:





           '''simple docstring'''

           A__   				=					parser.add_parser('''download''')
           download_parser.add_argument(
               '''--cache-dir'''      ,				type=UpperCAmelCase__      ,				default=UpperCAmelCase__      ,				help='''Path to location to store the models''')
           download_parser.add_argument(
               '''--force'''      ,				action='''store_true'''      ,				help='''Force the model to be download even if already in cache-dir''')
           download_parser.add_argument(
               '''--trust-remote-code'''      ,				action='''store_true'''      ,				help='''Whether or not to allow for custom models defined on the Hub in their own modeling files. Use only if you\'ve reviewed the code as it will execute on your local machine'''      ,				)
           download_parser.add_argument('''model'''      ,				type=UpperCAmelCase__      ,				help='''Name of the model to download''')
           download_parser.set_defaults(func=UpperCAmelCase__)


    def __init__(    self   :			Dict      ,				UpperCAmelCase__   :			str      ,				UpperCAmelCase__   :			str      ,				UpperCAmelCase__   :			bool      ,				UpperCAmelCase__   :			bool)		->int:





           '''simple docstring'''

           A__   				=					model
           A__   				=					cache
           A__   				=					force
           A__   				=					trust_remote_code





    def  SCREAMING_SNAKE_CASE							(    self   :			str)		->Union[str, Any]:





           '''simple docstring'''

           from ..models.auto import AutoModel, AutoTokenizer

           AutoModel.from_pretrained(
               self._model      ,				cache_dir=self._cache      ,				force_download=self._force      ,				trust_remote_code=self._trust_remote_code)
           AutoTokenizer.from_pretrained(
               self._model      ,				cache_dir=self._cache      ,				force_download=self._force      ,				trust_remote_code=self._trust_remote_code)


