




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      :
      UpperCAmelCase_        =		LEDConfig
      UpperCAmelCase_        =		{}
      UpperCAmelCase_        =		"""gelu"""
      def __init__(			self   :Union[str, Any] ,      lowerCamelCase   :Dict ,      lowerCamelCase   :Optional[Any]=13 ,      lowerCamelCase   :Optional[int]=7 ,      lowerCamelCase   :str=True ,      lowerCamelCase   :Union[str, Any]=False ,      lowerCamelCase   :List[Any]=99 ,      lowerCamelCase   :Tuple=32 ,      lowerCamelCase   :str=2 ,      lowerCamelCase   :Optional[Any]=4 ,      lowerCamelCase   :Optional[int]=37 ,      lowerCamelCase   :str=0.1 ,      lowerCamelCase   :Tuple=0.1 ,      lowerCamelCase   :List[Any]=20 ,      lowerCamelCase   :Optional[Any]=2 ,      lowerCamelCase   :Union[str, Any]=1 ,      lowerCamelCase   :Optional[int]=0 ,      lowerCamelCase   :Optional[Any]=4 ,      )     ->			List[Any]:
       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   :Dict      )     ->			Any:
       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(__A ,      __A ,      __A      )
       UpperCAmelCase__							 =						tf.concat(
           [tf.zeros_like(__A      )[:, :-1], tf.ones_like(__A      )[:, -1:]] ,      axis=-1 ,      )
       UpperCAmelCase__							 =						global_attention_mask
       return config, inputs_dict







      def 	UpperCAmelCase_						(			self   :int ,      lowerCamelCase   :Dict ,      lowerCamelCase   :Optional[int]      )     ->			List[Any]:
       UpperCAmelCase__							 =						TFLEDModel(config=__A      ).get_decoder()
       UpperCAmelCase__							 =						inputs_dict["input_ids"]

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

       # first forward pass
       UpperCAmelCase__							 =						model(__A ,      attention_mask=__A ,      use_cache=__A      )

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



def 						lowerCAmelCase		(		_lowerCAmelCase							:   Tuple						,							_lowerCAmelCase							:   Any						,							_lowerCAmelCase							:   Tuple						,							_lowerCAmelCase							:   str=None						,							_lowerCAmelCase							:   Optional[int]=None						,							_lowerCAmelCase							:   List[Any]=None						,							_lowerCAmelCase							:   Tuple=None						,							):




 """simple docstring"""







 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      (				lowerCAmelCase	,       lowerCAmelCase	,       unittest.TestCase		):
      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 	UpperCAmelCase_						(			self   :Any      )     ->			List[Any]:
       UpperCAmelCase__							 =						TFLEDModelTester(self      )
       UpperCAmelCase__							 =						ConfigTester(self ,      config_class=__A      )
      def 	UpperCAmelCase_						(			self   :int      )     ->			Union[str, Any]:
       self.config_tester.run_common_tests()
      def 	UpperCAmelCase_						(			self   :List[str]      )     ->			int:
       UpperCAmelCase__							 =						self.model_tester.prepare_config_and_inputs_for_common()
       self.model_tester.check_decoder_model_past_large_inputs(*__A      )
      def 	UpperCAmelCase_						(			self   :Optional[int]      )     ->			Optional[int]:
       UpperCAmelCase__ , 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   :str      ):
        UpperCAmelCase__							 =						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(lowerCamelCase   :Tuple      ):
        UpperCAmelCase__							 =						[t.numpy() for t in outputs.encoder_attentions]
        UpperCAmelCase__							 =						[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:
        UpperCAmelCase__							 =						True
        UpperCAmelCase__							 =						False
        UpperCAmelCase__							 =						False
        UpperCAmelCase__							 =						model_class(__A      )
        UpperCAmelCase__							 =						model(self._prepare_for_class(__A ,      __A      )      )
        UpperCAmelCase__							 =						len(__A      )
        self.assertEqual(config.output_hidden_states ,      __A      )
        check_encoder_attentions_output(__A      )

        if self.is_encoder_decoder:
         UpperCAmelCase__							 =						model_class(__A      )
         UpperCAmelCase__							 =						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"]
        UpperCAmelCase__							 =						True
        UpperCAmelCase__							 =						model_class(__A      )
        UpperCAmelCase__							 =						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
        UpperCAmelCase__							 =						True
        UpperCAmelCase__							 =						True
        UpperCAmelCase__							 =						model_class(__A      )
        UpperCAmelCase__							 =						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   :Optional[int]      )     ->			List[str]:
       pass







      def 	UpperCAmelCase_						(			self   :Dict      )     ->			Optional[Any]:
       # TODO: Head-masking not yet implement
       pass



def 						lowerCAmelCase		(		_lowerCAmelCase							:   Optional[Any]	):




 """simple docstring"""







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


_lowerCAmelCase						:   List[Any]		=   1E-4


