





'''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       __SCREAMING_SNAKE_CASE	:



     '''simple docstring'''






     UpperCAmelCase__       :						Tuple      =   LEDConfig
     UpperCAmelCase__       :						int      =   {}
     UpperCAmelCase__       :						Tuple      =   '''gelu'''

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






     def 	UpperCamelCase(					self						,   lowerCamelCase						,   lowerCamelCase	):
       _snake_case     		=				TFLEDModel(config=lowerCamelCase	).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(lowerCamelCase						,   attention_mask=lowerCamelCase						,   use_cache=lowerCamelCase	)

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







def 				snake_case_      (		SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__=None					,	SCREAMING_SNAKE_CASE__=None					,	SCREAMING_SNAKE_CASE__=None					,	SCREAMING_SNAKE_CASE__=None					,	):







  '''simple docstring'''



  if attention_mask is None:
    _snake_case     		=				tf.cast(tf.math.not_equal(SCREAMING_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       __SCREAMING_SNAKE_CASE	(						__UpperCamelCase					,    __UpperCamelCase					,    unittest.TestCase			):



     '''simple docstring'''






     UpperCAmelCase__       :						Optional[int]      =   (TFLEDForConditionalGeneration, TFLEDModel) if is_tf_available() else ()
     UpperCAmelCase__       :						List[str]      =   (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__       :						str      =   True
     UpperCAmelCase__       :						str      =   False
     UpperCAmelCase__       :						Optional[int]      =   False
     UpperCAmelCase__       :						int      =   False

     def 	UpperCamelCase(					self	):
       _snake_case     		=				TFLEDModelTester(self	)
       _snake_case     		=				ConfigTester(self						,   config_class=lowerCamelCase	)

     def 	UpperCamelCase(					self	):
       self.config_tester.run_common_tests()

     def 	UpperCamelCase(					self	):
       _snake_case     		=				self.model_tester.prepare_config_and_inputs_for_common()
       self.model_tester.check_decoder_model_past_large_inputs(*lowerCamelCase	)

     def 	UpperCamelCase(					self	):
       _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(lowerCamelCase	):
         _snake_case     		=				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	):
         _snake_case     		=				[t.numpy() for t in outputs.encoder_attentions]
         _snake_case     		=				[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:
         _snake_case     		=				True
         _snake_case     		=				False
         _snake_case     		=				False
         _snake_case     		=				model_class(lowerCamelCase	)
         _snake_case     		=				model(self._prepare_for_class(lowerCamelCase						,   lowerCamelCase	)	)
         _snake_case     		=				len(lowerCamelCase	)
         self.assertEqual(config.output_hidden_states						,   lowerCamelCase	)
         check_encoder_attentions_output(lowerCamelCase	)

         if self.is_encoder_decoder:
           _snake_case     		=				model_class(lowerCamelCase	)
           _snake_case     		=				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"]
         _snake_case     		=				True
         _snake_case     		=				model_class(lowerCamelCase	)
         _snake_case     		=				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
         _snake_case     		=				True
         _snake_case     		=				True
         _snake_case     		=				model_class(lowerCamelCase	)
         _snake_case     		=				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	):
       pass






     def 	UpperCamelCase(					self	):
       # TODO: Head-masking not yet implement
       pass







def 				snake_case_      (		SCREAMING_SNAKE_CASE__       ):







  '''simple docstring'''



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


__magic_name__   :						List[str]    =			1e-4




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



     '''simple docstring'''






     def 	UpperCamelCase(					self	):
       _snake_case     		=				TFLEDForConditionalGeneration.from_pretrained("allenai/led-base-16384"	).led

       # change to intended input here
       _snake_case     		=				_long_tensor([512 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]	)
       _snake_case     		=				_long_tensor([128 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]	)
       _snake_case     		=				prepare_led_inputs_dict(model.config						,   lowerCamelCase						,   lowerCamelCase	)
       _snake_case     		=				model(**lowerCamelCase	)[0]
       _snake_case     		=				(1, 1_024, 768)
       self.assertEqual(output.shape						,   lowerCamelCase	)
       # change to expected output here
       _snake_case     		=				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	):
       _snake_case     		=				TFLEDForConditionalGeneration.from_pretrained("allenai/led-base-16384"	)

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









'''simple docstring'''



from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__magic_name__   :						int    =			logging.get_logger(__name__)

__magic_name__   :						Tuple    =			{
    """xlm-mlm-en-2048""": """https://huggingface.co/xlm-mlm-en-2048/resolve/main/config.json""",
    """xlm-mlm-ende-1024""": """https://huggingface.co/xlm-mlm-ende-1024/resolve/main/config.json""",
    """xlm-mlm-enfr-1024""": """https://huggingface.co/xlm-mlm-enfr-1024/resolve/main/config.json""",
    """xlm-mlm-enro-1024""": """https://huggingface.co/xlm-mlm-enro-1024/resolve/main/config.json""",
    """xlm-mlm-tlm-xnli15-1024""": """https://huggingface.co/xlm-mlm-tlm-xnli15-1024/resolve/main/config.json""",
    """xlm-mlm-xnli15-1024""": """https://huggingface.co/xlm-mlm-xnli15-1024/resolve/main/config.json""",
    """xlm-clm-enfr-1024""": """https://huggingface.co/xlm-clm-enfr-1024/resolve/main/config.json""",
    """xlm-clm-ende-1024""": """https://huggingface.co/xlm-clm-ende-1024/resolve/main/config.json""",
    """xlm-mlm-17-1280""": """https://huggingface.co/xlm-mlm-17-1280/resolve/main/config.json""",
    """xlm-mlm-100-1280""": """https://huggingface.co/xlm-mlm-100-1280/resolve/main/config.json""",
}


class       __SCREAMING_SNAKE_CASE	(						__UpperCamelCase			):



     '''simple docstring'''






     UpperCAmelCase__       :						Union[str, Any]      =   '''xlm'''
     UpperCAmelCase__       :						Dict      =   {
         '''hidden_size''': '''emb_dim''',
         '''num_attention_heads''': '''n_heads''',
         '''num_hidden_layers''': '''n_layers''',
         '''n_words''': '''vocab_size''',  # For backward compatibility
     }

     def __init__(					self						,   lowerCamelCase=30_145						,   lowerCamelCase=2_048						,   lowerCamelCase=12						,   lowerCamelCase=16						,   lowerCamelCase=0.1						,   lowerCamelCase=0.1						,   lowerCamelCase=True						,   lowerCamelCase=False						,   lowerCamelCase=False						,   lowerCamelCase=False						,   lowerCamelCase=1						,   lowerCamelCase=True						,   lowerCamelCase=512						,   lowerCamelCase=2_048**-0.5						,   lowerCamelCase=1e-12						,   lowerCamelCase=0.02						,   lowerCamelCase=0						,   lowerCamelCase=1						,   lowerCamelCase=2						,   lowerCamelCase=3						,   lowerCamelCase=5						,   lowerCamelCase=True						,   lowerCamelCase="first"						,   lowerCamelCase=True						,   lowerCamelCase=None						,   lowerCamelCase=True						,   lowerCamelCase=0.1						,   lowerCamelCase=5						,   lowerCamelCase=5						,   lowerCamelCase=0						,   lowerCamelCase=0						,   lowerCamelCase=2						,   lowerCamelCase=0						,   **lowerCamelCase						,   ):
       _snake_case     		=				vocab_size
       _snake_case     		=				emb_dim
       _snake_case     		=				n_layers
       _snake_case     		=				n_heads
       _snake_case     		=				dropout
       _snake_case     		=				attention_dropout
       _snake_case     		=				gelu_activation
       _snake_case     		=				sinusoidal_embeddings
       _snake_case     		=				causal
       _snake_case     		=				asm
       _snake_case     		=				n_langs
       _snake_case     		=				use_lang_emb
       _snake_case     		=				layer_norm_eps
       _snake_case     		=				bos_index
       _snake_case     		=				eos_index
       _snake_case     		=				pad_index
       _snake_case     		=				unk_index
       _snake_case     		=				mask_index
       _snake_case     		=				is_encoder
       _snake_case     		=				max_position_embeddings
       _snake_case     		=				embed_init_std
       _snake_case     		=				init_std
       _snake_case     		=				summary_type
       _snake_case     		=				summary_use_proj
       _snake_case     		=				summary_activation
       _snake_case     		=				summary_proj_to_labels
       _snake_case     		=				summary_first_dropout
       _snake_case     		=				start_n_top
       _snake_case     		=				end_n_top
       _snake_case     		=				mask_token_id
       _snake_case     		=				lang_id

       if "n_words" in kwargs:
         _snake_case     		=				kwargs["n_words"]

       super().__init__(pad_token_id=lowerCamelCase						,   bos_token_id=lowerCamelCase						,   **lowerCamelCase	)







class       __SCREAMING_SNAKE_CASE	(						__UpperCamelCase			):



     '''simple docstring'''






     @property
     def 	UpperCamelCase(					self	):
       if self.task == "multiple-choice":
         _snake_case     		=				{0: "batch", 1: "choice", 2: "sequence"}
       else:
         _snake_case     		=				{0: "batch", 1: "sequence"}
       return OrderedDict(
           [
               ("input_ids", dynamic_axis),
               ("attention_mask", dynamic_axis),
               ("token_type_ids", dynamic_axis),
           ]	)


