

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




   lowerCamelCase__				=      LEDConfig
   lowerCamelCase__				=      {}
   lowerCamelCase__				=      "gelu"





   def __init__(       self					:			Tuple     , __lowerCamelCase					:			Optional[Any]     , __lowerCamelCase					:			List[str]=13     , __lowerCamelCase					:			str=7     , __lowerCamelCase					:			Any=True     , __lowerCamelCase					:			str=False     , __lowerCamelCase					:			Optional[Any]=99     , __lowerCamelCase					:			Any=32     , __lowerCamelCase					:			Tuple=2     , __lowerCamelCase					:			str=4     , __lowerCamelCase					:			Optional[int]=37     , __lowerCamelCase					:			Dict=0.1     , __lowerCamelCase					:			str=0.1     , __lowerCamelCase					:			Any=20     , __lowerCamelCase					:			Optional[int]=2     , __lowerCamelCase					:			List[Any]=1     , __lowerCamelCase					:			Optional[Any]=0     , __lowerCamelCase					:			Dict=4     , ):
       SCREAMING_SNAKE_CASE										=   parent
       SCREAMING_SNAKE_CASE										=   batch_size
       SCREAMING_SNAKE_CASE										=   seq_length
       SCREAMING_SNAKE_CASE										=   is_training
       SCREAMING_SNAKE_CASE										=   use_labels
       SCREAMING_SNAKE_CASE										=   vocab_size
       SCREAMING_SNAKE_CASE										=   hidden_size
       SCREAMING_SNAKE_CASE										=   num_hidden_layers
       SCREAMING_SNAKE_CASE										=   num_attention_heads
       SCREAMING_SNAKE_CASE										=   intermediate_size

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





   def      _snake_case		(       self					:			str			):
       SCREAMING_SNAKE_CASE										=   ids_tensor([self.batch_size, self.seq_length - 1]     , self.vocab_size			)
       SCREAMING_SNAKE_CASE										=   tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size			)     , 1			)
       SCREAMING_SNAKE_CASE										=   tf.concat([input_ids, eos_tensor]     , axis=1			)

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

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





   def      _snake_case		(       self					:			int     , __lowerCamelCase					:			str     , __lowerCamelCase					:			Optional[Any]			):
       SCREAMING_SNAKE_CASE										=   TFLEDModel(config=__lowerCamelCase			).get_decoder()
       SCREAMING_SNAKE_CASE										=   inputs_dict["input_ids"]

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

       # first forward pass
       SCREAMING_SNAKE_CASE										=   model(__lowerCamelCase     , attention_mask=__lowerCamelCase     , use_cache=__lowerCamelCase			)

       SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE										=   outputs.to_tuple()

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

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

       SCREAMING_SNAKE_CASE										=   model(__lowerCamelCase     , attention_mask=__lowerCamelCase			)[0]
       SCREAMING_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
       SCREAMING_SNAKE_CASE										=   int(ids_tensor((1,)     , output_from_past.shape[-1]			)			)
       SCREAMING_SNAKE_CASE										=   output_from_no_past[:, -3:, random_slice_idx]
       SCREAMING_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 	__a     (      A__			:  int			,  A__			:  Dict			,  A__			:  List[str]			,  A__			:  int=None			,  A__			:  List[Any]=None			,  A__			:  Optional[int]=None			,  A__			:  Dict=None			,  ):
    if attention_mask is None:
        SCREAMING_SNAKE_CASE										=   tf.cast(tf.math.not_equal(A__			,  config.pad_token_id       )			,  tf.inta       )
    if decoder_attention_mask is None:
        SCREAMING_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:
        SCREAMING_SNAKE_CASE										=   tf.ones((config.encoder_layers, config.encoder_attention_heads)       )
    if decoder_head_mask is None:
        SCREAMING_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 (   __snake_case					,       __snake_case					,       unittest.TestCase		):







   '''simple docstring'''




   lowerCamelCase__				=      (TFLEDForConditionalGeneration, TFLEDModel) if is_tf_available() else ()
   lowerCamelCase__				=      (TFLEDForConditionalGeneration,) if is_tf_available() else ()
   lowerCamelCase__				=      (
       {
           "conversational": TFLEDForConditionalGeneration,
           "feature-extraction": TFLEDModel,
           "summarization": TFLEDForConditionalGeneration,
           "text2text-generation": TFLEDForConditionalGeneration,
           "translation": TFLEDForConditionalGeneration,
       }
       if is_tf_available()
       else {}
   )
   lowerCamelCase__				=      True
   lowerCamelCase__				=      False
   lowerCamelCase__				=      False
   lowerCamelCase__				=      False





   def      _snake_case		(       self					:			str			):
       SCREAMING_SNAKE_CASE										=   TFLEDModelTester(self			)
       SCREAMING_SNAKE_CASE										=   ConfigTester(self     , config_class=__lowerCamelCase			)





   def      _snake_case		(       self					:			Tuple			):
       self.config_tester.run_common_tests()





   def      _snake_case		(       self					:			int			):
       SCREAMING_SNAKE_CASE										=   self.model_tester.prepare_config_and_inputs_for_common()
       self.model_tester.check_decoder_model_past_large_inputs(*__lowerCamelCase			)





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

       SCREAMING_SNAKE_CASE										=   True
       SCREAMING_SNAKE_CASE										=   self.model_tester.seq_length
       SCREAMING_SNAKE_CASE										=   self.model_tester.encoder_seq_length

       def check_decoder_attentions_output(__lowerCamelCase					:			Tuple			):
           SCREAMING_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					:			Dict			):
           SCREAMING_SNAKE_CASE										=   [t.numpy() for t in outputs.encoder_attentions]
           SCREAMING_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:
           SCREAMING_SNAKE_CASE										=   True
           SCREAMING_SNAKE_CASE										=   False
           SCREAMING_SNAKE_CASE										=   False
           SCREAMING_SNAKE_CASE										=   model_class(__lowerCamelCase			)
           SCREAMING_SNAKE_CASE										=   model(self._prepare_for_class(__lowerCamelCase     , __lowerCamelCase			)			)
           SCREAMING_SNAKE_CASE										=   len(__lowerCamelCase			)
           self.assertEqual(config.output_hidden_states     , __lowerCamelCase			)
           check_encoder_attentions_output(__lowerCamelCase			)

           if self.is_encoder_decoder:
               SCREAMING_SNAKE_CASE										=   model_class(__lowerCamelCase			)
               SCREAMING_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"]
           SCREAMING_SNAKE_CASE										=   True
           SCREAMING_SNAKE_CASE										=   model_class(__lowerCamelCase			)
           SCREAMING_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
           SCREAMING_SNAKE_CASE										=   True
           SCREAMING_SNAKE_CASE										=   True
           SCREAMING_SNAKE_CASE										=   model_class(__lowerCamelCase			)
           SCREAMING_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      _snake_case		(       self					:			List[Any]			):
       pass





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





