



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 	_lowercase							:





	"""simple docstring"""
	__A							=							LEDConfig
	__A							=							{}
	__A							=							'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		,       ):
						"""simple docstring"""

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

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




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

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

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

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




	def      UpperCamelCase_       (self		,       lowerCamelCase_		,       lowerCamelCase_		):
						"""simple docstring"""

						a					     =   TFLEDModel(config=lowerCamelCase_		).get_decoder()
						a					     =   inputs_dict["""input_ids"""]

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

						# first forward pass
						a					     =   model(lowerCamelCase_		,       attention_mask=lowerCamelCase_		,       use_cache=lowerCamelCase_		)

						a					     =   outputs.to_tuple()

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

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

						a					     =   model(lowerCamelCase_		,       attention_mask=lowerCamelCase_		)[0]
						a					     =   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
						a					     =   int(ids_tensor((1,)		,       output_from_past.shape[-1]		)		)
						a					     =   output_from_no_past[:, -3:, random_slice_idx]
						a					     =   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       :    Tuple ,     A       :    List[Any] ,     A       :    Dict ,     A       :    Tuple=None ,     A       :    Any=None ,     A       :    int=None ,     A       :    Union[str, Any]=None ,     )		->       Union[str, Any]:
					"""simple docstring"""





					if attention_mask is None:
										a					     =   tf.cast(tf.math.not_equal(lowerCamelCase__ ,     config.pad_token_id       ) ,     tf.inta       )
					if decoder_attention_mask is None:
										a					     =   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:
										a					     =   tf.ones((config.encoder_layers, config.encoder_attention_heads)       )
					if decoder_head_mask is None:
										a					     =   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 	_lowercase							(   __UpperCAmelCase,      __UpperCAmelCase,      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      UpperCamelCase_       (self		):
						"""simple docstring"""

						a					     =   TFLEDModelTester(self		)
						a					     =   ConfigTester(self		,       config_class=lowerCamelCase_		)




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

						self.config_tester.run_common_tests()




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

						a					     =   self.model_tester.prepare_config_and_inputs_for_common()
						self.model_tester.check_decoder_model_past_large_inputs(*lowerCamelCase_		)




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

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

						a					     =   True
						a					     =   self.model_tester.seq_length
						a					     =   self.model_tester.encoder_seq_length

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

											if self.is_encoder_decoder:
																a					     =   model_class(lowerCamelCase_		)
																a					     =   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"]
											a					     =   True
											a					     =   model_class(lowerCamelCase_		)
											a					     =   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
											a					     =   True
											a					     =   True
											a					     =   model_class(lowerCamelCase_		)
											a					     =   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		):
						"""simple docstring"""

						pass




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

						pass






def     a(      A       :    int       )		->       List[Any]:
					"""simple docstring"""





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


_lowercase: Optional[Any]      			=		1E-4







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





	"""simple docstring"""




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

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

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

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

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




import unittest

import numpy as np

from transformers import DistilBertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
						import jax.numpy as jnp

						from transformers.models.distilbert.modeling_flax_distilbert import (
						    FlaxDistilBertForMaskedLM,
						    FlaxDistilBertForMultipleChoice,
						    FlaxDistilBertForQuestionAnswering,
						    FlaxDistilBertForSequenceClassification,
						    FlaxDistilBertForTokenClassification,
						    FlaxDistilBertModel,
						)







