






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




 def __init__(				self       :			Dict ,	UpperCamelCase_       :			Optional[Any] ,	UpperCamelCase_       :			Dict=13 ,	UpperCamelCase_       :			Dict=7 ,	UpperCamelCase_       :			Tuple=True ,	UpperCamelCase_       :			List[Any]=False ,	UpperCamelCase_       :			List[str]=99 ,	UpperCamelCase_       :			List[str]=32 ,	UpperCamelCase_       :			List[str]=2 ,	UpperCamelCase_       :			List[Any]=4 ,	UpperCamelCase_       :			Tuple=37 ,	UpperCamelCase_       :			Optional[int]=0.1 ,	UpperCamelCase_       :			Union[str, Any]=0.1 ,	UpperCamelCase_       :			int=20 ,	UpperCamelCase_       :			Optional[int]=2 ,	UpperCamelCase_       :			Tuple=1 ,	UpperCamelCase_       :			List[Any]=0 ,	UpperCamelCase_       :			str=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 		lowerCAmelCase_     (				self       :			List[Any]				):


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





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


     """simple docstring"""



     __A         =	TFLEDModel(config=lowercase__				).get_decoder()
     __A         =	inputs_dict['''input_ids''']

     __A         =	input_ids[:1, :]
     __A         =	inputs_dict['''attention_mask'''][:1, :]
     __A         =	1

     # first forward pass
     __A         =	model(lowercase__ ,	attention_mask=lowercase__ ,	use_cache=lowercase__				)

     __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(lowercase__ ,	attention_mask=lowercase__				)[0]
     __A         =	model(lowercase__ ,	attention_mask=lowercase__ ,	past_key_values=lowercase__				)[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(lowercase__ ,	lowercase__ ,	rtol=1e-3				)



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			Tuple			,				__lowercase   :			Dict			,				__lowercase   :			int			,				__lowercase   :			int=None			,				__lowercase   :			List[Any]=None			,				__lowercase   :			Optional[Any]=None			,				__lowercase   :			List[Any]=None			,				)						->      List[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      (	a__			,      a__			,      unittest.TestCase  ):







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




 def 		lowerCAmelCase_     (				self       :			int				):


     """simple docstring"""



     __A         =	TFLEDModelTester(self				)
     __A         =	ConfigTester(self ,	config_class=lowercase__				)




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


     """simple docstring"""



     self.config_tester.run_common_tests()




 def 		lowerCAmelCase_     (				self       :			int				):


     """simple docstring"""



     __A         =	self.model_tester.prepare_config_and_inputs_for_common()
     self.model_tester.check_decoder_model_past_large_inputs(*lowercase__				)




 def 		lowerCAmelCase_     (				self       :			Optional[int]				):


     """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(UpperCamelCase_       :			Union[str, Any]				):
         __A         =	outputs.decoder_attentions
         self.assertEqual(len(lowercase__				) ,	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_       :			Any				):
         __A         =	[t.numpy() for t in outputs.encoder_attentions]
         __A         =	[t.numpy() for t in outputs.encoder_global_attentions]
         self.assertEqual(len(lowercase__				) ,	self.model_tester.num_hidden_layers				)
         self.assertEqual(len(lowercase__				) ,	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(lowercase__				)
         __A         =	model(self._prepare_for_class(lowercase__ ,	lowercase__				)				)
         __A         =	len(lowercase__				)
         self.assertEqual(config.output_hidden_states ,	lowercase__				)
         check_encoder_attentions_output(lowercase__				)

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

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

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

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




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


     """simple docstring"""



     pass





 def 		lowerCAmelCase_     (				self       :			int				):


     """simple docstring"""



     pass



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			Any    )						->      int:





    """simple docstring"""

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


__a      :       Dict				  =   1e-4






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







 '''simple docstring'''




 def 		lowerCAmelCase_     (				self       :			List[str]				):


     """simple docstring"""



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

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





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


     """simple docstring"""



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

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












import unittest

from transformers.testing_utils import CaptureStdout
from transformers.tools.python_interpreter import evaluate



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			Optional[int]    )						->      List[str]:





    """simple docstring"""

    return x + 2






class       __lowercase      (	unittest.TestCase  ):







     '''simple docstring'''




     def 		lowerCAmelCase_     (				self       :			List[Any]				):


         """simple docstring"""



         __A         =	"""x = 3"""
         __A         =	{}
         __A         =	evaluate(UpperCamelCase_ ,	{} ,	state=UpperCamelCase_				)
         assert result == 3
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3}				)

         __A         =	"""x = y"""
         __A         =	{"""y""": 5}
         __A         =	evaluate(UpperCamelCase_ ,	{} ,	state=UpperCamelCase_				)
         # evaluate returns the value of the last assignment.
         assert result == 5
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 5, """y""": 5}				)




     def 		lowerCAmelCase_     (				self       :			List[Any]				):


         """simple docstring"""



         __A         =	"""y = add_two(x)"""
         __A         =	{"""x""": 3}
         __A         =	evaluate(UpperCamelCase_ ,	{"""add_two""": add_two} ,	state=UpperCamelCase_				)
         assert result == 5
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """y""": 5}				)

         # Won't work without the tool
         with CaptureStdout() as out:
             __A         =	evaluate(UpperCamelCase_ ,	{} ,	state=UpperCamelCase_				)
         assert result is None
         assert "tried to execute add_two" in out.out




     def 		lowerCAmelCase_     (				self       :			Optional[Any]				):


         """simple docstring"""



         __A         =	"""x = 3"""
         __A         =	{}
         __A         =	evaluate(UpperCamelCase_ ,	{} ,	state=UpperCamelCase_				)
         assert result == 3
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3}				)




     def 		lowerCAmelCase_     (				self       :			Any				):


         """simple docstring"""



         __A         =	"""test_dict = {'x': x, 'y': add_two(x)}"""
         __A         =	{"""x""": 3}
         __A         =	evaluate(UpperCamelCase_ ,	{"""add_two""": add_two} ,	state=UpperCamelCase_				)
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """y""": 5}				)
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """test_dict""": {"""x""": 3, """y""": 5}}				)




     def 		lowerCAmelCase_     (				self       :			str				):


         """simple docstring"""



         __A         =	"""x = 3\ny = 5"""
         __A         =	{}
         __A         =	evaluate(UpperCamelCase_ ,	{} ,	state=UpperCamelCase_				)
         # evaluate returns the value of the last assignment.
         assert result == 5
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """y""": 5}				)




     def 		lowerCAmelCase_     (				self       :			Optional[int]				):


         """simple docstring"""



         __A         =	"""text = f'This is x: {x}.'"""
         __A         =	{"""x""": 3}
         __A         =	evaluate(UpperCamelCase_ ,	{} ,	state=UpperCamelCase_				)
         # evaluate returns the value of the last assignment.
         assert result == "This is x: 3."
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """text""": """This is x: 3."""}				)




     def 		lowerCAmelCase_     (				self       :			List[str]				):


         """simple docstring"""



         __A         =	"""if x <= 3:\n    y = 2\nelse:\n    y = 5"""
         __A         =	{"""x""": 3}
         __A         =	evaluate(UpperCamelCase_ ,	{} ,	state=UpperCamelCase_				)
         # evaluate returns the value of the last assignment.
         assert result == 2
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """y""": 2}				)

         __A         =	{"""x""": 8}
         __A         =	evaluate(UpperCamelCase_ ,	{} ,	state=UpperCamelCase_				)
         # evaluate returns the value of the last assignment.
         assert result == 5
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 8, """y""": 5}				)




     def 		lowerCAmelCase_     (				self       :			Optional[int]				):


         """simple docstring"""



         __A         =	"""test_list = [x, add_two(x)]"""
         __A         =	{"""x""": 3}
         __A         =	evaluate(UpperCamelCase_ ,	{"""add_two""": add_two} ,	state=UpperCamelCase_				)
         self.assertListEqual(UpperCamelCase_ ,	[3, 5]				)
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """test_list""": [3, 5]}				)




     def 		lowerCAmelCase_     (				self       :			Dict				):


         """simple docstring"""



         __A         =	"""y = x"""
         __A         =	{"""x""": 3}
         __A         =	evaluate(UpperCamelCase_ ,	{} ,	state=UpperCamelCase_				)
         assert result == 3
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """y""": 3}				)




     def 		lowerCAmelCase_     (				self       :			List[Any]				):


         """simple docstring"""



         __A         =	"""test_list = [x, add_two(x)]\ntest_list[1]"""
         __A         =	{"""x""": 3}
         __A         =	evaluate(UpperCamelCase_ ,	{"""add_two""": add_two} ,	state=UpperCamelCase_				)
         assert result == 5
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """test_list""": [3, 5]}				)

         __A         =	"""test_dict = {'x': x, 'y': add_two(x)}\ntest_dict['y']"""
         __A         =	{"""x""": 3}
         __A         =	evaluate(UpperCamelCase_ ,	{"""add_two""": add_two} ,	state=UpperCamelCase_				)
         assert result == 5
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """test_dict""": {"""x""": 3, """y""": 5}}				)





     def 		lowerCAmelCase_     (				self       :			Any				):


         """simple docstring"""



         __A         =	"""x = 0\nfor i in range(3):\n    x = i"""
         __A         =	{}
         __A         =	evaluate(UpperCamelCase_ ,	{"""range""": range} ,	state=UpperCamelCase_				)
         assert result == 2
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 2, """i""": 2}				)