def 	__a     (      A__			:  Any       ):
    return tf.constant(A__			,  dtype=tf.intaa       )


__A			:					Optional[int]						=			1e-4





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







   '''simple docstring'''




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

       # change to intended input here
       SCREAMING_SNAKE_CASE										=   _long_tensor([512 * [0, 31414, 232, 328, 740, 1140, 12695, 69]]			)
       SCREAMING_SNAKE_CASE										=   _long_tensor([128 * [0, 31414, 232, 328, 740, 1140, 12695, 69]]			)
       SCREAMING_SNAKE_CASE										=   prepare_led_inputs_dict(model.config     , __lowerCamelCase     , __lowerCamelCase			)
       SCREAMING_SNAKE_CASE										=   model(**__lowerCamelCase			)[0]
       SCREAMING_SNAKE_CASE										=   (1, 1024, 768)
       self.assertEqual(output.shape     , __lowerCamelCase			)
       # change to expected output here
       SCREAMING_SNAKE_CASE										=   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]     , __lowerCamelCase     , atol=1e-3			)





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

       # change to intended input here
       SCREAMING_SNAKE_CASE										=   _long_tensor([512 * [0, 31414, 232, 328, 740, 1140, 12695, 69]]			)
       SCREAMING_SNAKE_CASE										=   _long_tensor([128 * [0, 31414, 232, 328, 740, 1140, 12695, 69]]			)
       SCREAMING_SNAKE_CASE										=   prepare_led_inputs_dict(model.config     , __lowerCamelCase     , __lowerCamelCase			)
       SCREAMING_SNAKE_CASE										=   model(**__lowerCamelCase			)[0]
       SCREAMING_SNAKE_CASE										=   (1, 1024, model.config.vocab_size)
       self.assertEqual(output.shape     , __lowerCamelCase			)
       # change to expected output here
       SCREAMING_SNAKE_CASE										=   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]     , __lowerCamelCase     , atol=1e-3     , rtol=1e-3			)


from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


__A			:					List[Any]						=			{
    'configuration_git': ['GIT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'GitConfig', 'GitVisionConfig'],
    'processing_git': ['GitProcessor'],
}

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    __A			:					int						=			[
        'GIT_PRETRAINED_MODEL_ARCHIVE_LIST',
        'GitForCausalLM',
        'GitModel',
        'GitPreTrainedModel',
        'GitVisionModel',
    ]

if TYPE_CHECKING:
    from .configuration_git import GIT_PRETRAINED_CONFIG_ARCHIVE_MAP, GitConfig, GitVisionConfig
    from .processing_git import GitProcessor

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_git import (
            GIT_PRETRAINED_MODEL_ARCHIVE_LIST,
            GitForCausalLM,
            GitModel,
            GitPreTrainedModel,
            GitVisionModel,
        )

else:
    import sys

    __A			:					List[Any]						=			_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)