





"""simple docstring"""





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 		lowercase					:
      _SCREAMING_SNAKE_CASE					=      LEDConfig
      _SCREAMING_SNAKE_CASE					=      {}
      _SCREAMING_SNAKE_CASE					=      'gelu'





      def __init__(       self      ,						lowercase      ,						lowercase=13      ,						lowercase=7      ,						lowercase=True      ,						lowercase=False      ,						lowercase=99      ,						lowercase=32      ,						lowercase=2      ,						lowercase=4      ,						lowercase=37      ,						lowercase=0.1      ,						lowercase=0.1      ,						lowercase=20      ,						lowercase=2      ,						lowercase=1      ,						lowercase=0      ,						lowercase=4      ,						)		->     int:
            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 )		->     Optional[int]:
            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(lowercase      ,						lowercase      ,						lowercase )
            lowerCAmelCase  						=       tf.concat(
                [tf.zeros_like(lowercase )[:, :-1], tf.ones_like(lowercase )[:, -1:]]      ,						axis=-1      ,						)
            lowerCAmelCase  						=       global_attention_mask
            return config, inputs_dict





      def     _snake_case			(       self      ,						lowercase      ,						lowercase )		->     List[Any]:
            lowerCAmelCase  						=       TFLEDModel(config=lowercase ).get_decoder()
            lowerCAmelCase  						=       inputs_dict["""input_ids"""]

            lowerCAmelCase  						=       input_ids[:1, :]
            lowerCAmelCase  						=       inputs_dict["""attention_mask"""][:1, :]
            lowerCAmelCase  						=       1

            # first forward pass
            lowerCAmelCase  						=       model(lowercase      ,						attention_mask=lowercase      ,						use_cache=lowercase )

            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(lowercase      ,						attention_mask=lowercase )[0]
            lowerCAmelCase  						=       model(lowercase      ,						attention_mask=lowercase      ,						past_key_values=lowercase )[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(lowercase      ,						lowercase      ,						rtol=1e-3 )







def 			UpperCAmelCase__							(  SCREAMING_SNAKE_CASE     :						Optional[Any]							,						SCREAMING_SNAKE_CASE     :						Optional[Any]							,						SCREAMING_SNAKE_CASE     :						List[str]							,						SCREAMING_SNAKE_CASE     :						List[Any]=None							,						SCREAMING_SNAKE_CASE     :						Optional[int]=None							,						SCREAMING_SNAKE_CASE     :						Optional[int]=None							,						SCREAMING_SNAKE_CASE     :						Optional[int]=None							,						):







      '''simple docstring'''




      if attention_mask is None:
            lowerCAmelCase  						=       tf.cast(tf.math.not_equal(SCREAMING_SNAKE_CASE							,						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 		lowercase					(					_UpperCAmelCase     ,  _UpperCAmelCase     ,  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     _snake_case			(       self )		->     Optional[Any]:
            lowerCAmelCase  						=       TFLEDModelTester(self )
            lowerCAmelCase  						=       ConfigTester(self      ,						config_class=lowercase )





      def     _snake_case			(       self )		->     Any:
            self.config_tester.run_common_tests()





      def     _snake_case			(       self )		->     List[str]:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs_for_common()
            self.model_tester.check_decoder_model_past_large_inputs(*lowercase )





      def     _snake_case			(       self )		->     Union[str, 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(lowercase ):
                  lowerCAmelCase  						=       outputs.decoder_attentions
                  self.assertEqual(len(lowercase )      ,						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(lowercase ):
                  lowerCAmelCase  						=       [t.numpy() for t in outputs.encoder_attentions]
                  lowerCAmelCase  						=       [t.numpy() for t in outputs.encoder_global_attentions]
                  self.assertEqual(len(lowercase )      ,						self.model_tester.num_hidden_layers )
                  self.assertEqual(len(lowercase )      ,						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(lowercase )
                  lowerCAmelCase  						=       model(self._prepare_for_class(lowercase      ,						lowercase ) )
                  lowerCAmelCase  						=       len(lowercase )
                  self.assertEqual(config.output_hidden_states      ,						lowercase )
                  check_encoder_attentions_output(lowercase )

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

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

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

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





      @unittest.skip("""LED keeps using potentially symbolic tensors in conditionals and breaks tracing.""" )
      def     _snake_case			(       self )		->     Union[str, Any]:
            pass





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







def 			UpperCAmelCase__							(  SCREAMING_SNAKE_CASE     :						Union[str, Any]			):







      '''simple docstring'''




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


SCREAMING_SNAKE_CASE__					=   1e-4






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





      def     _snake_case			(       self )		->     int:
            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      ,						lowercase      ,						lowercase )
            lowerCAmelCase  						=       model(**lowercase )[0]
            lowerCAmelCase  						=       (1, 1_024, 768)
            self.assertEqual(output.shape      ,						lowercase )
            # 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]      ,						lowercase      ,						atol=1e-3 )





      def     _snake_case			(       self )		->     Optional[Any]:
            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      ,						lowercase      ,						lowercase )
            lowerCAmelCase  						=       model(**lowercase )[0]
            lowerCAmelCase  						=       (1, 1_024, model.config.vocab_size)
            self.assertEqual(output.shape      ,						lowercase )
            # 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]      ,						lowercase      ,						atol=1e-3      ,						rtol=1e-3 )













"""simple docstring"""





# This is the module that test_patching.py uses to test patch_submodule()

import os  # noqa: this is just for tests
import os as renamed_os  # noqa: this is just for tests
from os import path  # noqa: this is just for tests
from os import path as renamed_path  # noqa: this is just for tests
from os.path import join  # noqa: this is just for tests
from os.path import join as renamed_join  # noqa: this is just for tests


SCREAMING_SNAKE_CASE__					=   open  # noqa: we just need to have a builtin inside this module to test it properly