class 	_lowerCAmelCase							(			unittest.TestCase ):


	def __init__(self       ,      lowercase       ,      lowercase=13       ,      lowercase=7       ,      lowercase=True       ,      lowercase=True       ,      lowercase=True       ,      lowercase=True       ,      lowercase=99       ,      lowercase=32       ,      lowercase=5       ,      lowercase=4       ,      lowercase=37       ,      lowercase="gelu"       ,      lowercase=0.1       ,      lowercase=0.1       ,      lowercase=512       ,      lowercase=16       ,      lowercase=2       ,      lowercase=0.02       ,      lowercase=4       ,      ):
				A_	:						List[Any]	=     parent
				A_	:						str	=     batch_size
				A_	:						List[Any]	=     seq_length
				A_	:						Dict	=     is_training
				A_	:						List[Any]	=     use_attention_mask
				A_	:						Any	=     use_token_type_ids
				A_	:						Optional[int]	=     use_labels
				A_	:						Tuple	=     vocab_size
				A_	:						List[str]	=     hidden_size
				A_	:						List[str]	=     num_hidden_layers
				A_	:						Optional[Any]	=     num_attention_heads
				A_	:						int	=     intermediate_size
				A_	:						Optional[Any]	=     hidden_act
				A_	:						List[Any]	=     hidden_dropout_prob
				A_	:						Optional[Any]	=     attention_probs_dropout_prob
				A_	:						Any	=     max_position_embeddings
				A_	:						Union[str, Any]	=     type_vocab_size
				A_	:						int	=     type_sequence_label_size
				A_	:						Any	=     initializer_range
				A_	:						List[str]	=     num_choices


	def     _a	(self							):
				A_	:						List[str]	=     ids_tensor([self.batch_size, self.seq_length]       ,      self.vocab_size							)

				A_	:						Any	=     None
				if self.use_attention_mask:
							A_	:						Any	=     random_attention_mask([self.batch_size, self.seq_length]							)

				A_	:						Union[str, Any]	=     DistilBertConfig(
				    vocab_size=self.vocab_size       ,      dim=self.hidden_size       ,      n_layers=self.num_hidden_layers       ,      n_heads=self.num_attention_heads       ,      hidden_dim=self.intermediate_size       ,      hidden_act=self.hidden_act       ,      dropout=self.hidden_dropout_prob       ,      attention_dropout=self.attention_probs_dropout_prob       ,      max_position_embeddings=self.max_position_embeddings       ,      initializer_range=self.initializer_range       ,      tie_weights_=lowercase       ,      )

				return config, input_ids, attention_mask


	def     _a	(self							):
				A_	:						List[str]	=     self.prepare_config_and_inputs()
				A_,    A_,    A_	:						str	=     config_and_inputs
				A_	:						Any	=     {"""input_ids""": input_ids, """attention_mask""": attention_mask}
				return config, inputs_dict







@require_flax
class 	_lowerCAmelCase							(			__UpperCAmelCase      ,			unittest.TestCase ):
	__SCREAMING_SNAKE_CASE   :		Optional[Any]					=					(
	    (
	        FlaxDistilBertModel,
	        FlaxDistilBertForMaskedLM,
	        FlaxDistilBertForMultipleChoice,
	        FlaxDistilBertForQuestionAnswering,
	        FlaxDistilBertForSequenceClassification,
	        FlaxDistilBertForTokenClassification,
	        FlaxDistilBertForQuestionAnswering,
	    )
	    if is_flax_available()
	    else ()
	)


	def     _a	(self							):
				A_	:						Tuple	=     FlaxDistilBertModelTester(self							)


	@slow
	def     _a	(self							):
				for model_class_name in self.all_model_classes:
							A_	:						Union[str, Any]	=     model_class_name.from_pretrained("""distilbert-base-uncased"""							)
							A_	:						Any	=     model(np.ones((1, 1)							)							)
							self.assertIsNotNone(lowercase							)







@require_flax
class 	_lowerCAmelCase							(			unittest.TestCase ):


	@slow
	def     _a	(self							):
				A_	:						List[str]	=     FlaxDistilBertModel.from_pretrained("""distilbert-base-uncased"""							)
				A_	:						Optional[Any]	=     np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]							)
				A_	:						int	=     np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]							)
				A_	:						Optional[int]	=     model(lowercase       ,      attention_mask=lowercase							)[0]
				A_	:						Optional[Any]	=     (1, 11, 768)
				self.assertEqual(output.shape       ,      lowercase							)
				A_	:						Union[str, Any]	=     np.array([[[-0.16_39, 0.32_99, 0.16_48], [-0.17_46, 0.32_89, 0.17_10], [-0.18_84, 0.33_57, 0.18_10]]]							)

				self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]       ,      lowercase       ,      atol=1E-4							)							)