



'''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 				lowercase__      :
	lowercase__    							=     LEDConfig
	lowercase__    							=     {}
	lowercase__    							=     """gelu"""





	def __init__(					self   :							Tuple					,lowerCamelCase__   :							Tuple					,lowerCamelCase__   :							Any=13					,lowerCamelCase__   :							List[Any]=7					,lowerCamelCase__   :							Tuple=True					,lowerCamelCase__   :							Tuple=False					,lowerCamelCase__   :							Tuple=99					,lowerCamelCase__   :							List[Any]=32					,lowerCamelCase__   :							Optional[Any]=2					,lowerCamelCase__   :							Optional[Any]=4					,lowerCamelCase__   :							List[Any]=37					,lowerCamelCase__   :							Dict=0.1					,lowerCamelCase__   :							Optional[Any]=0.1					,lowerCamelCase__   :							Tuple=20					,lowerCamelCase__   :							Optional[int]=2					,lowerCamelCase__   :							List[Any]=1					,lowerCamelCase__   :							str=0					,lowerCamelCase__   :							Union[str, Any]=4					,):







							'''simple docstring'''
							_UpperCamelCase   :						Any												=  parent
							_UpperCamelCase   :						Optional[int]												=  batch_size
							_UpperCamelCase   :						int												=  seq_length
							_UpperCamelCase   :						Optional[int]												=  is_training
							_UpperCamelCase   :						Optional[int]												=  use_labels
							_UpperCamelCase   :						Optional[int]												=  vocab_size
							_UpperCamelCase   :						List[str]												=  hidden_size
							_UpperCamelCase   :						Optional[int]												=  num_hidden_layers
							_UpperCamelCase   :						str												=  num_attention_heads
							_UpperCamelCase   :						List[str]												=  intermediate_size

							_UpperCamelCase   :						Optional[int]												=  hidden_dropout_prob
							_UpperCamelCase   :						Any												=  attention_probs_dropout_prob
							_UpperCamelCase   :						Any												=  max_position_embeddings
							_UpperCamelCase   :						Dict												=  eos_token_id
							_UpperCamelCase   :						Tuple												=  pad_token_id
							_UpperCamelCase   :						Tuple												=  bos_token_id
							_UpperCamelCase   :						Any												=  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   :						int												=  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   :						Dict												=  (
							    self.seq_length + (self.attention_window - self.seq_length % self.attention_window) % self.attention_window
							)





	def       UpperCamelCase_    (					self   :							Tuple					):







							'''simple docstring'''
							_UpperCamelCase   :						int												=  ids_tensor([self.batch_size, self.seq_length - 1]					,self.vocab_size					)
							_UpperCamelCase   :						Union[str, Any]												=  tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size					)					,1					)
							_UpperCamelCase   :						Any												=  tf.concat([input_ids, eos_tensor]					,axis=1					)

							_UpperCamelCase   :						List[Any]												=  ids_tensor([self.batch_size, self.seq_length]					,self.vocab_size					)

							_UpperCamelCase   :						Any												=  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   :						int												=  prepare_led_inputs_dict(lowerCamelCase__					,lowerCamelCase__					,lowerCamelCase__					)
							_UpperCamelCase   :						Any												=  tf.concat(
							    [tf.zeros_like(lowerCamelCase__					)[:, :-1], tf.ones_like(lowerCamelCase__					)[:, -1:]]					,axis=-1					,)
							_UpperCamelCase   :						Optional[int]												=  global_attention_mask
							return config, inputs_dict





	def       UpperCamelCase_    (					self   :							str					,lowerCamelCase__   :							Tuple					,lowerCamelCase__   :							Tuple					):







							'''simple docstring'''
							_UpperCamelCase   :						Dict												=  TFLEDModel(config=lowerCamelCase__					).get_decoder()
							_UpperCamelCase   :						str												=  inputs_dict['input_ids']

							_UpperCamelCase   :						Union[str, Any]												=  input_ids[:1, :]
							_UpperCamelCase   :						str												=  inputs_dict['attention_mask'][:1, :]
							_UpperCamelCase   :						str												=  1

							# first forward pass
							_UpperCamelCase   :						int												=  model(lowerCamelCase__					,attention_mask=lowerCamelCase__					,use_cache=lowerCamelCase__					)

							_UpperCamelCase							,     _UpperCamelCase   :						int												=  outputs.to_tuple()

							# create hypothetical next token and extent to next_input_ids
							_UpperCamelCase   :						int												=  ids_tensor((self.batch_size, 3)					,config.vocab_size					)
							_UpperCamelCase   :						Optional[Any]												=  tf.cast(ids_tensor((self.batch_size, 3)					,2					)					,tf.inta					)

							# append to next input_ids and
							_UpperCamelCase   :						Optional[int]												=  tf.concat([input_ids, next_tokens]					,axis=-1					)
							_UpperCamelCase   :						Tuple												=  tf.concat([attention_mask, next_attn_mask]					,axis=-1					)

							_UpperCamelCase   :						Any												=  model(lowerCamelCase__					,attention_mask=lowerCamelCase__					)[0]
							_UpperCamelCase   :						Tuple												=  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
							_UpperCamelCase   :						List[str]												=  int(ids_tensor((1,)					,output_from_past.shape[-1]					)					)
							_UpperCamelCase   :						List[str]												=  output_from_no_past[:, -3:, random_slice_idx]
							_UpperCamelCase   :						Tuple												=  output_from_past[:, :, random_slice_idx]

							# test that outputs are equal for slice
							tf.debugging.assert_near(lowerCamelCase__					,lowerCamelCase__					,rtol=1E-3					)






def A__       (	UpperCAmelCase_    ,    UpperCAmelCase_    ,    UpperCAmelCase_    ,    UpperCAmelCase_=None    ,    UpperCAmelCase_=None    ,    UpperCAmelCase_=None    ,    UpperCAmelCase_=None    ,    ):
						if attention_mask is None:
												_UpperCamelCase   :						int												=  tf.cast(tf.math.not_equal(UpperCAmelCase_    ,    config.pad_token_id				)    ,    tf.inta				)
						if decoder_attention_mask is None:
												_UpperCamelCase   :						Optional[int]												=  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   :						Dict												=  tf.ones((config.encoder_layers, config.encoder_attention_heads)				)
						if decoder_head_mask is None:
												_UpperCamelCase   :						List[Any]												=  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__      (				lowercase		,   lowercase		,   unittest.TestCase  ):
	lowercase__    							=     (TFLEDForConditionalGeneration, TFLEDModel) if is_tf_available() else ()
	lowercase__    							=     (TFLEDForConditionalGeneration,) if is_tf_available() else ()
	lowercase__    							=     (
	    {
	        """conversational""": TFLEDForConditionalGeneration,
	        """feature-extraction""": TFLEDModel,
	        """summarization""": TFLEDForConditionalGeneration,
	        """text2text-generation""": TFLEDForConditionalGeneration,
	        """translation""": TFLEDForConditionalGeneration,
	    }
	    if is_tf_available()
	    else {}
	)
	lowercase__    							=     True
	lowercase__    							=     False
	lowercase__    							=     False
	lowercase__    							=     False





	def       UpperCamelCase_    (					self   :							Any					):







							'''simple docstring'''
							_UpperCamelCase   :						str												=  TFLEDModelTester(self					)
							_UpperCamelCase   :						int												=  ConfigTester(self					,config_class=lowerCamelCase__					)





	def       UpperCamelCase_    (					self   :							Optional[int]					):







							'''simple docstring'''
							self.config_tester.run_common_tests()





	def       UpperCamelCase_    (					self   :							Union[str, Any]					):







							'''simple docstring'''
							_UpperCamelCase   :						Optional[int]												=  self.model_tester.prepare_config_and_inputs_for_common()
							self.model_tester.check_decoder_model_past_large_inputs(*lowerCamelCase__					)





	def       UpperCamelCase_    (					self   :							Tuple					):







							'''simple docstring'''
							_UpperCamelCase							,     _UpperCamelCase   :						Dict												=  self.model_tester.prepare_config_and_inputs_for_common()
							_UpperCamelCase   :						Tuple												=  tf.zeros_like(inputs_dict['attention_mask']					)
							_UpperCamelCase   :						Union[str, Any]												=  2
							_UpperCamelCase   :						Tuple												=  tf.where(
							    tf.range(self.model_tester.seq_length					)[None, :] < num_global_attn_indices					,1					,inputs_dict['global_attention_mask']					,)

							_UpperCamelCase   :						Optional[int]												=  True
							_UpperCamelCase   :						int												=  self.model_tester.seq_length
							_UpperCamelCase   :						Optional[Any]												=  self.model_tester.encoder_seq_length

							def check_decoder_attentions_output(lowerCamelCase__   :							Any					):
													_UpperCamelCase   :						Union[str, Any]												=  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__   :							List[str]					):
													_UpperCamelCase   :						Optional[int]												=  [t.numpy() for t in outputs.encoder_attentions]
													_UpperCamelCase   :						Optional[int]												=  [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:
													_UpperCamelCase   :						str												=  True
													_UpperCamelCase   :						Dict												=  False
													_UpperCamelCase   :						str												=  False
													_UpperCamelCase   :						Optional[Any]												=  model_class(lowerCamelCase__					)
													_UpperCamelCase   :						Dict												=  model(self._prepare_for_class(lowerCamelCase__					,lowerCamelCase__					)					)
													_UpperCamelCase   :						Optional[int]												=  len(lowerCamelCase__					)
													self.assertEqual(config.output_hidden_states					,lowerCamelCase__					)
													check_encoder_attentions_output(lowerCamelCase__					)

													if self.is_encoder_decoder:
																			_UpperCamelCase   :						Optional[int]												=  model_class(lowerCamelCase__					)
																			_UpperCamelCase   :						Optional[int]												=  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"]
													_UpperCamelCase   :						int												=  True
													_UpperCamelCase   :						List[str]												=  model_class(lowerCamelCase__					)
													_UpperCamelCase   :						Dict												=  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
													_UpperCamelCase   :						Tuple												=  True
													_UpperCamelCase   :						Any												=  True
													_UpperCamelCase   :						List[Any]												=  model_class(lowerCamelCase__					)
													_UpperCamelCase   :						Optional[int]												=  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   :							Optional[int]					):







							'''simple docstring'''
							pass





	def       UpperCamelCase_    (					self   :							Optional[Any]					):







							'''simple docstring'''
							# TODO: Head-masking not yet implement
							pass






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


snake_case_					:      Optional[int]    			=    1e-4






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





	def       UpperCamelCase_    (					self   :							Union[str, Any]					):







							'''simple docstring'''
							_UpperCamelCase   :						str												=  TFLEDForConditionalGeneration.from_pretrained('allenai/led-base-16384'					).led

							# change to intended input here
							_UpperCamelCase   :						Dict												=  _long_tensor([512 * [0, 31414, 232, 328, 740, 1140, 12695, 69]]					)
							_UpperCamelCase   :						Tuple												=  _long_tensor([128 * [0, 31414, 232, 328, 740, 1140, 12695, 69]]					)
							_UpperCamelCase   :						Tuple												=  prepare_led_inputs_dict(model.config					,lowerCamelCase__					,lowerCamelCase__					)
							_UpperCamelCase   :						Optional[Any]												=  model(**lowerCamelCase__					)[0]
							_UpperCamelCase   :						Union[str, Any]												=  (1, 1024, 768)
							self.assertEqual(output.shape					,lowerCamelCase__					)
							# change to expected output here
							_UpperCamelCase   :						int												=  tf.convert_to_tensor(
							    [[2.3_0_5_0, 2.8_2_7_9, 0.6_5_3_1], [-1.8_4_5_7, -0.1_4_5_5, -3.5_6_6_1], [-1.0_1_8_6, 0.4_5_8_6, -2.2_0_4_3]]					,)
							tf.debugging.assert_near(output[:, :3, :3]					,lowerCamelCase__					,atol=1E-3					)





	def       UpperCamelCase_    (					self   :							Any					):







							'''simple docstring'''
							_UpperCamelCase   :						str												=  TFLEDForConditionalGeneration.from_pretrained('allenai/led-base-16384'					)

							# change to intended input here
							_UpperCamelCase   :						Dict												=  _long_tensor([512 * [0, 31414, 232, 328, 740, 1140, 12695, 69]]					)
							_UpperCamelCase   :						int												=  _long_tensor([128 * [0, 31414, 232, 328, 740, 1140, 12695, 69]]					)
							_UpperCamelCase   :						Optional[int]												=  prepare_led_inputs_dict(model.config					,lowerCamelCase__					,lowerCamelCase__					)
							_UpperCamelCase   :						Optional[int]												=  model(**lowerCamelCase__					)[0]
							_UpperCamelCase   :						Optional[Any]												=  (1, 1024, model.config.vocab_size)
							self.assertEqual(output.shape					,lowerCamelCase__					)
							# change to expected output here
							_UpperCamelCase   :						str												=  tf.convert_to_tensor(
							    [[3_3.6_5_0_7, 6.4_5_7_2, 1_6.8_0_8_9], [5.8_7_3_9, -2.4_2_3_8, 1_1.2_9_0_2], [-3.2_1_3_9, -4.3_1_4_9, 4.2_7_8_3]]					,)
							tf.debugging.assert_near(output[:, :3, :3]					,lowerCamelCase__					,atol=1E-3					,rtol=1E-3					)







'''simple docstring'''
import unittest

from transformers import AutoConfig, AutoTokenizer, BertConfig, TensorType, is_flax_available
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, require_flax, slow


if is_flax_available():
					import jax

					from transformers.models.auto.modeling_flax_auto import FlaxAutoModel
					from transformers.models.bert.modeling_flax_bert import FlaxBertModel
					from transformers.models.roberta.modeling_flax_roberta import FlaxRobertaModel






@require_flax
class 				lowercase__      (				unittest.TestCase  ):





	@slow
	def       UpperCamelCase_    (					self   :							Union[str, Any]					):







							'''simple docstring'''
							for model_name in ["bert-base-cased", "bert-large-uncased"]:
													with self.subTest(lowerCamelCase__					):
																			_UpperCamelCase   :						Optional[int]												=  AutoConfig.from_pretrained(lowerCamelCase__					)
																			self.assertIsNotNone(lowerCamelCase__					)
																			self.assertIsInstance(lowerCamelCase__					,lowerCamelCase__					)

																			_UpperCamelCase   :						Dict												=  FlaxAutoModel.from_pretrained(lowerCamelCase__					)
																			self.assertIsNotNone(lowerCamelCase__					)
																			self.assertIsInstance(lowerCamelCase__					,lowerCamelCase__					)





	@slow
	def       UpperCamelCase_    (					self   :							List[Any]					):







							'''simple docstring'''
							for model_name in ["roberta-base", "roberta-large"]:
													with self.subTest(lowerCamelCase__					):
																			_UpperCamelCase   :						Any												=  AutoConfig.from_pretrained(lowerCamelCase__					)
																			self.assertIsNotNone(lowerCamelCase__					)
																			self.assertIsInstance(lowerCamelCase__					,lowerCamelCase__					)

																			_UpperCamelCase   :						Any												=  FlaxAutoModel.from_pretrained(lowerCamelCase__					)
																			self.assertIsNotNone(lowerCamelCase__					)
																			self.assertIsInstance(lowerCamelCase__					,lowerCamelCase__					)





	@slow
	def       UpperCamelCase_    (					self   :							Union[str, Any]					):







							'''simple docstring'''
							for model_name in ["bert-base-cased", "bert-large-uncased"]:
													_UpperCamelCase   :						Optional[int]												=  AutoTokenizer.from_pretrained(lowerCamelCase__					)
													_UpperCamelCase   :						List[Any]												=  FlaxBertModel.from_pretrained(lowerCamelCase__					)
													_UpperCamelCase   :						Optional[int]												=  tokenizer('Do you support jax jitted function?'					,return_tensors=TensorType.JAX					)

													@jax.jit
													def eval(**lowerCamelCase__   :							Union[str, Any]					):
																			return model(**lowerCamelCase__					)

													eval(**lowerCamelCase__					).block_until_ready()





	@slow
	def       UpperCamelCase_    (					self   :							Dict					):







							'''simple docstring'''
							for model_name in ["roberta-base", "roberta-large"]:
													_UpperCamelCase   :						Optional[int]												=  AutoTokenizer.from_pretrained(lowerCamelCase__					)
													_UpperCamelCase   :						Tuple												=  FlaxRobertaModel.from_pretrained(lowerCamelCase__					)
													_UpperCamelCase   :						Optional[Any]												=  tokenizer('Do you support jax jitted function?'					,return_tensors=TensorType.JAX					)

													@jax.jit
													def eval(**lowerCamelCase__   :							Union[str, Any]					):
																			return model(**lowerCamelCase__					)

													eval(**lowerCamelCase__					).block_until_ready()





	def       UpperCamelCase_    (					self   :							Union[str, Any]					):







							'''simple docstring'''
							with self.assertRaisesRegex(
							    lowerCamelCase__					,'bert-base is not a local folder and is not a valid model identifier'					):
													_UpperCamelCase   :						int												=  FlaxAutoModel.from_pretrained('bert-base'					)





	def       UpperCamelCase_    (					self   :							str					):







							'''simple docstring'''
							with self.assertRaisesRegex(
							    lowerCamelCase__					,R'aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'					):
													_UpperCamelCase   :						Tuple												=  FlaxAutoModel.from_pretrained(lowerCamelCase__					,revision='aaaaaa'					)





	def       UpperCamelCase_    (					self   :							List[str]					):







							'''simple docstring'''
							with self.assertRaisesRegex(
							    lowerCamelCase__					,'hf-internal-testing/config-no-model does not appear to have a file named flax_model.msgpack'					,):
													_UpperCamelCase   :						List[Any]												=  FlaxAutoModel.from_pretrained('hf-internal-testing/config-no-model'					)





	def       UpperCamelCase_    (					self   :							List[str]					):







							'''simple docstring'''
							with self.assertRaisesRegex(lowerCamelCase__					,'Use `from_pt=True` to load this model'					):
													_UpperCamelCase   :						Tuple												=  FlaxAutoModel.from_pretrained('hf-internal-testing/tiny-bert-pt-only'					)