@slow
@require_tf
class 			_UpperCamelCase      (				unittest.TestCase		):
      def 	UpperCAmelCase_						(			self   :Optional[int]      )     ->			Any:
       UpperCAmelCase__							 =						TFLEDForConditionalGeneration.from_pretrained("allenai/led-base-16384"      ).led

       # change to intended input here
       UpperCAmelCase__							 =						_long_tensor([512 * [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69]]      )
       UpperCAmelCase__							 =						_long_tensor([128 * [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69]]      )
       UpperCAmelCase__							 =						prepare_led_inputs_dict(model.config ,      __A ,      __A      )
       UpperCAmelCase__							 =						model(**__A      )[0]
       UpperCAmelCase__							 =						(1, 1024, 768)
       self.assertEqual(output.shape ,      __A      )
       # change to expected output here
       UpperCAmelCase__							 =						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      )     ->			Union[str, Any]:
       UpperCAmelCase__							 =						TFLEDForConditionalGeneration.from_pretrained("allenai/led-base-16384"      )

       # change to intended input here
       UpperCAmelCase__							 =						_long_tensor([512 * [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69]]      )
       UpperCAmelCase__							 =						_long_tensor([128 * [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69]]      )
       UpperCAmelCase__							 =						prepare_led_inputs_dict(model.config ,      __A ,      __A      )
       UpperCAmelCase__							 =						model(**__A      )[0]
       UpperCAmelCase__							 =						(1, 1024, model.config.vocab_size)
       self.assertEqual(output.shape ,      __A      )
       # change to expected output here
       UpperCAmelCase__							 =						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 typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tensorflow_text_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


_lowerCAmelCase						:   Union[str, Any]		=   {
    "configuration_bert": ["BERT_PRETRAINED_CONFIG_ARCHIVE_MAP", "BertConfig", "BertOnnxConfig"],
    "tokenization_bert": ["BasicTokenizer", "BertTokenizer", "WordpieceTokenizer"],
}

try:
  if not is_tokenizers_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _lowerCAmelCase						:   str		=   ["BertTokenizerFast"]

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _lowerCAmelCase						:   int		=   [
      "BERT_PRETRAINED_MODEL_ARCHIVE_LIST",
      "BertForMaskedLM",
      "BertForMultipleChoice",
      "BertForNextSentencePrediction",
      "BertForPreTraining",
      "BertForQuestionAnswering",
      "BertForSequenceClassification",
      "BertForTokenClassification",
      "BertLayer",
      "BertLMHeadModel",
      "BertModel",
      "BertPreTrainedModel",
      "load_tf_weights_in_bert",
  ]

try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _lowerCAmelCase						:   Dict		=   [
      "TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST",
      "TFBertEmbeddings",
      "TFBertForMaskedLM",
      "TFBertForMultipleChoice",
      "TFBertForNextSentencePrediction",
      "TFBertForPreTraining",
      "TFBertForQuestionAnswering",
      "TFBertForSequenceClassification",
      "TFBertForTokenClassification",
      "TFBertLMHeadModel",
      "TFBertMainLayer",
      "TFBertModel",
      "TFBertPreTrainedModel",
  ]
try:
  if not is_tensorflow_text_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _lowerCAmelCase						:   str		=   ["TFBertTokenizer"]

try:
  if not is_flax_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _lowerCAmelCase						:   Optional[Any]		=   [
      "FlaxBertForCausalLM",
      "FlaxBertForMaskedLM",
      "FlaxBertForMultipleChoice",
      "FlaxBertForNextSentencePrediction",
      "FlaxBertForPreTraining",
      "FlaxBertForQuestionAnswering",
      "FlaxBertForSequenceClassification",
      "FlaxBertForTokenClassification",
      "FlaxBertModel",
      "FlaxBertPreTrainedModel",
  ]

if TYPE_CHECKING:
  from .configuration_bert import BERT_PRETRAINED_CONFIG_ARCHIVE_MAP, BertConfig, BertOnnxConfig
  from .tokenization_bert import BasicTokenizer, BertTokenizer, WordpieceTokenizer

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_bert_fast import BertTokenizerFast

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_bert import (
        BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        BertForMaskedLM,
        BertForMultipleChoice,
        BertForNextSentencePrediction,
        BertForPreTraining,
        BertForQuestionAnswering,
        BertForSequenceClassification,
        BertForTokenClassification,
        BertLayer,
        BertLMHeadModel,
        BertModel,
        BertPreTrainedModel,
        load_tf_weights_in_bert,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_bert import (
        TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFBertEmbeddings,
        TFBertForMaskedLM,
        TFBertForMultipleChoice,
        TFBertForNextSentencePrediction,
        TFBertForPreTraining,
        TFBertForQuestionAnswering,
        TFBertForSequenceClassification,
        TFBertForTokenClassification,
        TFBertLMHeadModel,
        TFBertMainLayer,
        TFBertModel,
        TFBertPreTrainedModel,
    )

  try:
    if not is_tensorflow_text_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_bert_tf import TFBertTokenizer

  try:
    if not is_flax_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_flax_bert import (
        FlaxBertForCausalLM,
        FlaxBertForMaskedLM,
        FlaxBertForMultipleChoice,
        FlaxBertForNextSentencePrediction,
        FlaxBertForPreTraining,
        FlaxBertForQuestionAnswering,
        FlaxBertForSequenceClassification,
        FlaxBertForTokenClassification,
        FlaxBertModel,
        FlaxBertPreTrainedModel,
    )

else:
  import sys

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





