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"""






  lowerCAmelCase__		:	Any    		= LEDConfig
  lowerCAmelCase__		:	List[Any]    		= {}
  lowerCAmelCase__		:	Union[str, Any]    		= 'gelu'






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




    '''simple docstring'''




    __UpperCAmelCase						    =       parent
    __UpperCAmelCase						    =       batch_size
    __UpperCAmelCase						    =       seq_length
    __UpperCAmelCase						    =       is_training
    __UpperCAmelCase						    =       use_labels
    __UpperCAmelCase						    =       vocab_size
    __UpperCAmelCase						    =       hidden_size
    __UpperCAmelCase						    =       num_hidden_layers
    __UpperCAmelCase						    =       num_attention_heads
    __UpperCAmelCase						    =       intermediate_size

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






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




    '''simple docstring'''




    __UpperCAmelCase						    =       ids_tensor([self.batch_size, self.seq_length - 1]				,				self.vocab_size   )
    __UpperCAmelCase						    =       tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size   )				,				1   )
    __UpperCAmelCase						    =       tf.concat([input_ids, eos_tensor]				,				axis=1   )

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

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






  def 	_UpperCAmelCase				(				self:			str				,				__lowerCAmelCase:			Dict				,				__lowerCAmelCase:			Dict   ) -> Tuple:




    '''simple docstring'''




    __UpperCAmelCase						    =       TFLEDModel(config=__lowerCAmelCase   ).get_decoder()
    __UpperCAmelCase						    =       inputs_dict['input_ids']

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

    # first forward pass
    __UpperCAmelCase						    =       model(__lowerCAmelCase				,				attention_mask=__lowerCAmelCase				,				use_cache=__lowerCAmelCase   )

    __UpperCAmelCase						    =       outputs.to_tuple()

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

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

    __UpperCAmelCase						    =       model(__lowerCAmelCase				,				attention_mask=__lowerCAmelCase   )[0]
    __UpperCAmelCase						    =       model(__lowerCAmelCase				,				attention_mask=__lowerCAmelCase				,				past_key_values=__lowerCAmelCase   )[0]

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

    # select random slice
    __UpperCAmelCase						    =       int(ids_tensor((1,)				,				output_from_past.shape[-1]   )   )
    __UpperCAmelCase						    =       output_from_no_past[:, -3:, random_slice_idx]
    __UpperCAmelCase						    =       output_from_past[:, :, random_slice_idx]

    # test that outputs are equal for slice
    tf.debugging.assert_near(__lowerCAmelCase				,				__lowerCAmelCase				,				rtol=1E-3   )






def       __lowerCAmelCase      (						A_   : Optional[Any]						,	A_   : str						,	A_   : List[str]						,	A_   : str=None						,	A_   : Tuple=None						,	A_   : List[Any]=None						,	A_   : Tuple=None						,	)    ->  str:




  if attention_mask is None:
    __UpperCAmelCase						    =       tf.cast(tf.math.not_equal(__lowercase						,	config.pad_token_id					)						,	tf.inta					)
  if decoder_attention_mask is None:
    __UpperCAmelCase						    =       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:
    __UpperCAmelCase						    =       tf.ones((config.encoder_layers, config.encoder_attention_heads)					)
  if decoder_head_mask is None:
    __UpperCAmelCase						    =       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__    (		__A		,	__A		,	unittest.TestCase					):



  """simple docstring"""






  lowerCAmelCase__		:	List[Any]    		= (TFLEDForConditionalGeneration, TFLEDModel) if is_tf_available() else ()
  lowerCAmelCase__		:	Optional[int]    		= (TFLEDForConditionalGeneration,) if is_tf_available() else ()
  lowerCAmelCase__		:	int    		= (
      {
          'conversational': TFLEDForConditionalGeneration,
          'feature-extraction': TFLEDModel,
          'summarization': TFLEDForConditionalGeneration,
          'text2text-generation': TFLEDForConditionalGeneration,
          'translation': TFLEDForConditionalGeneration,
      }
      if is_tf_available()
      else {}
  )
  lowerCAmelCase__		:	Dict    		= True
  lowerCAmelCase__		:	Optional[Any]    		= False
  lowerCAmelCase__		:	List[Any]    		= False
  lowerCAmelCase__		:	int    		= False






  def 	_UpperCAmelCase				(				self:			Optional[int]   ) -> Dict:




    '''simple docstring'''




    __UpperCAmelCase						    =       TFLEDModelTester(self   )
    __UpperCAmelCase						    =       ConfigTester(self				,				config_class=__lowerCAmelCase   )






  def 	_UpperCAmelCase				(				self:			Tuple   ) -> Dict:




    '''simple docstring'''




    self.config_tester.run_common_tests()






  def 	_UpperCAmelCase				(				self:			Optional[Any]   ) -> Optional[int]:




    '''simple docstring'''




    __UpperCAmelCase						    =       self.model_tester.prepare_config_and_inputs_for_common()
    self.model_tester.check_decoder_model_past_large_inputs(*__lowerCAmelCase   )






  def 	_UpperCAmelCase				(				self:			Any   ) -> Any:




    '''simple docstring'''




    __UpperCAmelCase						    =       self.model_tester.prepare_config_and_inputs_for_common()
    __UpperCAmelCase						    =       tf.zeros_like(inputs_dict["attention_mask"]   )
    __UpperCAmelCase						    =       2
    __UpperCAmelCase						    =       tf.where(
        tf.range(self.model_tester.seq_length   )[None, :] < num_global_attn_indices				,				1				,				inputs_dict["global_attention_mask"]				,				)

    __UpperCAmelCase						    =       True
    __UpperCAmelCase						    =       self.model_tester.seq_length
    __UpperCAmelCase						    =       self.model_tester.encoder_seq_length

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

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

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

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

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






  @unittest.skip("LED keeps using potentially symbolic tensors in conditionals and breaks tracing."   )
  def 	_UpperCAmelCase				(				self:			Dict   ) -> int:




    '''simple docstring'''




    pass






  def 	_UpperCAmelCase				(				self:			int   ) -> Tuple:




    '''simple docstring'''




    pass






def       __lowerCAmelCase      (						A_   : Union[str, Any]					)    ->  Dict:
  return tf.constant(__lowercase						,	dtype=tf.intaa					)


a_							=      1e-4





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



  """simple docstring"""






  def 	_UpperCAmelCase				(				self:			List[Any]   ) -> int:




    '''simple docstring'''




    __UpperCAmelCase						    =       TFLEDForConditionalGeneration.from_pretrained("allenai/led-base-16384"   ).led

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






  def 	_UpperCAmelCase				(				self:			str   ) -> Any:




    '''simple docstring'''




    __UpperCAmelCase						    =       TFLEDForConditionalGeneration.from_pretrained("allenai/led-base-16384"   )

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




import inspect
import unittest


class 			UpperCAmelCase__    (		unittest.TestCase					):



  """simple docstring"""






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




    '''simple docstring'''




    try:
      import diffusers  # noqa: F401
    except ImportError:
      assert False






  def 	_UpperCAmelCase				(				self:			Dict   ) -> Dict:




    '''simple docstring'''




    import diffusers
    from diffusers.dependency_versions_table import deps

    __UpperCAmelCase						    =       inspect.getmembers(__lowerCAmelCase				,				inspect.isclass   )

    for cls_name, cls_module in all_classes:
      if "dummy_" in cls_module.__module__:
        for backend in cls_module._backends:
          if backend == "k_diffusion":
            __UpperCAmelCase						    =       "k-diffusion"
          elif backend == "invisible_watermark":
            __UpperCAmelCase						    =       "invisible-watermark"
          assert backend in deps, F'''{backend} is not in the deps table!'''



