

from typing import Tuple, Union

from ...modeling_outputs import BackboneOutput
from ...modeling_utils import PreTrainedModel
from ...utils import is_timm_available, is_torch_available, requires_backends
from ...utils.backbone_utils import BackboneMixin
from .configuration_timm_backbone import TimmBackboneConfig


if is_timm_available():
		import timm


if is_torch_available():
		from torch import Tensor


class SCREAMING_SNAKE_CASE  (		__snake_case							,						__snake_case     ):






			"""simple docstring"""




			__A						  =  """pixel_values"""
			__A						  =  False
			__A						  =  TimmBackboneConfig


			def __init__(    self ,				__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(self ,				'timm'							)
				super().__init__(__UpperCamelCase							)
				snake_case_      		= config

				if config.backbone is None:
					raise ValueError('backbone is not set in the config. Please set it to a timm model name.'							)

				if config.backbone not in timm.list_models():
					raise ValueError(f"""backbone {config.backbone} is not supported by timm."""							)

				if hasattr(__UpperCamelCase ,				'out_features'							) and config.out_features is not None:
					raise ValueError('out_features is not supported by TimmBackbone. Please use out_indices instead.'							)

				snake_case_      		= getattr(__UpperCamelCase ,				'use_pretrained_backbone' ,				__UpperCamelCase							)
				if pretrained is None:
					raise ValueError('use_pretrained_backbone is not set in the config. Please set it to True or False.'							)

				# We just take the final layer by default. This matches the default for the transformers models.
				snake_case_      		= config.out_indices if getattr(__UpperCamelCase ,				'out_indices' ,				__UpperCamelCase							) is not None else (-1,)

				snake_case_      		= timm.create_model(
				    config.backbone ,				pretrained=__UpperCamelCase ,				features_only=config.features_only ,				in_chans=config.num_channels ,				out_indices=__UpperCamelCase ,				**__UpperCamelCase ,				)
				# These are used to control the output of the model when called. If output_hidden_states is True, then
				# return_layers is modified to include all layers.
				snake_case_      		= self._backbone.return_layers
				snake_case_      		= {layer['module']: str(__UpperCamelCase							) for i, layer in enumerate(self._backbone.feature_info.info							)}
				super()._init_backbone(__UpperCamelCase							)


			@classmethod
			def 			__lowerCAmelCase (    cls ,				__UpperCamelCase ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(cls ,				['vision', 'timm']							)
				from ...models.timm_backbone import TimmBackboneConfig

				snake_case_      		= kwargs.pop('config' ,				TimmBackboneConfig()							)

				snake_case_      		= kwargs.pop('use_timm_backbone' ,				__UpperCamelCase							)
				if not use_timm:
					raise ValueError('use_timm_backbone must be True for timm backbones'							)

				snake_case_      		= kwargs.pop('num_channels' ,				config.num_channels							)
				snake_case_      		= kwargs.pop('features_only' ,				config.features_only							)
				snake_case_      		= kwargs.pop('use_pretrained_backbone' ,				config.use_pretrained_backbone							)
				snake_case_      		= kwargs.pop('out_indices' ,				config.out_indices							)
				snake_case_      		= TimmBackboneConfig(
				    backbone=__UpperCamelCase ,				num_channels=__UpperCamelCase ,				features_only=__UpperCamelCase ,				use_pretrained_backbone=__UpperCamelCase ,				out_indices=__UpperCamelCase ,				)
				return super()._from_config(__UpperCamelCase ,				**__UpperCamelCase							)


			def 			__lowerCAmelCase (    self ,				__UpperCamelCase							):

				"""simple docstring"""


				pass







			def 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase=None ,				__UpperCamelCase=None ,				__UpperCamelCase=None ,				**__UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= return_dict if return_dict is not None else self.config.use_return_dict
				snake_case_      		= (
				    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
				)
				snake_case_      		= output_attentions if output_attentions is not None else self.config.output_attentions

				if output_attentions:
					raise ValueError('Cannot output attentions for timm backbones at the moment'							)

				if output_hidden_states:
					# We modify the return layers to include all the stages of the backbone
					snake_case_      		= self._all_layers
					snake_case_      		= self._backbone(__UpperCamelCase ,				**__UpperCamelCase							)
					snake_case_      		= self._return_layers
					snake_case_      		= tuple(hidden_states[i] for i in self.out_indices							)
				else:
					snake_case_      		= self._backbone(__UpperCamelCase ,				**__UpperCamelCase							)
					snake_case_      		= None

				snake_case_      		= tuple(__UpperCamelCase							)
				snake_case_      		= tuple(__UpperCamelCase							) if hidden_states is not None else None

				if not return_dict:
					snake_case_      		= (feature_maps,)
					if output_hidden_states:
						snake_case_      		= output + (hidden_states,)
					return output

				return BackboneOutput(feature_maps=__UpperCamelCase ,				hidden_states=__UpperCamelCase ,				attentions=__UpperCamelCase							)



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




			__A						  =  LEDConfig
			__A						  =  {}
			__A						  =  """gelu"""


			def __init__(    self ,				__UpperCamelCase ,				__UpperCamelCase=13 ,				__UpperCamelCase=7 ,				__UpperCamelCase=True ,				__UpperCamelCase=False ,				__UpperCamelCase=99 ,				__UpperCamelCase=32 ,				__UpperCamelCase=2 ,				__UpperCamelCase=4 ,				__UpperCamelCase=37 ,				__UpperCamelCase=0.1 ,				__UpperCamelCase=0.1 ,				__UpperCamelCase=20 ,				__UpperCamelCase=2 ,				__UpperCamelCase=1 ,				__UpperCamelCase=0 ,				__UpperCamelCase=4 ,				):

				"""simple docstring"""


				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 			__lowerCAmelCase (    self							):

				"""simple docstring"""


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







			def 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= TFLEDModel(config=__UpperCamelCase							).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(__UpperCamelCase ,				attention_mask=__UpperCamelCase ,				use_cache=__UpperCamelCase							)

				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(__UpperCamelCase ,				attention_mask=__UpperCamelCase							)[0]
				snake_case_      		= model(__UpperCamelCase ,				attention_mask=__UpperCamelCase ,				past_key_values=__UpperCamelCase							)[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(__UpperCamelCase ,				__UpperCamelCase ,				rtol=1E-3							)







def a(lowercase__			, lowercase__			, lowercase__			, lowercase__=None			, lowercase__=None			, lowercase__=None			, lowercase__=None			, ):





	'''simple docstring'''
	if attention_mask is None:
		snake_case_      		= tf.cast(tf.math.not_equal(lowercase__			, 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  (		__snake_case							,						__snake_case							,						unittest.TestCase							):






			"""simple docstring"""




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


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= TFLEDModelTester(self							)
				snake_case_      		= ConfigTester(self ,				config_class=__UpperCamelCase							)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				self.config_tester.run_common_tests()


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.model_tester.prepare_config_and_inputs_for_common()
				self.model_tester.check_decoder_model_past_large_inputs(*__UpperCamelCase							)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


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

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

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

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

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


			@unittest.skip('LED keeps using potentially symbolic tensors in conditionals and breaks tracing.'							)
			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				pass







			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				pass







def a(lowercase__	):





	'''simple docstring'''
	return tf.constant(lowercase__			, dtype=tf.intaa	)


A											=     1e-4






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






			"""simple docstring"""




			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= TFLEDForConditionalGeneration.from_pretrained('allenai/led-base-16384'							).led

				# change to intended input here
				snake_case_      		= _long_tensor([5_12 * [0, 3_14_14, 2_32, 3_28, 7_40, 11_40, 1_26_95, 69]]							)
				snake_case_      		= _long_tensor([1_28 * [0, 3_14_14, 2_32, 3_28, 7_40, 11_40, 1_26_95, 69]]							)
				snake_case_      		= prepare_led_inputs_dict(model.config ,				__UpperCamelCase ,				__UpperCamelCase							)
				snake_case_      		= model(**__UpperCamelCase							)[0]
				snake_case_      		= (1, 10_24, 7_68)
				self.assertEqual(output.shape ,				__UpperCamelCase							)
				# 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] ,				__UpperCamelCase ,				atol=1E-3							)







			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= TFLEDForConditionalGeneration.from_pretrained('allenai/led-base-16384'							)

				# change to intended input here
				snake_case_      		= _long_tensor([5_12 * [0, 3_14_14, 2_32, 3_28, 7_40, 11_40, 1_26_95, 69]]							)
				snake_case_      		= _long_tensor([1_28 * [0, 3_14_14, 2_32, 3_28, 7_40, 11_40, 1_26_95, 69]]							)
				snake_case_      		= prepare_led_inputs_dict(model.config ,				__UpperCamelCase ,				__UpperCamelCase							)
				snake_case_      		= model(**__UpperCamelCase							)[0]
				snake_case_      		= (1, 10_24, model.config.vocab_size)
				self.assertEqual(output.shape ,				__UpperCamelCase							)
				# 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] ,				__UpperCamelCase ,				atol=1E-3 ,				rtol=1E-3							)
